[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