SIGN IN SIGN UP

[Bugfix] Fix wait parity for explicit mbarriers in pipelined loops (#3087)

A synchronous T.gemm(mbar=...) inside T.Pipelined waited on the wrong
mbarrier parity. InjectSoftwarePipeline annotated every copy/im2col/gemm
call with the ring parity (k / num_stages) % 2, but that law belongs to
the compiler-generated pipeline_mbar ring: an explicit mbarrier whose
arrive happens inside a tile op is deliberately excluded from
ExpandPipelineBarriers, so it remains a single slot that flips once per
iteration. The auto-emitted wait asked for a parity that had already
completed, returned immediately, and the next iteration's copies
overwrote a shared-memory stage the still-running MMA was reading
(wrong results observed at num_stages >= 3).

Annotate only the ops that consume the phase -- those that auto-emit a
paired arrive+wait on an explicit mbarrier: sync gemm with an mbar, and
copies carrying an explicit "barrier" that are not is_tma_copy -- and
annotate them with the logical iteration parity (k - min) % 2, written
against the original loop var. EmitImpl's existing per-block
Substitute(loop_var -> normalized_access_index) specializes the
expression for every prologue/body/epilogue clone, since the tirx
mutators already rewrite PrimExprs inside call annotations.
Pipeline-managed TMA copies keep their ring parity in the explicit
waits the pass inserts, and the explicit-async T.tcgen05_gemm never
auto-waits, so neither is annotated.

Drop the matching ring branch from LowerTileOp's fallback: it keyed on
the "num_stages" annotation, which InjectSoftwarePipeline strips, so it
was dead on pipelined loops and wrong on unpipelined ones -- a
compiler-generated per-site barrier flips once per loop iteration
regardless of pipeline depth. The fallback is now the zero-based epoch
of the nearest enclosing serial loop.

Add transform-level tests asserting the annotated phase and the emitted
wait parities on sm_90a, so the fix is covered without sm_100 hardware;
the loops start at a nonzero min to pin the epoch normalization. Add a
tcgen05 correctness test running the failing kernel at num_stages 2
and 3.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Y
Yongqi Zhuo committed
728f2c88a19ad3c936d2fa1768e38bcfa38d5edb
Parent: 234c0bc
Committed by GitHub <noreply@github.com> on 8/26/2026, 7:18:10 AM