[AMD] Warp-pipeline: back-to-back loop optimization & flat (unrolled) pipeline support#9929
Merged
Conversation
…d loops
When two warp-pipelined loops execute consecutively, ConvertWarpPipeline
previously emitted a full reconverge/re-phase-shift/pre-barrier sequence
between them:
scf.for { loop 1 }
cond_barrier(warpLow) ← post-loop reconverge
ttg.barrier local ← pre-barrier for loop 2
cond_barrier(warpHigh) ← pre-loop phase shift
scf.for { loop 2 }
The post-loop reconverge and pre-loop phase shift are complementary
predicates on the same counter-based S_BARRIER, so they cancel out.
The intervening ttg.barrier local is redundant when loop 1's
wrap-around cluster barrier already includes a local fence (i.e. the
dependency analysis determined an LDS read/write hazard exists across
the wrap-around point). In that case, all pending LDS writes are
already resolved before loop 1 yields, and ModuleMembarAnalysis will
not need to insert additional barriers between the loops.
This patch adds a post-processing pass (eliminateRedundantCondBarriers)
that detects this pattern and erases the three redundant ops, reducing
the barrier overhead to:
scf.for { loop 1 }
scf.for { loop 2 }
cond_barrier(warpLow) ← final reconverge only
The pass runs after all scf.for loops have been converted (patternFor)
but before execute_regions are inlined (patternInline), preserving the
scf.for / cond_barrier adjacency needed for pattern matching.
Also updates the f16_gemm_warp_pipeline_gfx1250.py example to use
range() (producing scf.for) instead of static_range() (which unrolls
at the Python level) for the epilogue loop, and wraps its stages in
warp_pipeline_stage annotations so the back-to-back optimization can
apply.
Extend the warp-pipeline infrastructure to handle loops unrolled at the Python level (e.g. via static_range/ttgl.static_range). Previously, warp-pipelining only worked with scf.for loops. Unrolled loops produce flat sequences of border markers in the IR that were silently ignored. Three main changes: 1. WarpPipeliner: add createFlatPipeline() Scans each block for triton.warp_pipeline.border markers outside scf.for. Groups the operations between borders into clusters and wraps each in an scf.execute_region with triton.warp_pipeline.stage, triton.warp_pipeline.priority, and no_inline attributes — the same representation createPipeline() produces for loop bodies. 2. ConvertWarpPipeline: add processUnrolledPipelineRegions() + emitPipelinedFlat() After the existing patternFor converts scf.for loops, this new pass walks each function block for contiguous sequences of flat scf.execute_region ops (with triton.warp_pipeline.stage). For each sequence it emits the full barrier structure: pre-barrier, phase shift (cond_barrier warpHigh), linear dependency analysis for cluster barriers (no wrap-around since the sequence is finite), priority management (s_setprio), and post-sequence reconverge (cond_barrier warpLow). The execute_regions are then inlined by the existing InlineWarpPipelineExecuteRegionPattern. Also extends eliminateRedundantCondBarriers() to handle the case where a pipelined scf.for is immediately followed by a flat pipeline (instead of only scf.for → scf.for). When the first loop's wrap-around barrier includes a local fence, the intervening reconverge + pre-barrier + phase-shift are redundant and eliminated. 3. Gluon frontend: assert warp_pipeline_stage is inside a for loop Since the compiler now supports flat border markers, there is a risk that users place warp_pipeline_stage outside any loop, which has no meaningful pipelining semantics. A for_loop_depth counter is added to GluonSemantic and incremented/decremented in code_generator's visit_For (covering both range and static_range). warp_pipeline_stage asserts for_loop_depth > 0 at exit. The f16 GEMM example kernel is updated to use ttgl.static_range for the epilogue loop, exercising the new flat pipeline path end-to-end. Lit tests added for both WarpPipeliner (flat_pipeline_example) and ConvertWarpPipeline (flat_pipeline_backend, back_to_back_for_then_flat).
Factor out the duplicated pre-barrier + phase-shift setup and the post-pipeline reconverge logic from emitPipelinedFor and emitPipelinedFlat into shared helpers emitPipelinePrelude and emitPipelinePostlude. NFC.
Unify the duplicated pairwise dependency analysis from emitPipelinedFor (circular/wrap-around) and emitPipelinedFlat (linear) into a single analyzePipelineDependencies function parameterized by `bool circular`. NFC.
ThomasRaoux
reviewed
Apr 6, 2026
Comment on lines
+70
to
+75
| # Warp-pipelining is a loop optimization: stages must be declared | ||
| # inside a for loop (range or static_range). Allowing stages outside | ||
| # a loop would produce border markers with no well-defined iteration | ||
| # structure, breaking the phase-shift/reconvergence contract. | ||
| assert getattr(self._semantic, 'for_loop_depth', 0) > 0, ("warp_pipeline_stage must be used inside a for loop " | ||
| "(range or static_range)") |
Collaborator
There was a problem hiding this comment.
can we verify that on IR instead?
Contributor
Author
There was a problem hiding this comment.
The for_loop_depth check exists because static_range unrolls at the Python level - by the time the compiler sees the IR, there's no loop structure left, so we can't determine from the IR alone whether the code originated from a loop. For range (dynamic), we could simply check whether the parent op is scf.for, but static_range has no such anchor. Tracking depth for both loop types uniformly looked simpler than special-casing static_range alone.
For context: the compiler already correctly groups contiguous warp_pipeline_stage blocks and drops sequences with fewer than two stages, regardless of whether they came from a loop. However, warp-pipelining is only beneficial inside a loop - the phase-shift overhead amortizes over iterations. Outside a loop, it adds barrier overhead with no pipelining gain. This check is a conservative guard rail against accidental misuse, not a correctness requirement.
To summarize, potential other options are : 1) separate check for static/dynamic range, 2) don't check it.
Collaborator
There was a problem hiding this comment.
it seems weird that something works with static loops but still depends on the loop.
Contributor
Author
There was a problem hiding this comment.
You're right - if the compiler handles flat sequences correctly regardless of origin, there's no reason to enforce loop context at the frontend. The natural guard (< 2 stages → no pipeline emitted) is sufficient. I'll remove the for_loop_depth assertion.
jungpark-mlir
requested review from
antiagainst,
ptillet and
zhanglx13
as code owners
…rier exists emitPipelinedFlat unconditionally inserted a new cluster barrier (s_barrier) at every stage boundary, ignoring pre-existing barrier ops (e.g., async_wait) between execute_regions. This produced two barriers at the same boundary. Mirror the emitPipelinedFor logic: scan between consecutive stages for existing barrier ops and wrap them with sched_barriers instead of inserting a new one.
jungpark-mlir
changed the title
Two changes to analyzePipelineDependencies and eliminateRedundantCondBarriers: 1. Adjacent-stage check: the inner loop previously started at distance 2 (next = src + 2 + offset), so consecutive clusters sharing an LDS allocation never got a LOCAL barrier. Add a preliminary loop that checks clusterInfo[src] against clusterInfo[src+1] and sets bars[src+1] when they intersect. This prevents ModuleMembarAnalysis from inserting a redundant ttg.barrier local inside the pipeline. 2. Cross-pipeline analysis: when eliminating redundant cond_barriers between back-to-back pipelines, run analyzePipelineDependencies on the merged cluster sequence to verify no LDS hazard exists at the boundary. If the boundary needs a barrier (adjacent or distance-2+), the optimization is skipped. Lit tests: - back_to_back_cross_dep_kept: shared-buffer RAW at boundary → kept - back_to_back_no_dep_elimination: loop B has no LDS → eliminated - back_to_back_dep_covered_elimination: 3-stage loop A with internal barrier covering the cross-pipeline dep → eliminated - adjacent_stage_lds_dep: 3-stage pipeline verifying LOCAL barrier between adjacent stages with RAW dependency
collectNextPipelineClusters stopped at the first intervening sched_barrier / cluster barrier for flat pipelines, so only b_0 was ever visible to isCrossPipelineSafe. A cross-pipeline LDS dep involving a later flat stage (b_1, b_2, ...) was missed and the boundary cond_barrier / prelude ttg.barrier local / phase-shift cond_barrier triplet could be wrongly eliminated. Split the collection into collectLoopClusters / collectFlatClusters and walk past intra-pipeline glue (sched_barrier, s_setprio, cluster barriers, pre-existing async waits) so every flat stage is collected. Also thread B's materialized barrier flags into isCrossPipelineSafe so the merged analysis sees B's actual internal LOCAL barriers instead of relying on re-discovery from all-false placeholders. Add a lit test (@cross_pipeline_dep_in_b1) that fails without the fix.
Fold the adjacent (distance-1) and longer-distance phases of the warp-pipeline LDS dependency analysis into one loop that sweeps `dist` from 1 to maxDist. A single `wrap()` helper handles modular arithmetic, and a single `isCovered()` lambda replaces the two near-identical `isSynced` bodies for circular and linear modes. Also drop the redundant final iteration in circular mode (the old `offset == N - 1` step corresponds to `dist == 1` after wrap and only re-walked already-handled adjacent pairs). Behavior is preserved: same `(src, dst)` pairs are visited in the same order, `barrierLoc` resolves to the same slot ((dist == 1) ? dst : wrap(dst - 1)), and the coverage walk inspects the same `(src, barrierLoc]` range. No bar pattern changes. Add a thorough doc block describing the pipeline layout, the goal, the placement choice, the coverage check semantics, and the iteration order, since the conventions are easy to miss.
- Drop the stale "distance-2+ check" reference in isCrossPipelineSafe; reword to match the unified single-distance sweep in analyzePipelineDependencies. - Make emitPipelinedFor and emitPipelinedFlat use parallel section numbering (1..5) and parallel "Circular ..." / "Linear ..." headings for the dependency-analysis step. - Remove a duplicate analysis comment inside emitPipelinedFlat.
jungpark-mlir
changed the title
antiagainst
requested changes
Apr 25, 2026
antiagainst
approved these changes
Apr 30, 2026
Member
antiagainst
left a comment
antiagainst
left a comment
There was a problem hiding this comment.
Now looks good to me overall; just two final nits. @ThomasRaoux would you like to take another look?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ConvertWarpPipelineto handle sequences ofscf.execute_regionops outsidescf.for(produced byWarpPipeliner::createFlatPipeline). Emit pre-barrier, phase shift, cluster barriers, and reconverge around them.analyzePipelineDependencieson the merged sequence. Skips the optimization when a dependency is found.analyzePipelineDependencies. The existing loop only checked pairs at distance 2+, so consecutive clusters sharing an LDS allocation never got a LOCAL barrier — causingModuleMembarAnalysisto insert a redundantttg.barrier localinside the pipeline.analyzePipelineDependencies,emitPipelinePrelude/Postlude, andemitClusterBarrierhelpers.Test plan
back_to_back_cross_dep_kept: shared-buffer RAW at boundary → barriers keptback_to_back_no_dep_elimination: loop B has no LDS → barriers eliminatedback_to_back_dep_covered_elimination: 3-stage loop A with internal barrier covering the cross-pipeline dep → barriers eliminatedback_to_back_for_then_flat: pipelined loop + flat pipeline → barriers eliminatedadjacent_stage_lds_dep: 3-stage pipeline verifying LOCAL barrier between adjacent stages with RAW dependencyflat_pipeline_existing_barrier: pre-existingasync_waitwrapped withsched_barrier