[Bugfix] Support loop-dependent conditions in IfThenElse within T.Pipelined#1799
Conversation
This commit applies the same strategy used for LetStmt to IfThenElse conditions: 1.Introduced IfWrapper struct to track if conditions that depend on the loop variable 2.Added dependency detection that checks if an if condition uses: - The pipeline loop variable directly, OR - Any variable transitively dependent on the loop variable 3.Loop-dependent conditions are pushed inside each pipeline stage with the loop variable properly substituted for that iteration
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughAdds IfWrapper support to the pipeline injection flow: collects if-conditions depending on the pipeline loop variable, threads them through PipelineInjector → PipelineRewriter, and applies per-stage conditional substitutions during pipeline construction alongside existing LetWrapper handling. Changes
Sequence Diagram(s)sequenceDiagram
participant Inj as PipelineInjector
participant Coll as TraversalLogic
participant Rw as PipelineRewriter
participant Build as BuildPipeline
Inj->>Coll: Analyze For-node body (collect LetWrapper & IfWrapper)
Coll->>Coll: Identify IfThenElse nodes whose conditions depend on loop var
Coll->>Coll: Track transitive LetStmt deps and group wrappers
Coll->>Inj: Return loop_var_if_wrappers vector
Inj->>Rw: Construct PipelineRewriter with loop_var_if_wrappers
Rw->>Rw: Store loop_var_if_wrappers_ member
Rw->>Build: Rewrite loop body per stage
Build->>Build: Substitute IfWrapper conditions per stage context
Build->>Build: Apply LetWrapper substitutions per stage
Build->>Rw: Return transformed stages
Rw->>Inj: Return rewritten pipeline
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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
🤖 Fix all issues with AI agents
In `@src/transform/inject_pipeline.cc`:
- Around line 872-885: The current code applies let-wrappers
(loop_var_let_wrappers_) then if-wrappers (loop_var_if_wrappers_) separately
which reverses nesting and frees let-bound vars used by if conditions; replace
these two separate vectors with a single ordered vector of variant wrappers
(e.g., CondWrapper holding either LetWrapper or IfWrapper, stored in
loop_var_wrappers_) and, when injecting, iterate loop_var_wrappers_ in reverse
and apply each element by type (apply let semantics for LetWrapper and create
IfThenElse with Substitute on iw.condition using pipeline_loop_->loop_var and
normalized_access_index for IfWrapper) so the original source nesting is
preserved and conditions see their let-bound vars. Ensure you update usages of
loop_var_if_wrappers_ and loop_var_let_wrappers_ to the unified
loop_var_wrappers_ and maintain the existing Substitute logic and iw.span
handling.
🧹 Nitpick comments (1)
src/transform/inject_pipeline.cc (1)
1106-1131: Transitive dependency detection looks correct.The logic properly detects conditions that depend on the loop variable either directly or transitively through let-bound variables. This mirrors the existing LetStmt handling.
Consider extracting the duplicate
dependent_varsconstruction (also at lines 1144-1148) into a helper lambda for maintainability:♻️ Optional refactor
+ auto build_dependent_vars = [&]() { + std::unordered_set<const VarNode *> vars; + vars.insert(op->loop_var.get()); + for (const auto &lw : loop_var_let_wrappers) { + vars.insert(lw.var.get()); + } + return vars; + }; + // In IfThenElse handling: - std::unordered_set<const VarNode *> dependent_vars; - dependent_vars.insert(op->loop_var.get()); - for (const auto &lw : loop_var_let_wrappers) { - dependent_vars.insert(lw.var.get()); - } + auto dependent_vars = build_dependent_vars();
| // Similarly, handle If-wrappers whose conditions depend on the | ||
| // pipeline loop var. | ||
| if (!loop_var_if_wrappers_.empty()) { | ||
| BlockNode *n = new_block.CopyOnWrite(); | ||
| Stmt inner = n->body; | ||
| for (auto it = loop_var_if_wrappers_.rbegin(); | ||
| it != loop_var_if_wrappers_.rend(); ++it) { | ||
| const auto &iw = *it; | ||
| PrimExpr substituted_condition = Substitute( | ||
| iw.condition, {{pipeline_loop_->loop_var, normalized_access_index}}); | ||
| inner = IfThenElse(substituted_condition, inner, Stmt(), iw.span); | ||
| } | ||
| n->body = inner; | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's examine the actual code at the mentioned lines
cat -n src/transform/inject_pipeline.cc | sed -n '850,900p'Repository: tile-ai/tilelang
Length of output: 2522
🏁 Script executed:
# Also check what's before line 872 to understand let-wrapper handling
cat -n src/transform/inject_pipeline.cc | sed -n '859,885p'Repository: tile-ai/tilelang
Length of output: 1380
🏁 Script executed:
# Search for the data structures used: loop_var_if_wrappers_, loop_var_let_wrappers_
rg -n "loop_var_if_wrappers_|loop_var_let_wrappers_" src/transform/inject_pipeline.cc | head -30Repository: tile-ai/tilelang
Length of output: 622
🏁 Script executed:
# Find the test case mentioned
find . -name "test_tilelang_issue_1263.py" -o -name "*issue_1263*"Repository: tile-ai/tilelang
Length of output: 110
🏁 Script executed:
# Search for dependent_vars logic around line 1115
cat -n src/transform/inject_pipeline.cc | sed -n '1100,1130p'Repository: tile-ai/tilelang
Length of output: 1783
🏁 Script executed:
# Check if there are any guards or conditions that prevent dependent conditions
rg -n "dependent" src/transform/inject_pipeline.cc -B 2 -A 2Repository: tile-ai/tilelang
Length of output: 2627
🏁 Script executed:
cat -n testing/python/issue/test_tilelang_issue_1263.pyRepository: tile-ai/tilelang
Length of output: 2370
🏁 Script executed:
# Search for where loop_var_if_wrappers and loop_var_let_wrappers are used together
rg -n "loop_var_if_wrappers|loop_var_let_wrappers" src/transform/inject_pipeline.cc -B 1 -A 1 | head -80Repository: tile-ai/tilelang
Length of output: 2838
🏁 Script executed:
# Check if there's any special handling to ensure condition substitution happens correctly
rg -n "substituted_condition" src/transform/inject_pipeline.cc -B 3 -A 3Repository: tile-ai/tilelang
Length of output: 526
🏁 Script executed:
# Check LetWrapper and IfWrapper structure definitions
rg -n "struct LetWrapper|struct IfWrapper" src/transform/inject_pipeline.cc -A 5Repository: tile-ai/tilelang
Length of output: 207
🏁 Script executed:
# Check if there's any post-processing of conditions or variables after substitution
rg -n "UsesVar|HasVar|GetVars" src/transform/inject_pipeline.cc -B 2 -A 2Repository: tile-ai/tilelang
Length of output: 550
🏁 Script executed:
# Check how Substitute is implemented and what it does with nested variables
rg -n "class Substitute|struct Substitute" --type cpp | head -5Repository: tile-ai/tilelang
Length of output: 42
🏁 Script executed:
# Look for any variable renaming or binding adjustments in the pipeline building
rg -n "RenameVar|VarReplacer|VarSubstitute" src/transform/inject_pipeline.cc -B 2 -A 2Repository: tile-ai/tilelang
Length of output: 42
Fix variable scoping when if-conditions depend on let-bound variables.
The sequential application of let-wrappers (lines 859-870) followed by if-wrappers (lines 872-885) creates incorrect nesting. When a condition depends on let-bound variables (detected at line 1115), the final structure becomes:
IfThenElse(id2 > 1,
LetStmt(id2, ids2[id],
LetStmt(id, ids[i], body)))
Here, the condition references id2 which is bound inside the then-branch, and id2's value references id which is bound even deeper. Both become free variables. The original structure should preserve:
LetStmt(id, ids[i],
LetStmt(id2, ids2[id],
IfThenElse(id2 > 1, body)))
The issue affects test case _test_kernel_if_cond at line 43 where if id2 > 1 depends on transitively loop-dependent variables.
Use a unified wrapper vector of variant type applied in collection order to preserve original nesting:
🔧 Suggested approach
struct CondWrapper {
std::variant<LetWrapper, IfWrapper> data;
};
std::vector<CondWrapper> loop_var_wrappers_;
// Apply in reverse to preserve original nesting
for (auto it = loop_var_wrappers_.rbegin(); it != loop_var_wrappers_.rend(); ++it) {
if (auto* lw = std::get_if<LetWrapper>(&it->data)) {
// apply let
} else if (auto* iw = std::get_if<IfWrapper>(&it->data)) {
// apply if
}
}🤖 Prompt for AI Agents
In `@src/transform/inject_pipeline.cc` around lines 872 - 885, The current code
applies let-wrappers (loop_var_let_wrappers_) then if-wrappers
(loop_var_if_wrappers_) separately which reverses nesting and frees let-bound
vars used by if conditions; replace these two separate vectors with a single
ordered vector of variant wrappers (e.g., CondWrapper holding either LetWrapper
or IfWrapper, stored in loop_var_wrappers_) and, when injecting, iterate
loop_var_wrappers_ in reverse and apply each element by type (apply let
semantics for LetWrapper and create IfThenElse with Substitute on iw.condition
using pipeline_loop_->loop_var and normalized_access_index for IfWrapper) so the
original source nesting is preserved and conditions see their let-bound vars.
Ensure you update usages of loop_var_if_wrappers_ and loop_var_let_wrappers_ to
the unified loop_var_wrappers_ and maintain the existing Substitute logic and
iw.span handling.
2 participants
This PR extends the pipeline injection transform to handle
IfThenElsestatements whose conditions depend on the pipeline loop variable, fixing a regression similar to #1210 but for conditional statements.Problem
When an
ifstatement with a condition depending on the pipeline loop variable (e.g.,if i > 1:) was placed inside aT.Pipelinedblock, the condition may be incorrectly hoisted outside the pipeline loop, causing the loop variable to become undefined and resulting in the error:The
IfThenElsehandling ininject_pipeline.ccalways added conditions torewrap_fns(wrapping them outside the pipeline), without checking if the condition uses the loop variable. This is the same bug that was fixed forLetStmtin PR #1212, but was not addressed for conditional statementsExample:
Before this PR, above code would be transformed as:
Solution
This PR applies the same strategy used for
LetStmttoIfThenElseconditions:IfWrapperstruct to track if conditions that depend on the loop variableKey Changes
inject_pipeline.cc:IfWrapperstruct to store if conditions and spansPipelineRewriterto accept and processloop_var_if_wrappersIfThenElsedetection logic to check for loop variable dependenciesRelated Issues/PRs:
#1210
#1212
#1263
Summary by CodeRabbit
New Features
Tests