Skip to content

Letstmt analysis bug fix #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
silentCoder-dev merged 1 commit into from
Apr 1, 2026
Merged

Letstmt analysis bug fix #6

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 17 additions & 17 deletions src/transform/auto_schedule.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,16 +434,16 @@ bool AssignWarpgroupIdsGlobal(IRStructure *root, bool enable_warp_partition) {
// Builder that collects ScheduleUnits from IRStructure
class ScheduleUnitBuilder {
public:
bool Build(IRStructure *root) {
bool Build(std::shared_ptr<IRStructure> &root) {
ScheduleRecursive(root, {});

// Global warpgroup id assignment from the top level
return AssignWarpgroupIdsGlobal(root, enable_warp_partition_);
return AssignWarpgroupIdsGlobal(root.get(), enable_warp_partition_);
}

// New recursive scheduling function that replaces Collect method
// Directly schedules the entire IRStructure tree recursively in place
void ScheduleRecursive(IRStructure *node,
void ScheduleRecursive(std::shared_ptr<IRStructure> &node,
const std::set<Buffer> &used_buffers);

// Z3-based scheduler that calls Python implementation via FFI
Expand Down Expand Up @@ -1172,7 +1172,7 @@ class ScheduleUnitBuilder {

// Implementation of ScheduleRecursive function
void ScheduleUnitBuilder::ScheduleRecursive(
IRStructure *node, const std::set<Buffer> &used_buffers) {
std::shared_ptr<IRStructure> &node, const std::set<Buffer> &used_buffers) {
if (!node)
return;

Expand Down Expand Up @@ -1235,7 +1235,7 @@ void ScheduleUnitBuilder::ScheduleRecursive(
// TaskNode: no further scheduling needed
return;
} else if (node->IsSequence()) {
auto seq = static_cast<SequenceNode *>(node);
auto seq = static_cast<SequenceNode *>(node.get());

// First, recursively schedule all children
std::vector<std::shared_ptr<IRStructure>> seq_children, origin_children;
Expand All @@ -1252,13 +1252,13 @@ void ScheduleUnitBuilder::ScheduleRecursive(
}
}
}
ScheduleRecursive(child.get(), child_used_buffers);
ScheduleRecursive(child, child_used_buffers);
}

seq->children = ChildrenScheduleHelper(origin_children);
return;
} else if (node->IsControl()) {
auto ctrl = static_cast<ControlNode *>(node);
auto ctrl = static_cast<ControlNode *>(node.get());

// Now schedule the ControlNode's internal tasks (if any) as a unit
// The body should now be a SequenceNode containing the tasks
Expand All @@ -1279,13 +1279,13 @@ void ScheduleUnitBuilder::ScheduleRecursive(
}
}
}
ScheduleRecursive(child.get(), child_used_buffers);
ScheduleRecursive(child, child_used_buffers);
}
Z3SchedulePythonLoop(ctrl, used_buffers);
} else if (ctrl->child->IsWrapper()) {
auto wrapper = static_cast<WrapperNode *>(ctrl->child.get());
std::vector<std::shared_ptr<IRStructure>> origin_children;
GatherTaskNodes({wrapper->child}, origin_children);
GatherTaskNodes({wrapper->task, wrapper->child}, origin_children);

@coderabbitai coderabbitai Bot Apr 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Use the wrapper-aware gather helper here.

GatherTaskNodes assumes every input entry is non-null, but WrapperNode::child is optional in this file's builder (see Lines 1508-1511 and 1526-1529). Constructing {wrapper->task, wrapper->child} can therefore feed a null shared_ptr into GatherTaskNodes and crash on the first dereference. Reusing GatherTaskNodesSingle(...) keeps the existing wrapper/null handling intact.

🛠️ Suggested fix 
-        GatherTaskNodes({wrapper->task, wrapper->child}, origin_children);
+        GatherTaskNodesSingle(ctrl->child, origin_children);
-    GatherTaskNodes({wrapper->task, wrapper->child}, origin_children);
+    GatherTaskNodesSingle(node, origin_children);

Also applies to: 1312-1312

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/transform/auto_schedule.cc` at line 1288, The call to
GatherTaskNodes({wrapper->task, wrapper->child}, origin_children) can pass a
null shared_ptr from WrapperNode::child and lead to a crash; replace it with the
wrapper-aware helper GatherTaskNodesSingle so null children are handled safely
(e.g., call GatherTaskNodesSingle(wrapper->task, wrapper->child,
origin_children)), and make the same replacement for the other occurrence that
mirrors this pattern (the second call around the later occurrence).

for (auto &child : origin_children) {
auto child_used_buffers = used_buffers;
for (auto &other_child : origin_children) {
Expand All @@ -1298,18 +1298,18 @@ void ScheduleUnitBuilder::ScheduleRecursive(
}
}
}
ScheduleRecursive(child.get(), child_used_buffers);
ScheduleRecursive(child, child_used_buffers);
}
Z3SchedulePythonLoop(ctrl, used_buffers);
} else {
ScheduleRecursive(ctrl->child.get(), used_buffers);
ScheduleRecursive(ctrl->child, used_buffers);
}
}
return;
} else if (node->IsWrapper()) {
auto wrapper = static_cast<WrapperNode *>(node);
auto wrapper = static_cast<WrapperNode *>(node.get());
std::vector<std::shared_ptr<IRStructure>> origin_children;
GatherTaskNodes({wrapper->child}, origin_children);
GatherTaskNodes({wrapper->task, wrapper->child}, origin_children);
for (auto &child : origin_children) {
auto child_used_buffers = used_buffers;
for (auto &other_child : origin_children) {
Expand All @@ -1322,15 +1322,15 @@ void ScheduleUnitBuilder::ScheduleRecursive(
}
}
}
ScheduleRecursive(child.get(), child_used_buffers);
ScheduleRecursive(child, child_used_buffers);
}
auto seq_node = std::make_shared<SequenceNode>();
seq_node->children = ChildrenScheduleHelper(origin_children);
*node = *seq_node;
node = seq_node;
return;
}

LOG(FATAL) << "[ScheduleRecursive] Unknown IRStructure type" << node;
LOG(FATAL) << "[ScheduleRecursive] Unknown IRStructure type" << node.get();
}

// Mutator to update thread extent in AttrStmt nodes
Expand Down Expand Up @@ -1878,7 +1878,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) {
}
unit_builder.SetEnableWarpPartition(config.enable_warp_partition);
unit_builder.SetSharedMemoryLimit(config.shared_memory_limit);
bool double_thread = unit_builder.Build(ir_structure.get());
bool double_thread = unit_builder.Build(ir_structure);

if (!config.enable_warpgroup_partition) {
Stmt new_body = ConvertIRStructureToStmt(ir_structure.get(), enable_epi);
Expand Down