SIGN IN SIGN UP

[Transform] Add VerifyBufferInit, a general buffer-initialization check (#2956)

* [Op] Add read-before-write semantics to tile operators

GetAccessRegions() answers a may-read question: could this operator read this
region? It is deliberately conservative, which is correct for the dependency
analysis it was built for. There is no way to ask the stricter question of
whether an operator consumes a region.s preexisting contents before
establishing a value of its own.

GetReadBeforeWriteRegions() asks it. The default returns GetAccessRegions()
.reads, so operators whose read set does not depend on an argument value need
no change, and existing consumers of GetAccessRegions() are untouched.

Gemm and GemmSP override it. GetAccessRegions() lists the accumulator as a
read whenever the clear cannot be proven (!is_one(clearAccum_)), which is
right for pipelining: a conditional clear genuinely is a read dependency. It
is wrong for a read-before-write question, because the common pipelined idiom
clear_accum=(k == 0) is neither literal true nor literal false. The overrides
use is_zero() instead.

Note that operators such as ReduceOp need no override: its clear is a plain
bool rather than a PrimExpr, so the default is already exact.

* [Transform] Add VerifyBufferInit, a general buffer-initialization check

Reading a non-global buffer that nothing has written yields whatever those
locations last held. That is undefined behaviour: it can look correct on one
architecture and produce NaN on another (issue #2936).

VerifyBufferInit walks a PrimFunc before LowerTileOp, where tile operators
still declare their access regions, so writes performed by T.copy, T.clear and
friends are visible rather than opaque. It warns and returns the function
unchanged; it never fails a build.

The analysis is conservative by construction, in three ways:

- A call the pass cannot interpret may write through its arguments, so it is
  assumed to. TCallEffectKind distinguishes these from pure calls; an
  operator with no registered effect kind is assumed to write. This also
  covers tl.access_ptr, which keeps a BufferLoad rather than a data Var.
- Shared memory ignores source order. A warp-specialized kernel fills a shared
  buffer from a producer branch that appears after the consumer branch that
  reads it; the two run concurrently and coordinate through barriers, so
  source order is not an approximation of happens-before. For shared scopes
  the question is whether any *other* operation writes the buffer at all.
  Per-thread storage keeps order-sensitive tracking, where source order is
  that thread.s execution order.
- Global memory and mbarrier scopes are never reported.

Excluding the reading operator.s own write is what keeps #2936 reported: a
gemm.s only writer of its accumulator is itself.

* [Transform] Enable the buffer-initialization check in every pipeline

Wired into all five backends immediately after LayoutReducer, and on by
default with tl.disable_buffer_init_check to opt out.

The slot matters in four ways. LegalizeNegativeIndex has run, so buffer
indices are canonical. LowerTileOp has not, so tile operators still carry
their access regions. InjectSoftwarePipeline has not, so loop bodies are in
source order rather than rotated across stages. And LayoutReducer has run:
before it, a tl.tileop.finalize_reducer call carries a single argument while
its builder reads args[1], so ParseOperator throws on it. Every pre-existing
ParseOperator caller sits after that point, which is why the incomplete form
had not been parsed before.

22 tests cover the gemm cases, the general cases, the two scope-ordering
rules, report aggregation, and the config plumbing end to end.

* [Transform] Route both buffer-init walks through one opaque-write rule

The whole-function collector and the ordered walk each decided for
themselves what an uninterpretable call may write, and they disagreed.
The collector recorded any BufferLoad argument, so a read-only
tl.access_ptr counted as a writer and could suppress a real finding on
shared memory, while a write through tvm_access_ptr passed a data Var
rather than a BufferLoad and was missed entirely. It also applied the
rule to pure calls, whose arguments only read.

Both now share ForEachOpaqueWrite and the same escape predicate, so the
access-mask and pointer-form handling cannot drift apart again.

Also adds a regression test for the unguarded loop-carried read. That
shape is reported and the report is correct: the first iteration reads
the fragment before the write at the bottom of the body reaches it. The
case the analysis does not decide is a read the loop guards out of its
own first iteration, which is path sensitivity rather than ordering.

Verified: 23 verifier tests pass; testing/python/language is 997 passed
with the same 5 pre-existing failures and zero warnings; all 284 example
scripts emit one warning, the known uninitialized lse_max_local in
examples/flash_decoding/example_mha_inference.py.

* [Transform] Record a store's destination after evaluating it

A BufferStore reads its own right-hand side before it establishes
anything, so marking the destination written first hid every
self-referential read: `x[i] = x[i] + 1` with nothing written earlier
went unreported.

I had defended that ordering on the grounds that flagging
`idx = T.if_then_else(cond, i, idx)` turns the check into
definite-assignment analysis over conventional idioms. Looking at the
two examples that motivated it, that defence does not survive.
examples/grouped_gemm/example_grouped_gemm_fwd_ptr.py writes
`cur_batch_idx = 0` before exactly this loop; example_grouped_gemm_fwd.py
and example_grouped_gemm_bwd.py do not. The initializer is already there
in the sibling, so this is an oversight in two of three copies rather
than an idiom the check misreads.

Cross-thread scopes are unaffected: they ignore source order and ask
only whether another node writes the buffer.

Verified: 24 verifier tests pass; testing/python/language is 997 passed
with the same 5 pre-existing failures and zero warnings; all 284 example
scripts yield 3 warnings in 2 places, the lse_max_local reduce
destination and the two cur_batch_idx sites, with no false positives.

Note the two grouped_gemm examples also fail their own torch.allclose
check, but that is unrelated: initializing cur_batch_idx leaves the
mismatch at an unchanged 0.25 max absolute error over ~3.4k of 1572864
elements.

* [Transform] Discount a store's own write when checking its reads

CheckRead already took the reading node, so a cross-thread scope could
ask whether some other node writes the buffer. That is the rule keeping
a gemm accumulating into an untouched accumulator reported. Loads passed
no reader, so a store that read its own destination counted itself as
its own initializer: `s[i] = s[i] + A[i]` on shared memory was silent
while the identical code on a fragment was reported.

The enclosing store is now published while its subexpressions are
visited. A store is recorded as a writer of nothing but its own
destination -- an opaque call inside the value is recorded under the
call node -- so discounting it cannot suppress a report for any other
buffer.

Also covers the reducer v2 idiom introduced by #2940. reducer_init,
reducer_update and finalize_reducer inherit the default
GetReadBeforeWriteRegions and need no override, which is the same
result the pass already gets for T.reduce.

* [Transform] Exempt a parameter buffer from the check at any scope

SeedParams treats every parameter as written by the caller, but the
cross-thread branch of CheckRead asks only whether another node in the
body writes the buffer, and the caller is not one. A shared-scope
parameter therefore reported while an identical local-scope parameter
stayed silent.

Record the parameters in their own set and admit them ahead of the
scope-dependent question, so WrittenByAnotherNode is left untouched for
every buffer the body has to establish itself.

* [Transform] Print the buffer name and widen the parse fallback

Streaming the Buffer sent a one-word diagnostic through the reflection
repr printer. The name is a String with a direct stream overload and is
what the message wanted either way, so print it directly.

TryParseOperator caught tvm::ffi::Error, which every builder failure
seen so far raises. It derives from std::exception, so catching the base
extends the documented degrade-rather-than-abort rule to the failures
that would not.

* [Example] Initialize the buffers the new check reports

Three sites, all read before anything writes them. None is new and none
is known to misbehave; they are benign on the inputs these examples use,
which is why nothing has surfaced them.

grouped_gemm fwd and bwd allocate `cur_batch_idx` and immediately carry
it forward through `T.if_then_else(cond, i, cur_batch_idx)`, so a false
predicate on the first iteration returns whatever the variable held.
`cur_batch_idx = 0` is what example_grouped_gemm_fwd_ptr.py already does,
and has done since its first commit (#1923) — a second author writing the
same loop initialized it unprompted.

flash_decoding passes an uninitialized `lse_max_local` to
`T.reduce_max(..., clear=False)`, which makes the destination a read by
the operator's own semantics. It is the only write to that buffer and it
sits two lines below `T.clear(lse_logsum_local)`, so `clear=True` is what
was meant.

Verified: all 284 example scripts now compile with zero warnings.
example_mha_inference.py still reports "All checks passed!" and its
output is bit-identical to the current form on an RTX 3060 Ti.

The two grouped_gemm examples still fail their own torch.allclose check.
That is a separate, pre-existing issue untouched here: the mismatch is
unchanged at 0.25 max absolute error over ~3.4k of 1572864 elements
whether or not cur_batch_idx is initialized.

* [Refactor] Read the buffer-init config inside the pass

The pass function already receives the PassContext, so it can gate
itself on tl.disable_buffer_init_check rather than have every backend
pipeline consult a Python helper first. This matches ThreadSync, which
reads tl.disable_thread_storage_sync the same way and is scheduled
unconditionally.

Drops should_enable_buffer_init_check and the five gated call sites.
The two tests that asserted on the helper now call the pass directly,
which also covers the case the Python gate never did: invoking
VerifyBufferInit outside the pipeline under a disabling PassContext.
R
Ryan Lei committed
baf8a16348c835294220fc997788dfe20abf1351
Parent: ddc2c54
Committed by GitHub <noreply@github.com> on 8/17/2026, 4:12:36 AM