[BugFix][Transform] Close transitive pipeline dependencies - #2822
[BugFix][Transform] Close transitive pipeline dependencies#2822JayceSu98 wants to merge 1 commit into
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
📝 WalkthroughWalkthroughPipeline dependency construction now filters disjoint buffer regions. Pipeline planning uses unified transitive propagation across buffer and scalar dependencies, preserves direct copy anchors, rejects TMA transitive producers, and adds regression coverage. ChangesPipeline dependency planning
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PipelinePlanner
participant CopyAnchorAnalysis
participant ProducerWorklist
participant InjectSoftwarePipeline
participant CUDA
PipelinePlanner->>CopyAnchorAnalysis: analyze copy stages and direct consumers
CopyAnchorAnalysis->>ProducerWorklist: seed direct copy anchors
ProducerWorklist->>ProducerWorklist: traverse buffer and scalar dependencies
ProducerWorklist->>InjectSoftwarePipeline: assign stages and producer groups
InjectSoftwarePipeline->>CUDA: lower the sm_80 regression kernel
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
testing/python/transform/test_tilelang_transform_pipeline_planning.py (1)
849-855: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider relational assertions instead of exact order lists.
orders == [3, 0, 1, 2, 4]pins the exact rotation chosen by the planner, so any future (still-legal) reordering breaks the test even though the property under test — the disjoint write is not pulled into the copy's producer group — still holds.test_pipeline_planning_resolves_each_multi_root_anchor_before_minimumalready demonstrates the more durable style (orders[1] < orders[2]). Keeping the exactstagescheck plus relational order checks would preserve intent with less churn.Based on learnings, transform tests should assert structural/behavioral patterns rather than exact literals.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_pipeline_planning.py` around lines 849 - 855, Update the assertions in the test around _collect_pipeline_loop_annotations to keep the exact stages validation but replace the full orders list comparison with relational checks that enforce the disjoint write remains outside the copy’s producer group, following the durable ordering style used by test_pipeline_planning_resolves_each_multi_root_anchor_before_minimum.Source: Learnings
src/transform/inject_pipeline.cc (1)
529-549: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRegion-disjointness filtering looks correct and conservative.
Different data vars → no edge, same data var with distinct
Bufferviews or differing ranks → conservative conflict, per-dimIntSetintersection only prunes provably disjoint ranges. Note thatarith::Intersectwithout anAnalyzerwill simply fail to prove disjointness for symbolic ranges, which is the safe direction here.Minor DRY note: this is now the third near-identical region-overlap helper (
PipelineRewriter::MayConflictat Line 1471 in this file, andMayConflictinsrc/transform/pipeline_planning.cc). Consolidating into one shared helper would keep the three copies from diverging as the region semantics evolve.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/transform/inject_pipeline.cc` around lines 529 - 549, Consolidate the duplicated region-overlap logic in RegionsMayConflict, PipelineRewriter::MayConflict, and the MayConflict helper in pipeline_planning.cc into one shared helper. Update each caller to reuse it while preserving the current conservative behavior for distinct buffers, differing ranks or sizes, and symbolic ranges.testing/python/issue/test_tilelang_issue_2720.py (1)
6-19: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptionally assert the resulting stage/order, not just that lowering succeeds.
The linked issue asks the regression to verify "valid producer ordering and stages". As written this only guards against an exception; a schedule regression that still lowers would pass.
test_pipeline_planning_propagates_buffer_producer_through_bindcovers the annotations, so this is a nice-to-have here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/issue/test_tilelang_issue_2720.py` around lines 6 - 19, Strengthen the regression around _producer_bind_kernel by asserting the lowered result’s producer ordering and pipeline stage assignments, rather than only verifying that lowering succeeds. Reuse the existing lowering or inspection helpers and expected annotations covered by test_pipeline_planning_propagates_buffer_producer_through_bind, while keeping this test focused on the resulting schedule.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@testing/python/issue/test_tilelang_issue_2720.py`:
- Around line 22-25: Add the tilelang.testing.requires_cuda decorator to
test_pipeline_planning_tracks_transitive_copy_index_producer so this sm_80
target lowering test is skipped in CPU-only environments, matching the gating
used for issue 2307.
---
Nitpick comments:
In `@src/transform/inject_pipeline.cc`:
- Around line 529-549: Consolidate the duplicated region-overlap logic in
RegionsMayConflict, PipelineRewriter::MayConflict, and the MayConflict helper in
pipeline_planning.cc into one shared helper. Update each caller to reuse it
while preserving the current conservative behavior for distinct buffers,
differing ranks or sizes, and symbolic ranges.
In `@testing/python/issue/test_tilelang_issue_2720.py`:
- Around line 6-19: Strengthen the regression around _producer_bind_kernel by
asserting the lowered result’s producer ordering and pipeline stage assignments,
rather than only verifying that lowering succeeds. Reuse the existing lowering
or inspection helpers and expected annotations covered by
test_pipeline_planning_propagates_buffer_producer_through_bind, while keeping
this test focused on the resulting schedule.
In `@testing/python/transform/test_tilelang_transform_pipeline_planning.py`:
- Around line 849-855: Update the assertions in the test around
_collect_pipeline_loop_annotations to keep the exact stages validation but
replace the full orders list comparison with relational checks that enforce the
disjoint write remains outside the copy’s producer group, following the durable
ordering style used by
test_pipeline_planning_resolves_each_multi_root_anchor_before_minimum.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 000b9710-a70d-4f82-a1ef-69e29bb181d5
📒 Files selected for processing (4)
src/transform/inject_pipeline.ccsrc/transform/pipeline_planning.cctesting/python/issue/test_tilelang_issue_2720.pytesting/python/transform/test_tilelang_transform_pipeline_planning.py
| def test_pipeline_planning_tracks_transitive_copy_index_producer(): | ||
| target = tilelang.tvm.target.Target({"kind": "cuda", "arch": "sm_80"}) | ||
| with target: | ||
| tilelang.lower(_producer_bind_kernel, target=target) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# How do other lowering-only tests in testing/python/issue handle CUDA gating?
fd -t f 'test_tilelang_issue_.*\.py' testing/python/issue --exec sh -c '
if rg -qn "tilelang\.lower\(" "$1"; then
echo "=== $1"
rg -n "requires_cuda|requires_cuda_compute_version|tilelang\.lower\(|Target\(" "$1" | head -20
fi
' _ {}Repository: tile-ai/tilelang
Length of output: 526
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== target test files"
fd -t f 'test_tilelang_issue_.*\.py' testing/python/issue | sed -n '1,80p'
echo
echo "=== issue 2307 file"
sed -n '1,80p' testing/python/issue/test_tilelang_issue_2307.py
echo
echo "=== issue 2720 file"
sed -n '1,80p' testing/python/issue/test_tilelang_issue_2720.py
echo
echo "=== locate testing decorators and tilelang lower definitions"
rg -n "def requires_cuda|requires_cuda_compute_version|class requires_cuda|tilelang\.lower|def lower" testing python tilelang -S | sed -n '1,200p'Repository: tile-ai/tilelang
Length of output: 26942
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== requires decorators"
sed -n '60,130p' tilelang/testing/__init__.py
echo
echo "=== lower public function"
sed -n '245,355p' tilelang/engine/lower.py
echo
echo "=== call graph-ish: usages of lower_to_host_device_ir"
rg -n "lower_to_host_device_ir|CompileDevice|compile_device|requires_cuda\\(" tilelang/testing tilelang/engine tilelang/transform tilelang/cuda -S | sed -n '1,240p'
echo
echo "=== focused tests around copy indexing and pipeline/producer consumer"
rg -n "index_shared|producer_bind|Pipelined|transitive|copy_index|requires_cuda\\(" testing/python/issue/test_tile*issue*2*.py testing/python/components test_tilelang_issue_*.py tilelang -S | sed -n '1,160p'Repository: tile-ai/tilelang
Length of output: 9753
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== full testing requires definitions"
sed -n '1,170p' tilelang/testing/__init__.py
echo
echo "=== lower.py helpers"
sed -n '1,270p' tilelang/engine/lower.py
echo
echo "=== device codegen helpers"
rg -n "def resolve_device_codegen|class DeviceCodeGen|class .*CUDA|resolve_pipeline|def .*Pipeline|device_codegen_without_compile" tilelang -S
echo
echo "=== targeted file slice for relevant issue tests"
python3 - <<'PY'
from pathlib import Path
for p in sorted(Path("testing/python/issue").glob("test_tilelang_issue_*.py")):
txt=p.read_text()
if "tilelang.lower" in txt:
print(f"\n== {p}")
for i,line in enumerate(txt.splitlines(),1): line=line.strip()
if "requires_cuda" in line or "tilelang.lower" in line or 'Target(' in line or "Target(" in line:
print(f"{i}: {line}")
PYRepository: tile-ai/tilelang
Length of output: 19210
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== lowering path exact"
sed -n '259,342p' tilelang/engine/lower.py
echo
echo "=== target resolution for explicit cuda dict target"
rg -n "def determine_target\(|kind.*name|target_is_cuda|cuda" tilelang/backend/target.py tilelang/cuda/target.py -S
echo
echo "=== read target resolution file"
sed -n '1,140p' tilelang/backend/target.py
echo
echo "=== read cuda target detection"
sed -n '1,120p' tilelang/cuda/target.py
echo
echo "=== parse CUDA issue lower-only tests for gating"
python3 - <<'PY'
from pathlib import Path
for p in sorted(Path("testing/python/issue").glob("test_tilelang_issue_*.py")):
txt = p.read_text()
if "tilelang.lower" in txt:
lines = txt.splitlines()
print(f"\n== {p}")
for i, line in enumerate(lines, 1):
if any(s in line for s in ("requires_cuda", "tilelang.lower", "Target(")):
print(f"{i}: {line}")
PYRepository: tile-ai/tilelang
Length of output: 14024
Add CUDA gating to this lowered CUDA lower call.
This uses a sm_80 CUDA device target and calls tilelang.lower, which runs target-specific pipeline lowering before the default device codegen path. Mirror issue 2307 and protect the test with @tilelang.testing.requires_cuda so CPU-only CI excludes it.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@testing/python/issue/test_tilelang_issue_2720.py` around lines 22 - 25, Add
the tilelang.testing.requires_cuda decorator to
test_pipeline_planning_tracks_transitive_copy_index_producer so this sm_80
target lowering test is skipped in CPU-only environments, matching the gating
used for issue 2307.
cf92691 to
e09b252
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/transform/inject_pipeline.cc (1)
529-549: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConservative filtering looks correct; consider reusing one region-intersection helper.
The disjointness test only drops an edge when
arith::Intersect(...).IsNothing()proves no overlap, so symbolic/unprovable cases stay conservative — good. However this per-dimension loop duplicatesPipelineRewriter::MayConflict(Line 1471 in this file) andtl::MayConflictinsrc/transform/pipeline_planning.cc. Extracting one shared free helper would avoid the three copies drifting apart.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/transform/inject_pipeline.cc` around lines 529 - 549, Extract the per-dimension region-overlap logic from RegionsMayConflict, PipelineRewriter::MayConflict, and tl::MayConflict into one shared free helper. Update all three call sites to reuse it while preserving conservative behavior: return false only when disjointness is proven, and retain conflicts for symbolic or unprovable intersections.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/transform/pipeline_planning.cc`:
- Around line 654-666: Introduce a shared buffer-conflict predicate matching
buffers by their ->data Var, treating differing views or region ranks as
potentially conflicting before calling MayConflict, which must only receive
equal-rank regions. In src/transform/pipeline_planning.cc lines 654-666, replace
both same_as comparisons in the RAW and WAW find_if predicates with this
predicate; likewise update FindBufferProducerIndices at lines 704-705 to use it,
while preserving the existing conflict handling.
- Around line 754-766: The relax_producer path must avoid aborting on transitive
TMA-copy producers and instead fall back to a legal non-pipelined schedule for
the current loop. Replace the ICHECK in the relax_producer lambda with the
existing fallback behavior used by the ROCm path around num_stages, ensuring
pipelining is disabled and planning continues without assigning an illegal TMA
producer chain.
---
Nitpick comments:
In `@src/transform/inject_pipeline.cc`:
- Around line 529-549: Extract the per-dimension region-overlap logic from
RegionsMayConflict, PipelineRewriter::MayConflict, and tl::MayConflict into one
shared free helper. Update all three call sites to reuse it while preserving
conservative behavior: return false only when disjointness is proven, and retain
conflicts for symbolic or unprovable intersections.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 13298fe5-3057-448e-9278-c936ec01c7d7
📒 Files selected for processing (4)
src/transform/inject_pipeline.ccsrc/transform/pipeline_planning.cctesting/python/issue/test_tilelang_issue_2720.pytesting/python/transform/test_tilelang_transform_pipeline_planning.py
| return r->buffer.same_as(read->buffer) && | ||
| MayConflict(r->region, read->region); | ||
| }) != pinfo.writes.end()) { | ||
| pinfo.last_use_stmt_index = std::max(pinfo.last_use_stmt_index, i); | ||
| } | ||
| } | ||
|
|
||
| if (!pinfo.IsCopyStage()) { | ||
| continue; | ||
| } | ||
|
|
||
| for (const BufferRegion &write : (*pipeline_stage_infos)[i].writes) { | ||
| if (std::find_if(pinfo.writes.begin(), pinfo.writes.end(), | ||
| [&](const BufferRegion &r) { | ||
| return r->buffer == write->buffer && | ||
| return r->buffer.same_as(write->buffer) && | ||
| MayConflict(r->region, write->region); | ||
| }) != pinfo.writes.end()) { |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Buffer conflicts are keyed on Buffer::same_as in both the last-use analysis and the producer worklist. Two distinct Buffer objects that alias one data Var compare as non-conflicting, so RAW edges can be dropped (and the WAW LOG(FATAL) guard bypassed) at both sites. src/transform/inject_pipeline.cc Lines 531-540 keeps such aliases conservatively conflicting, so the two passes now disagree about the same IR. Introduce one shared predicate that matches on ->data, returns "may conflict" when the views or region ranks differ, and only then falls through to MayConflict (which ICHECKs equal ranks).
src/transform/pipeline_planning.cc#L654-L666: replace bothsame_ascomparisons in the RAW and WAWfind_ifpredicates with the shared predicate.src/transform/pipeline_planning.cc#L704-L705: replace thewrite->buffer.same_as(read->buffer) && MayConflict(...)condition inFindBufferProducerIndiceswith the same shared predicate.
📍 Affects 1 file
src/transform/pipeline_planning.cc#L654-L666(this comment)src/transform/pipeline_planning.cc#L704-L705
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/transform/pipeline_planning.cc` around lines 654 - 666, Introduce a
shared buffer-conflict predicate matching buffers by their ->data Var, treating
differing views or region ranks as potentially conflicting before calling
MayConflict, which must only receive equal-rank regions. In
src/transform/pipeline_planning.cc lines 654-666, replace both same_as
comparisons in the RAW and WAW find_if predicates with this predicate; likewise
update FindBufferProducerIndices at lines 704-705 to use it, while preserving
the existing conflict handling.
| auto relax_producer = [&](int producer_idx, int root_anchor) { | ||
| auto &producer = (*pipeline_stage_infos)[producer_idx]; | ||
| ICHECK(!producer.IsTmaCopy()) | ||
| << "Pipeline planning does not support a TMA copy as a transitive " | ||
| "producer of another pipeline copy because a stage-0 TMA " | ||
| "consumer requires an explicit wait"; | ||
| if (!producer.IsCopyStage()) { | ||
| producer.producer_for_copy = true; | ||
| } | ||
| bool changed = false; | ||
| if (!producer->producer_for_copy) { | ||
| producer->producer_for_copy = true; | ||
| producer->last_use_stmt_index = consumer_last_use; | ||
| changed = true; | ||
| } else if (!producer->IsLastUseStmtIndexValid() || | ||
| consumer_last_use < producer->last_use_stmt_index) { | ||
| producer->last_use_stmt_index = consumer_last_use; | ||
| changed = true; | ||
| if (anchor_candidates[producer_idx].insert(root_anchor).second) { | ||
| worklist.emplace_back(producer_idx, root_anchor); | ||
| } | ||
| return changed; | ||
| }; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Hard ICHECK abort instead of the graceful fallback the linked issue asks for.
Issue #2720 requires falling back to a non-pipelined (or otherwise legal) schedule when no legal pipelined assignment can be constructed. A transitive TMA producer chain currently aborts the whole compilation, so any kernel that happens to hit this shape becomes uncompilable rather than degrading to a sequential loop. The worklist has enough information here to instead bail out of pipelining for this loop (as the ROCm path at Lines 1140-1156 already does by stripping num_stages).
Want me to sketch that fallback, or open a follow-up issue to track it?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/transform/pipeline_planning.cc` around lines 754 - 766, The
relax_producer path must avoid aborting on transitive TMA-copy producers and
instead fall back to a legal non-pipelined schedule for the current loop.
Replace the ICHECK in the relax_producer lambda with the existing fallback
behavior used by the ROCm path around num_stages, ensuring pipelining is
disabled and planning continues without assigning an illegal TMA producer chain.
Unify scalar and buffer producer discovery in a dependency worklist, resolve multi-root scheduling anchors without merging async copy groups, and conservatively reject unsupported TMA producer chains. Align software-pipeline dependency validation with region-aware same-buffer hazards and add fixed-point, copy-chain, multi-root, TMA, and issue regressions. Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
e09b252 to
7c81745
Compare
1 participant
Fixes #2720
Summary
PipelinePlanningpropagated buffer and scalar producers separately. A mixed chain such asbuffer store -> scalar Bind -> copystopped at the Bind, leaving the buffer store in a later stage.InjectSoftwarePipelinethen rejected the schedule.Planning now uses one backward worklist across scalar use-def and buffer RAW edges. It keeps every copy-root candidate until the anchors are resolved, then selects the earliest scheduling sink. The planner and injector both need this dependency information, so region precision is also kept during injection instead of treating every write through the same Buffer as overlapping.
Changes
InjectSoftwarePipelineretain individual write regions so disjoint regions of the same Buffer do not create false dependency edges; distinct alias views remain conservative.Review Notes
(statement, root anchor)pair is enqueued at most once.last_use_stmt_indexis the resolved scheduling anchor.direct_copy_last_use_stmt_indexremains the async commit-group anchor.Validation
testing/python/issue/test_tilelang_issue_pipeline_transitive_copy_dependencies.pylowers successfully for CUDAsm_80.clang-format, Python compilation, andgit diff --checkpassed.