Introduce letstmt defined variables in auto schedule pipeline#3
Merged
silentCoder-dev merged 7 commits intoApr 1, 2026
Merged
Conversation
Denverjin
force-pushed
the
auto-schedule
branch
from
6db9e1d to
c92bc1a
Compare
Denverjin
force-pushed
the
auto-schedule
branch
from
c92bc1a to
3bea503
Compare
…lelang into auto-schedule
Owner
There was a problem hiding this comment.
The debugging info should be removed
| } | ||
| } | ||
|
|
||
| // LOG(INFO) << "Cloned task: " << cloned_task->stmts[0]; |
Owner
There was a problem hiding this comment.
The debugging info should be removed
📝 WalkthroughWalkthroughThis PR refactors the auto-scheduling infrastructure to support shared pointer ownership semantics and variable-level dependency tracking. It switches IR node children from Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 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)
Comment |
Owner
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/transform/auto_schedule/ir_structure.cc (1)
337-349:⚠️ Potential issue | 🟠 Major
ScheduleUnit::Clone()drops scheduled side statements.Lines 337-349 clone the child and stage, but not
before/after. Any clone-based path loses the inserted prologue/epilogue statements carried by the unit and can emit different IR than the source tree.Suggested fix
std::shared_ptr<IRStructure> ScheduleUnit::Clone() const { auto new_unit = std::make_shared<ScheduleUnit>(); // Copy var and value (TVM objects with reference counting) new_unit->stage = stage; + new_unit->before = before; + new_unit->after = after; // Clone child if exists if (child) { new_unit->child = child->Clone(); } // Copy latency and II new_unit->SetLatency(GetLatency()); new_unit->SetII(GetII()); + new_unit->SetStartTime(GetStartTime()); return new_unit; }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/auto_schedule/ir_structure.cc` around lines 337 - 349, ScheduleUnit::Clone currently copies stage, child, latency and II but omits the scheduled prologue/epilogue statements (before/after), so cloned units lose inserted side statements. Update ScheduleUnit::Clone (the method) to also copy the unit's before and after fields (or call the existing accessors like GetBefore()/GetAfter() or setters like SetBefore()/SetAfter() if present) into new_unit; if the before/after are AST/TVM nodes and need deep copy use their Clone/Copy method, otherwise perform a shallow copy/pointer assignment consistent with how stage and child are handled so the cloned ScheduleUnit preserves the scheduled side statements.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/transform/auto_schedule.cc`:
- Around line 856-862: The cloned TaskNode created as cloned_task currently only
contains the cloned_let_stmt and lacks dependency metadata, so
SolveConflictVar()'s HasDependency() treats it as independent; update the clone
creation (where cloned_task is made and cloned_let_stmt is pushed) to also
populate the cloned task's dependency info by copying/deriving read/write
regions and var usages from the original node_i_let_stmt (or its parent task)
into cloned_task (e.g., read_regions, write_regions, and any var sets used by
HasDependency()); ensure stage_map[cloned_task.get()] remains set to rem_stage_j
after these metadata fields are filled so later passes see correct dependencies.
In `@src/transform/auto_schedule/ir_structure.h`:
- Around line 297-299: The current containWarpgroupId(int id) always returns
true when ContainsLoopBreak() is true, causing loop_break tasks (which should be
neutral-only and have warpgroup_id_ == -1) to be duplicated into
wg0_structure/wg1_structure; change containWarpgroupId to treat loop_break tasks
as neutral by returning false when ContainsLoopBreak() is true, otherwise return
warpgroup_id_ == id; update the containWarpgroupId implementation in the
IRStructure class so that the ContainsLoopBreak() check yields false instead of
true and confirm ApplyWarpgroupPartitionToIRStructure() will continue to emit
those tasks via the neutral clone path.
In `@src/transform/auto_schedule/memory_detector.h`:
- Around line 266-270: VisitStmt_(const LetStmtNode *op) records the let binding
before traversing the initializer, so reads inside op->value are never recorded;
change the order to traverse the initializer first (visit op->value with the
appropriate visitor), then insert into let_bindings_ and call
UpdateWriteVar(op->var), then traverse op->body via StmtExprVisitor::VisitStmt;
reference symbols: VisitStmt_, LetStmtNode, let_bindings_, UpdateWriteVar, and
StmtExprVisitor::VisitStmt.
In `@src/transform/if_condition_extract.cc`:
- Around line 59-88: The rewrite currently alters control-flow by merging inner
if conditions and by splitting SeqStmt into independently guarded statements;
fix it by (1) changing bind_cond_var so it does NOT merge conditions for nested
IfThenElseNodes but instead wraps the entire sentence under cond (i.e., return
IfThenElse(cond, sentence) rather than combining cond & if_sentence->condition),
(2) changing bind_cond_var_body so that when body is a SeqStmtNode you do NOT
iterate and wrap each element individually but instead wrap the whole SeqStmt as
a single guarded statement (preserve LetStmt/AttrStmt scope), and (3) ensure the
parent else branch is guarded with the negation of cond_var (use !cond_var when
inserting else_case) so the else executes only when cond is false; adjust uses
of bind_cond_var, bind_cond_var_body, cond_var, then_case, and else_case
accordingly.
---
Outside diff comments:
In `@src/transform/auto_schedule/ir_structure.cc`:
- Around line 337-349: ScheduleUnit::Clone currently copies stage, child,
latency and II but omits the scheduled prologue/epilogue statements
(before/after), so cloned units lose inserted side statements. Update
ScheduleUnit::Clone (the method) to also copy the unit's before and after fields
(or call the existing accessors like GetBefore()/GetAfter() or setters like
SetBefore()/SetAfter() if present) into new_unit; if the before/after are
AST/TVM nodes and need deep copy use their Clone/Copy method, otherwise perform
a shallow copy/pointer assignment consistent with how stage and child are
handled so the cloned ScheduleUnit preserves the scheduled side statements.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3e7ea5c8-bbe9-4130-8d6b-b413b4ec70ec
📒 Files selected for processing (7)
src/transform/auto_schedule.ccsrc/transform/auto_schedule/ir_structure.ccsrc/transform/auto_schedule/ir_structure.hsrc/transform/auto_schedule/memory_detector.hsrc/transform/if_condition_extract.cctilelang/engine/phase.pytilelang/transform/__init__.py
silentCoder-dev
merged commit 22ecc7b
into
silentCoder-dev:auto-schedule
3 checks passed
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
2 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 by CodeRabbit
Release Notes