[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