diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index 225fcc19bd..3b25a8222b 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -69,6 +69,38 @@ namespace tl { using namespace tir; using ffi::GetRef; +// Extract all sequencial task nodes from the IR structure tree +void GatherTaskNodesSingle( + const std::shared_ptr &node, + std::vector> &task_nodes); +void GatherTaskNodes(const std::vector> &nodes, + std::vector> &task_nodes) { + for (const auto &node : nodes) { + if (node->IsTask()) { + task_nodes.emplace_back(node); + } else if (node->IsSequence()) { + auto seq = static_cast(node.get()); + GatherTaskNodes(seq->children, task_nodes); + } else if (node->IsWrapper()) { + auto wrapper = static_cast(node.get()); + if (wrapper->task) + task_nodes.emplace_back(wrapper->task); + if (wrapper->child) + GatherTaskNodesSingle(wrapper->child, task_nodes); + } else if (node->IsControl()) { + task_nodes.emplace_back(node); + } else { + LOG(FATAL) << "Unknown node type in GatherTaskNodes"; + } + } +} + +void GatherTaskNodesSingle( + const std::shared_ptr &node, + std::vector> &task_nodes) { + return GatherTaskNodes({node}, task_nodes); +} + // Forward declaration Stmt ApplyWarpgroupPartitionToIRStructure( IRStructure *root, IterVar thread_var, std::vector &barrier_buffers, @@ -576,6 +608,22 @@ class ScheduleUnitBuilder { // Z3-based scheduler for loops that calls Python implementation via FFI // with distance-aware dependencies void Z3SchedulePythonLoop(ControlNode *ctrl) { + if (ctrl->child == nullptr) { + LOG(WARNING) + << "Z3SchedulePythonLoop called on a control node without child"; + return; + } + std::vector> flat_children; + if (!ctrl->child->IsSequence()) { + GatherTaskNodesSingle(ctrl->child, flat_children); + } else { + auto seq_node = static_cast(ctrl->child.get()); + GatherTaskNodes(seq_node->children, flat_children); + } + auto seq_node = std::make_shared(); + seq_node->children = flat_children; + ctrl->child = std::move(seq_node); + auto seq_body = static_cast(ctrl->child.get()); std::vector nodes; nodes.reserve(seq_body->children.size()); @@ -752,7 +800,7 @@ class ScheduleUnitBuilder { node_to_index[nodes[i]] = i; } - std::vector> reordered_children; + std::vector> reordered_children; reordered_children.reserve(sorted_nodes.size()); for (IRStructure *sorted_node : sorted_nodes) { auto it = node_to_index.find(sorted_node); @@ -765,8 +813,87 @@ class ScheduleUnitBuilder { seq_body->children = std::move(reordered_children); } + // Reorder & Copy Let-defined variables + auto IsVarDecl = [&](IRStructure *node) -> bool { + if (!node) + return false; + + if (node->IsTask()) { + auto task = static_cast(node); + if (task->stmts.size() == 1) { + return task->stmts[0].as() != nullptr; + } + } + return false; + }; + auto SolveConflictVar = [&]() -> bool { + for (int i = 0; i < n; ++i) + if (IsVarDecl(seq_body->children[i].get())) { + for (int j = 0; j < n; ++j) { + if (i == j) + continue; + + auto node_i = seq_body->children[i].get(); + auto node_j = seq_body->children[j].get(); + int rem_stage_j = stage_map[node_j]; + + if (!HasDependency(node_i, node_j)) + continue; + + // LOG(INFO) << "[ScheduleRecursive] Conflict var detection between + // " << i << " and " << j; + + if (stage_map[node_j] == stage_map[node_i]) + continue; + + auto node_i_task = static_cast(node_i); + auto node_i_let_stmt = node_i_task->stmts[0].as(); + + auto iter = ctrl->control->loop_var; + auto step = ctrl->control->step.has_value() + ? ctrl->control->step.value() + : 1; + auto cloned_value = node_i_let_stmt->value; + auto cloned_let_stmt = + LetStmt(node_i_let_stmt->var.copy_with_suffix(""), cloned_value, + Evaluate(0)); + auto cloned_task = std::make_shared(); + cloned_task->stmts.push_back(cloned_let_stmt); + stage_map[cloned_task.get()] = rem_stage_j; + + for (int k = j; k < n; ++k) { + auto node_k = seq_body->children[k].get(); + auto task_k = static_cast(node_k); + if (rem_stage_j != stage_map[node_k]) + continue; + if (HasDependency(node_i, node_k)) { + for (size_t id = 0; id < task_k->stmts.size(); ++id) { + task_k->stmts[id] = Substitute( + task_k->stmts[id], + {{node_i_let_stmt->var, cloned_let_stmt->var}}); + } + task_k->SubstituteVar(node_i_let_stmt->var, + cloned_let_stmt->var); + stage_map[node_k] = rem_stage_j; + } + } + + seq_body->children.insert(seq_body->children.begin() + j, + std::move(cloned_task)); + n += 1; + return true; // Conflict resolved, restart the loop + } + } + return false; + }; + // Resolve conflicts until no more conflicts exist or max iterations reached + // (to avoid infinite loop) + int conflict_count = 0; + while (SolveConflictVar() && ++conflict_count < 100) + ; + for (auto &node : seq_body->children) { - auto unit = std::make_unique(); + auto unit = std::make_shared(); unit->stage = stage_map[node.get()]; unit->child = std::move(node); node = std::move(unit); @@ -827,6 +954,9 @@ class ScheduleUnitBuilder { return a->buffer.same_as(b->buffer); } + // Check if two variables are the same + bool SameVar(const Var &a, const Var &b) const { return a.same_as(b); } + // Check if two IRStructures have data dependency (excluding read-after-read) bool HasDependency(const IRStructure *a, const IRStructure *b) const { // Check if either node contains loop_break (if it's a TaskNode) @@ -873,6 +1003,12 @@ class ScheduleUnitBuilder { return true; } } + for (const auto &write_var_a : a->GetWriteVars()) { + for (const auto &read_var_b : b->GetReadVars()) { + if (SameVar(write_var_a, read_var_b)) + return true; + } + } return false; } @@ -973,21 +1109,13 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { if (!node) return; - if (node->IsTask()) { - // TaskNode: no further scheduling needed - return; - } else if (node->IsSequence()) { - auto seq = static_cast(node); - - // First, recursively schedule all children - for (size_t i = 0; i < seq->children.size(); ++i) { - ScheduleRecursive(seq->children[i].get()); - } - + auto ChildrenScheduleHelper = + [&](std::vector> origin_children) + -> std::vector> { // Now collect child nodes for potential scheduling std::vector child_nodes; - child_nodes.reserve(seq->children.size()); - for (const auto &child : seq->children) { + child_nodes.reserve(origin_children.size()); + for (const auto &child : origin_children) { child_nodes.push_back(child.get()); } @@ -1011,7 +1139,7 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { } // Reorder children according to sorted_nodes - std::vector> reordered_children; + std::vector> reordered_children; reordered_children.reserve(sorted_nodes.size()); for (IRStructure *sorted_node : sorted_nodes) { @@ -1021,41 +1149,73 @@ void ScheduleUnitBuilder::ScheduleRecursive(IRStructure *node) { "children mapping"; } size_t old_idx = it->second; - reordered_children.push_back(std::move(seq->children[old_idx])); + reordered_children.emplace_back(origin_children[old_idx]); } // Move reordered children back - seq->children = std::move(reordered_children); + origin_children = reordered_children; } - for (auto &node : seq->children) { - auto unit = std::make_unique(); + for (auto &node : origin_children) { + auto unit = std::make_shared(); unit->stage = -1; - unit->child = std::move(node); - node = std::move(unit); + unit->child = std::shared_ptr(node); + node = unit; + } + return origin_children; + }; + + if (node->IsTask()) { + // TaskNode: no further scheduling needed + return; + } else if (node->IsSequence()) { + auto seq = static_cast(node); + + // First, recursively schedule all children + std::vector> seq_children, origin_children; + GatherTaskNodes(seq->children, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); } + + seq->children = ChildrenScheduleHelper(origin_children); return; } else if (node->IsControl()) { auto ctrl = static_cast(node); // Now schedule the ControlNode's internal tasks (if any) as a unit // The body should now be a SequenceNode containing the tasks - if (ctrl->child && ctrl->child->IsSequence()) { - auto seq_body = static_cast(ctrl->child.get()); - for (const auto &child : seq_body->children) { - ScheduleRecursive(child.get()); + if (ctrl->child) { + if (ctrl->child->IsSequence()) { + auto seq_body = static_cast(ctrl->child.get()); + std::vector> origin_children; + GatherTaskNodes(seq_body->children, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); + } + Z3SchedulePythonLoop(ctrl); + } else if (ctrl->child->IsWrapper()) { + auto wrapper = static_cast(ctrl->child.get()); + std::vector> origin_children; + GatherTaskNodes({wrapper->child}, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); + } + Z3SchedulePythonLoop(ctrl); + } else { + ScheduleRecursive(ctrl->child.get()); } - Z3SchedulePythonLoop(ctrl); - } else { - ScheduleRecursive(ctrl->child.get()); } return; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - if (wrapper->child) { - ScheduleRecursive(wrapper->child.get()); - wrapper->SetII(wrapper->child->GetII()); - wrapper->SetLatency(wrapper->child->GetLatency()); - } + std::vector> origin_children; + GatherTaskNodes({wrapper->child}, origin_children); + for (auto &child : origin_children) { + ScheduleRecursive(child.get()); + } + auto seq_node = std::make_shared(); + seq_node->children = ChildrenScheduleHelper(origin_children); + *node = *(seq_node.get()); return; } @@ -1139,7 +1299,7 @@ class TilelangRootBodyReplacer : public StmtMutator { // Visitor to build IRStructure from TIR statements class IRStructureBuilder : public StmtVisitor { public: - std::unique_ptr Build(const Stmt &stmt, int64_t thread_count = 1, + std::shared_ptr Build(const Stmt &stmt, int64_t thread_count = 1, Target target = Target()) { thread_count_ = thread_count; target_ = target; @@ -1149,7 +1309,7 @@ class IRStructureBuilder : public StmtVisitor { << "IRStructureBuilder: root_ is null after visiting statement. " << "This may indicate an unhandled statement type."; // Return an empty TaskNode as fallback - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(stmt); return task_node; } @@ -1158,7 +1318,7 @@ class IRStructureBuilder : public StmtVisitor { protected: void VisitStmt_(const SeqStmtNode *op) override { - auto seq_node = std::make_unique(); + auto seq_node = std::make_shared(); for (size_t i = 0; i < op->seq.size(); i++) { VisitStmt(op->seq[i]); @@ -1173,7 +1333,7 @@ class IRStructureBuilder : public StmtVisitor { // Determine if this is a sequential or parallel for if (op->kind == ForKind::kSerial) { // Sequential For -> ControlNode - auto control_node = std::make_unique(); + auto control_node = std::make_shared(); control_node->control = GetRef(op); // Process the loop body @@ -1183,17 +1343,10 @@ class IRStructureBuilder : public StmtVisitor { } else { } - // Absorb WrapperNode chain into ControlNode.wrappers - while (control_node->child && control_node->child->IsWrapper()) { - auto *wrapper = static_cast(control_node->child.get()); - control_node->wrappers.push_back(wrapper->wrapper); - control_node->child = std::move(wrapper->child); - } - root_ = std::move(control_node); } else { // Parallel For -> TaskNode - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); // Analyze the loop body for resource usage @@ -1205,7 +1358,7 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const EvaluateNode *op) override { // Evaluate statement (usually a Call) -> TaskNode - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); // Analyze the expression for resource usage @@ -1216,10 +1369,11 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const IfThenElseNode *op) override { // If statement -> treat as TaskNode for now (could be refined later) - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); AnalyzeMemoryExpr(op->condition, task_node.get()); + AnalyzeResourceUsage(Evaluate(op->condition), task_node.get(), true); // Analyze both branches for resource usage AnalyzeResourceUsage(op->then_case, task_node.get()); @@ -1232,8 +1386,12 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const LetStmtNode *op) override { // Wrapper statement -> WrapperNode - auto wrapper_node = std::make_unique(); + auto wrapper_node = std::make_shared(); wrapper_node->wrapper = GetRef(op); + auto task_node = std::make_shared(); + task_node->stmts.push_back(GetLetDecl(op)); + AnalyzeResourceUsage(GetLetDecl(op), task_node.get()); + wrapper_node->task = std::move(task_node); // Process the wrapperbody VisitStmt(op->body); @@ -1246,8 +1404,12 @@ class IRStructureBuilder : public StmtVisitor { void VisitStmt_(const AttrStmtNode *op) override { // Wrapper statement -> WrapperNode - auto wrapper_node = std::make_unique(); + auto wrapper_node = std::make_shared(); wrapper_node->wrapper = GetRef(op); + auto task_node = std::make_shared(); + task_node->stmts.push_back(GetAttrDecl(op)); + AnalyzeResourceUsage(GetAttrDecl(op), task_node.get()); + wrapper_node->task = std::move(task_node); // Process the wrapperbody VisitStmt(op->body); @@ -1259,7 +1421,7 @@ class IRStructureBuilder : public StmtVisitor { } void VisitStmt_(const WhileNode *op) override { - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); // Analyze condition and body for resource usage @@ -1275,18 +1437,19 @@ class IRStructureBuilder : public StmtVisitor { // TilelangRootBodyExtractor If we encounter it here, it means we're // processing the entire function body (not extracted), which should only // happen when there's no tilelang_root block - auto task_node = std::make_unique(); + auto task_node = std::make_shared(); task_node->stmts.push_back(GetRef(op)); AnalyzeResourceUsage(GetRef(op), task_node.get()); root_ = std::move(task_node); } private: - std::unique_ptr root_; + std::shared_ptr root_; int64_t thread_count_ = 1; Target target_; - void AnalyzeResourceUsage(const Stmt &stmt, TaskNode *task_node) { + void AnalyzeResourceUsage(const Stmt &stmt, TaskNode *task_node, + bool only_variables = false) { // Recursively analyze statements to determine resource usage struct ResourceAnalyzer : public StmtExprVisitor { TaskNode *task_node; @@ -1307,7 +1470,7 @@ class IRStructureBuilder : public StmtVisitor { std::vector tensor_core_shapes; // GemmInst: the resolved tensor core instruction (single, asserted - // unique) + // shared) GemmInst gemm_inst{GemmInst::kMMA}; bool has_gemm_inst{false}; @@ -1336,7 +1499,8 @@ class IRStructureBuilder : public StmtVisitor { // Check if this is a TMA copy operation if (op->op.same_as(copy_op)) { - bool found_global = false; + bool found_global = false, found_shared = false; + int idx_global = -1, idx_shared = -1; for (unsigned idx = 0; idx != 2; ++idx) { auto region = Downcast(op->args[idx]); if (const auto *buffer_load = @@ -1346,13 +1510,21 @@ class IRStructureBuilder : public StmtVisitor { MemoryType mem_type = GetMemoryTypeFromScope(scope); if (mem_type == MemoryType::kGlobal) { found_global = true; - if (idx == 0) { - found_tma_load = true; - } + idx_global = idx; + } + if (mem_type == MemoryType::kShared) { + found_shared = true; + idx_shared = idx; } } } - found_tma |= found_global; + found_tma = false; + if (found_global && found_shared) { + if (idx_global == 0 && idx_shared == 1) + found_tma = true; + if (idx_global == 1 && idx_shared == 0) + found_tma = true; + } } else if (op->op.same_as(gemm_py_op) || op->op.same_as(gemm_op)) { found_tensor = true; @@ -1428,37 +1600,41 @@ class IRStructureBuilder : public StmtVisitor { ResourceAnalyzer analyzer(task_node, target_, thread_count_); analyzer(stmt); - // Set task node flags based on what was found - if (analyzer.found_tma) { - task_node->SetUsesTMACore(true); - if (analyzer.found_tma_load) { - task_node->SetHasTMALoad(true); + if (!only_variables) { + // Set task node flags based on what was found + if (analyzer.found_tma) { + task_node->SetUsesTMACore(true); + if (analyzer.found_tma_load) { + task_node->SetHasTMALoad(true); + } } - } - if (analyzer.found_tensor) { - task_node->SetUsesTensorCore(true); - // Set Tensor Core shape information if available - for (const auto &shape : analyzer.tensor_core_shapes) { - if (shape.m > 0 && shape.n > 0 && shape.k > 0) { - task_node->AddTensorCoreShape(shape.m, shape.n, shape.k); + if (analyzer.found_tensor) { + task_node->SetUsesTensorCore(true); + // Set Tensor Core shape information if available + for (const auto &shape : analyzer.tensor_core_shapes) { + if (shape.m > 0 && shape.n > 0 && shape.k > 0) { + task_node->AddTensorCoreShape(shape.m, shape.n, shape.k); + } + } + // Set GemmInst information + if (analyzer.has_gemm_inst) { + task_node->SetGemmInst(analyzer.gemm_inst); } } - // Set GemmInst information - if (analyzer.has_gemm_inst) { - task_node->SetGemmInst(analyzer.gemm_inst); + // If neither TMA nor Tensor core was used, and CUDA operations were + // found, set CUDA core flag + if (!analyzer.found_tma && !analyzer.found_tensor) { + task_node->SetUsesCUDACore(true); } } - // If neither TMA nor Tensor core was used, and CUDA operations were found, - // set CUDA core flag - if (!analyzer.found_tma && !analyzer.found_tensor) { - task_node->SetUsesCUDACore(true); - } // Analyze memory access regions MemoryAccessDetector memory_detector; memory_detector.Analyze(stmt); std::vector read_regions = memory_detector.GetReadRegions(); std::vector write_regions = memory_detector.GetWriteRegions(); + std::vector read_vars = memory_detector.GetReadVars(); + std::vector write_vars = memory_detector.GetWriteVars(); // Merge with existing regions (avoid duplicates) for (const auto ®ion : read_regions) { @@ -1469,6 +1645,14 @@ class IRStructureBuilder : public StmtVisitor { task_node->AddWriteRegion(region); } + for (const auto &var : read_vars) { + task_node->AddReadVar(var); + } + + for (const auto &var : write_vars) { + task_node->AddWriteVar(var); + } + // Estimate latency and initiation interval for this task LatencyEstimator latency_estimator; latency_estimator.SetThreadCount(thread_count_); @@ -1498,6 +1682,8 @@ class IRStructureBuilder : public StmtVisitor { } }; +Stmt ReNestLetStmts(const Stmt &stmt); + // The main pass function tvm::transform::Pass AutoSchedule(const bool enable_epi) { using namespace tir::transform; @@ -1588,6 +1774,8 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { TilelangRootBodyReplacer replacer(new_body); final_body = replacer(func->body); + final_body = ReNestLetStmts(final_body); + // Create a new PrimFunc with the updated body auto new_func = PrimFunc(func->params, final_body, func->ret_type, func->buffer_map, func->attrs); @@ -1699,6 +1887,10 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { final_body = RewriteAllocBuffers(final_body, buffer_infos); } + // LOG(INFO) << final_body << std::endl; + + final_body = ReNestLetStmts(final_body); + // Create a new PrimFunc with the updated body auto new_func = PrimFunc(func->params, final_body, func->ret_type, func->buffer_map, func->attrs); @@ -1708,55 +1900,315 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { return CreatePrimFuncPass(pass_func, 0, "tl.AutoSchedule", {}); } -// Helper function to clone IRStructure with warpgroup filter -std::unique_ptr -CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { - if (!node || !node->containWarpgroupId(warpgroup_id)) +// Helper: check if a TaskNode is a LetDecl (single LetStmt with empty body) +static bool IsLetDeclTask(const TaskNode *task) { + return task->stmts.size() == 1 && task->stmts[0].as() != nullptr; +} + +// Helper: check if an IRStructure node is a LetDecl task (or a ScheduleUnit +// wrapping one) +static bool IsLetDeclNode(const IRStructure *node) { + if (!node) + return false; + if (node->IsTask()) { + return IsLetDeclTask(static_cast(node)); + } + if (node->IsScheduleUnit()) { + auto unit = static_cast(node); + return unit->child && unit->child->IsTask() && + IsLetDeclTask(static_cast(unit->child.get())); + } + return false; +} + +// Helper: check if an IRStructure subtree contains any LetDecl tasks +static bool ContainsLetDecl(const IRStructure *node) { + if (!node) + return false; + if (IsLetDeclNode(node)) + return true; + if (node->IsSequence()) { + auto seq = static_cast(node); + for (const auto &child : seq->children) { + if (ContainsLetDecl(child.get())) + return true; + } + } else if (node->IsControl()) { + auto ctrl = static_cast(node); + return ContainsLetDecl(ctrl->child.get()); + } else if (node->IsWrapper()) { + auto wrapper = static_cast(node); + return ContainsLetDecl(wrapper->child.get()); + } else if (node->IsScheduleUnit()) { + auto unit = static_cast(node); + return ContainsLetDecl(unit->child.get()); + } + return false; +} + +// Helper function to clone IRStructure with warpgroup filter. +std::shared_ptr +CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id, + Map &var_remap) { + if (!node) return nullptr; if (node->IsTask()) { auto task = static_cast(node); - return task->Clone(); + + // LetDecl tasks are always included in every warp group clone. + // Create a fresh variable copy so the two warp groups use different names. + if (IsLetDeclTask(task)) { + const auto *let = task->stmts[0].as(); + auto new_var = let->var.copy_with_suffix(""); + // Substitute previously renamed variables in the value expression. + PrimExpr new_value = + var_remap.empty() ? let->value : Substitute(let->value, var_remap); + var_remap.Set(let->var, new_var); + auto new_task = std::make_shared(); + new_task->stmts.push_back(LetStmt(new_var, new_value, Evaluate(0))); + return new_task; + } + + // Non-LetDecl tasks: only include if warp group matches + if (!node->containWarpgroupId(warpgroup_id)) + return nullptr; + auto cloned = task->Clone(); + // Substitute renamed LetDecl variables in task statements + if (!var_remap.empty()) { + auto ct = static_cast(cloned.get()); + for (size_t i = 0; i < ct->stmts.size(); ++i) { + ct->stmts[i] = Substitute(ct->stmts[i], var_remap); + } + } + return cloned; } else if (node->IsSequence()) { + // A SequenceNode is included if it contains the target warp group + // OR if it contains LetDecl tasks (which are always needed). + if (!node->containWarpgroupId(warpgroup_id) && !ContainsLetDecl(node)) + return nullptr; auto seq = static_cast(node); - auto new_seq = std::make_unique(); + auto new_seq = std::make_shared(); for (const auto &child : seq->children) { - auto new_child = - CloneIRStructureWithWarpgroupFilter(child.get(), warpgroup_id); + auto new_child = CloneIRStructureWithWarpgroupFilter( + child.get(), warpgroup_id, var_remap); if (new_child) { new_seq->children.push_back(std::move(new_child)); } } + if (new_seq->children.empty()) + return nullptr; return new_seq; } else if (node->IsControl()) { + // A ControlNode is included if it contains the target warp group + // OR if it contains LetDecl tasks. + if (!node->containWarpgroupId(warpgroup_id) && !ContainsLetDecl(node)) + return nullptr; auto ctrl = static_cast(node); - auto new_ctrl = std::make_unique(); + auto new_ctrl = std::make_shared(); new_ctrl->control = ctrl->control; new_ctrl->SetPromote(ctrl->hasPromote()); - new_ctrl->wrappers = ctrl->wrappers; - new_ctrl->child = - CloneIRStructureWithWarpgroupFilter(ctrl->child.get(), warpgroup_id); + new_ctrl->child = CloneIRStructureWithWarpgroupFilter( + ctrl->child.get(), warpgroup_id, var_remap); return new_ctrl; } else if (node->IsWrapper()) { + if (!node->containWarpgroupId(warpgroup_id) && !ContainsLetDecl(node)) + return nullptr; auto wrapper = static_cast(node); - auto new_wrapper = std::make_unique(); + auto new_wrapper = std::make_shared(); + // Keep the wrapper statement as-is (do NOT rename LetStmt wrappers here; + // only LetDecl TaskNodes get renamed). new_wrapper->wrapper = wrapper->wrapper; - new_wrapper->child = - CloneIRStructureWithWarpgroupFilter(wrapper->child.get(), warpgroup_id); + new_wrapper->child = CloneIRStructureWithWarpgroupFilter( + wrapper->child.get(), warpgroup_id, var_remap); return new_wrapper; } else if (node->IsScheduleUnit()) { auto unit = static_cast(node); - auto new_unit = std::make_unique(); - new_unit->before[warpgroup_id] = unit->before[warpgroup_id]; - new_unit->after[warpgroup_id] = unit->after[warpgroup_id]; + bool child_is_let_decl = IsLetDeclNode(unit->child.get()); + + // Include the ScheduleUnit if the child is a LetDecl or the warp group + // matches. + if (!child_is_let_decl && !node->containWarpgroupId(warpgroup_id)) + return nullptr; + + auto new_unit = std::make_shared(); new_unit->stage = unit->stage; - new_unit->child = - CloneIRStructureWithWarpgroupFilter(unit->child.get(), warpgroup_id); + new_unit->child = CloneIRStructureWithWarpgroupFilter( + unit->child.get(), warpgroup_id, var_remap); + + if (!child_is_let_decl) { + // Copy before/after for the target warp group + new_unit->before[warpgroup_id] = unit->before[warpgroup_id]; + new_unit->after[warpgroup_id] = unit->after[warpgroup_id]; + // Substitute renamed LetDecl variables in before/after stmts + if (!var_remap.empty()) { + for (auto &s : new_unit->before[warpgroup_id]) { + s = Substitute(s, var_remap); + } + for (auto &s : new_unit->after[warpgroup_id]) { + s = Substitute(s, var_remap); + } + } + } return new_unit; } LOG(FATAL); } +// Entry point overload — creates a fresh var_remap per call +std::shared_ptr +CloneIRStructureWithWarpgroupFilter(IRStructure *node, int warpgroup_id) { + Map var_remap; + return CloneIRStructureWithWarpgroupFilter(node, warpgroup_id, var_remap); +} + +// Simple visitor to collect all Var references from statements/expressions +class VarRefCollector : public StmtExprVisitor { +public: + std::unordered_set vars; + void VisitExpr_(const VarNode *op) override { vars.insert(op); } +}; + +// Remove LetDecl definitions whose variables are not referenced by any +// non-LetDecl task in the IR tree. After warp-group partitioning some +// LetDecl values may access out-of-bounds indices because the consumer +// tasks that used them ended up in the other warp group. +std::shared_ptr +RemoveUnusedLetDecls(std::shared_ptr root) { + if (!root) + return nullptr; + + // Phase 1: Collect LetDecl definitions and variable references from + // non-LetDecl nodes (task stmts and ScheduleUnit before/after). + struct LetDeclEntry { + const VarNode *var; + PrimExpr value; + }; + std::vector let_decls; + std::unordered_set referenced_vars; + + std::function collect = + [&](const IRStructure *node) { + if (!node) + return; + if (node->IsTask()) { + auto task = static_cast(node); + if (IsLetDeclTask(task)) { + const auto *let = task->stmts[0].as(); + let_decls.push_back({let->var.get(), let->value}); + } else { + VarRefCollector collector; + for (const auto &stmt : task->stmts) { + collector(stmt); + } + referenced_vars.insert(collector.vars.begin(), + collector.vars.end()); + } + } else if (node->IsSequence()) { + for (const auto &child : + static_cast(node)->children) { + collect(child.get()); + } + } else if (node->IsControl()) { + collect(static_cast(node)->child.get()); + } else if (node->IsWrapper()) { + collect(static_cast(node)->child.get()); + } else if (node->IsScheduleUnit()) { + auto unit = static_cast(node); + collect(unit->child.get()); + VarRefCollector collector; + for (const auto &stmts : unit->before) { + for (const auto &s : stmts) + collector(s); + } + for (const auto &stmts : unit->after) { + for (const auto &s : stmts) + collector(s); + } + referenced_vars.insert(collector.vars.begin(), collector.vars.end()); + } + }; + collect(root.get()); + + // Phase 2: Transitive closure — if a LetDecl var is referenced, + // all vars in its value expression are transitively referenced too. + bool changed = true; + while (changed) { + changed = false; + for (const auto &entry : let_decls) { + if (referenced_vars.count(entry.var)) { + VarRefCollector collector; + collector(entry.value); + for (const auto *v : collector.vars) { + if (!referenced_vars.count(v)) { + referenced_vars.insert(v); + changed = true; + } + } + } + } + } + + // Phase 3: Filter the tree — remove LetDecl tasks for unused vars. + std::function( + const std::shared_ptr &)> + filter_tree = [&](const std::shared_ptr &node) + -> std::shared_ptr { + if (!node) + return nullptr; + if (node->IsTask()) { + if (IsLetDeclTask(static_cast(node.get()))) { + const auto *let = static_cast(node.get()) + ->stmts[0] + .as(); + if (!referenced_vars.count(let->var.get())) { + return nullptr; // Remove unused LetDecl + } + } + return node; + } else if (node->IsSequence()) { + auto seq = static_cast(node.get()); + auto new_seq = std::make_shared(); + for (const auto &child : seq->children) { + auto filtered = filter_tree(child); + if (filtered) + new_seq->children.push_back(std::move(filtered)); + } + if (new_seq->children.empty()) + return nullptr; + return new_seq; + } else if (node->IsControl()) { + auto ctrl = static_cast(node.get()); + auto new_ctrl = std::make_shared(); + new_ctrl->control = ctrl->control; + new_ctrl->SetPromote(ctrl->hasPromote()); + new_ctrl->child = filter_tree(ctrl->child); + if (!new_ctrl->child) + return nullptr; + return new_ctrl; + } else if (node->IsWrapper()) { + auto wrapper = static_cast(node.get()); + auto new_wrapper = std::make_shared(); + new_wrapper->wrapper = wrapper->wrapper; + new_wrapper->child = filter_tree(wrapper->child); + return new_wrapper; + } else if (node->IsScheduleUnit()) { + auto unit = static_cast(node.get()); + auto new_unit = std::make_shared(); + new_unit->stage = unit->stage; + new_unit->before = unit->before; + new_unit->after = unit->after; + new_unit->child = filter_tree(unit->child); + if (!new_unit->child) + return nullptr; + return new_unit; + } + return node; + }; + + return filter_tree(root); +} + class SimtCopyDetector : public StmtExprVisitor { public: static bool Detect(const Stmt &stmt) { @@ -1869,16 +2321,6 @@ Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi) { LOG(FATAL); } Stmt body = SeqStmt::Flatten(stmts); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - body = LetStmt(let->var, let->value, body); - } else if (const auto *attr = (*it).as()) { - body = AttrStmt(attr->node, attr->attr_key, attr->value, body); - } - } // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; filtered_annotations.erase("num_stages"); @@ -1969,32 +2411,31 @@ Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi) { stmts.push_back(stmt); } } - Map substitution; - PrimExpr condition = - And(loop_var < loop_extent, loop_var >= loop_start); - if (unit->stage == min_stages) { - condition = loop_var >= loop_start; - } - if (unit->stage == max_stages) { - condition = loop_var < loop_extent; - } - Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + Map substitution, substitution_cond; substitution.Set(loop_var, loop_var - loop_step * (max_stages - unit->stage)); - steady.push_back(Substitute(stmt, substitution)); - } - Stmt new_body = SeqStmt::Flatten(steady); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - new_body = LetStmt(let->var, let->value, new_body); - } else if (const auto *attr = (*it).as()) { - new_body = - AttrStmt(attr->node, attr->attr_key, attr->value, new_body); + substitution_cond.Set( + loop_var, + Max(loop_start, + Min(loop_start + loop_extent - loop_step, + loop_var - loop_step * (max_stages - unit->stage)))); + if (IsLetDeclNode(unit->child.get())) { + Stmt stmt = SeqStmt::Flatten(stmts); + steady.push_back(Substitute(stmt, substitution_cond)); + } else { + PrimExpr condition = + And(loop_var < loop_extent, loop_var >= loop_start); + if (unit->stage == min_stages) { + condition = loop_var >= loop_start; + } + if (unit->stage == max_stages) { + condition = loop_var < loop_extent; + } + Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + steady.push_back(Substitute(stmt, substitution)); } } + Stmt new_body = SeqStmt::Flatten(steady); auto new_var = loop_var.copy_with_suffix(""); // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; @@ -2198,16 +2639,6 @@ Stmt ApplyWarpgroupPartitionToIRStructure( LOG(FATAL); } Stmt body = SeqStmt::Flatten(stmts); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - body = LetStmt(let->var, let->value, body); - } else if (const auto *attr = (*it).as()) { - body = AttrStmt(attr->node, attr->attr_key, attr->value, body); - } - } // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; filtered_annotations.erase("num_stages"); @@ -2298,32 +2729,31 @@ Stmt ApplyWarpgroupPartitionToIRStructure( stmts.push_back(stmt); } } - Map substitution; - PrimExpr condition = - And(loop_var < loop_extent, loop_var >= loop_start); - if (unit->stage == min_stages) { - condition = loop_var >= loop_start; - } - if (unit->stage == max_stages) { - condition = loop_var < loop_extent; - } - Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + Map substitution, substitution_cond; substitution.Set(loop_var, loop_var - loop_step * (max_stages - unit->stage)); - steady.push_back(Substitute(stmt, substitution)); - } - Stmt new_body = SeqStmt::Flatten(steady); - // Re-wrap body with ControlNode wrappers (innermost first = reverse - // order) - for (auto it = ctrl->wrappers.rbegin(); it != ctrl->wrappers.rend(); - ++it) { - if (const auto *let = (*it).as()) { - new_body = LetStmt(let->var, let->value, new_body); - } else if (const auto *attr = (*it).as()) { - new_body = - AttrStmt(attr->node, attr->attr_key, attr->value, new_body); + substitution_cond.Set( + loop_var, + Max(loop_start, + Min(loop_start + loop_extent - loop_step, + loop_var - loop_step * (max_stages - unit->stage)))); + if (IsLetDeclNode(unit->child.get())) { + Stmt stmt = SeqStmt::Flatten(stmts); + steady.push_back(Substitute(stmt, substitution_cond)); + } else { + PrimExpr condition = + And(loop_var < loop_extent, loop_var >= loop_start); + if (unit->stage == min_stages) { + condition = loop_var >= loop_start; + } + if (unit->stage == max_stages) { + condition = loop_var < loop_extent; + } + Stmt stmt = IfThenElse(condition, SeqStmt::Flatten(stmts)); + steady.push_back(Substitute(stmt, substitution)); } } + Stmt new_body = SeqStmt::Flatten(steady); auto new_var = loop_var.copy_with_suffix(""); // Filter out "num_stages" annotation Map filtered_annotations = ctrl->control->annotations; @@ -2396,11 +2826,11 @@ Stmt ApplyWarpgroupPartitionToIRStructure( // Helper function to clone IRStructure filtering tasks with warpgroup_id == // -1 (neutral tasks) - std::function(IRStructure *)> + std::function(IRStructure *)> clone_neutral_filter; clone_neutral_filter = [&clone_neutral_filter]( - IRStructure *node) -> std::unique_ptr { + IRStructure *node) -> std::shared_ptr { if (!node) return nullptr; @@ -2409,19 +2839,19 @@ Stmt ApplyWarpgroupPartitionToIRStructure( if (task->GetWarpgroupId() == -1) { return task->Clone(); } else { - auto new_task = std::make_unique(); + auto new_task = std::make_shared(); // Empty statements return new_task; } } else if (node->IsSequence()) { auto seq = static_cast(node); - auto new_seq = std::make_unique(); + auto new_seq = std::make_shared(); for (const auto &child : seq->children) { if (child) { auto node = static_cast(child.get()); auto new_node = clone_neutral_filter(node->child.get()); if (new_node) { - auto new_unit = std::make_unique(); + auto new_unit = std::make_shared(); new_unit->child = std::move(new_node); new_seq->children.push_back(std::move(new_unit)); } @@ -2430,7 +2860,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return new_seq; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - auto new_wrapper = std::make_unique(); + auto new_wrapper = std::make_shared(); new_wrapper->child = clone_neutral_filter(wrapper->child.get()); if (new_wrapper->child) { return new_wrapper; @@ -2453,13 +2883,13 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return false; }; - std::function( + std::function( IRStructure *, const std::function &, int)> clone_neutral_filter_with_top_level; clone_neutral_filter_with_top_level = [&clone_neutral_filter_with_top_level, &clone_neutral_filter]( IRStructure *node, const std::function &include_top_level, - int top_level_index) -> std::unique_ptr { + int top_level_index) -> std::shared_ptr { if (!node) return nullptr; @@ -2467,13 +2897,13 @@ Stmt ApplyWarpgroupPartitionToIRStructure( if (include_top_level(top_level_index)) { return clone_neutral_filter(node); } else { - auto new_task = std::make_unique(); + auto new_task = std::make_shared(); // Empty statements return new_task; } } else if (node->IsSequence()) { auto seq = static_cast(node); - auto new_seq = std::make_unique(); + auto new_seq = std::make_shared(); int child_index = 0; for (const auto &child : seq->children) { if (child) { @@ -2484,7 +2914,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( schedule_unit->child.get(), include_top_level, next_top_level_index); if (new_node) { - auto new_unit = std::make_unique(); + auto new_unit = std::make_shared(); new_unit->child = std::move(new_node); new_seq->children.push_back(std::move(new_unit)); } @@ -2494,7 +2924,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return new_seq; } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - auto new_wrapper = std::make_unique(); + auto new_wrapper = std::make_shared(); new_wrapper->child = clone_neutral_filter_with_top_level( wrapper->child.get(), include_top_level, top_level_index); if (new_wrapper->child) { @@ -2544,8 +2974,10 @@ Stmt ApplyWarpgroupPartitionToIRStructure( root, is_epi_top_level_index, -1) : nullptr; - auto wg0_structure = CloneIRStructureWithWarpgroupFilter(root, 0); - auto wg1_structure = CloneIRStructureWithWarpgroupFilter(root, 1); + auto wg0_structure = + RemoveUnusedLetDecls(CloneIRStructureWithWarpgroupFilter(root, 0)); + auto wg1_structure = + RemoveUnusedLetDecls(CloneIRStructureWithWarpgroupFilter(root, 1)); bool wg_pro_neutral_has_stmts = wg_pro_neutral_structure @@ -2675,6 +3107,113 @@ Stmt ApplyWarpgroupPartitionToIRStructure( return combined_stmt; } +// Re-write LetStmt to nest them properly +// Example transformation: +// SeqStmt { +// let x = 42 { Evaluate(0) } // standalone, empty body +// let y = x+1 { Evaluate(0) } // standalone, empty body +// compute(x, y) // actual work +// store(result) +// } +// becomes: +// let x = 42 { +// let y = x+1 { +// SeqStmt { +// compute(x, y) +// store(result) +// } +// } +// } +class LetStmtNester : public StmtMutator { +public: + Stmt VisitStmt_(const SeqStmtNode *op) override { + Array stmts; + for (const auto &stmt : op->seq) { + stmts.push_back(this->VisitStmt(stmt)); + } + + Array flat_stmts; + for (const auto &stmt : stmts) { + if (const auto *inner_seq = stmt.as()) { + for (const auto &inner_stmt : inner_seq->seq) { + flat_stmts.push_back(inner_stmt); + } + } else { + flat_stmts.push_back(stmt); + } + } + stmts = flat_stmts; + + for (int i = static_cast(stmts.size()) - 2; i >= 0; --i) { + if (const auto *let = stmts[i].as()) { + if (IsEmptyBody(let->body)) { + Stmt absorbed_body = CollectRemaining(stmts, i + 1); + stmts = TruncateAndReplace( + stmts, i, LetStmt(let->var, let->value, absorbed_body)); + } + } else if (const auto *attr = stmts[i].as()) { + if (IsEmptyBody(attr->body)) { + Stmt absorbed_body = CollectRemaining(stmts, i + 1); + stmts = TruncateAndReplace( + stmts, i, + AttrStmt(attr->node, attr->attr_key, attr->value, absorbed_body)); + } + } + } + + if (stmts.empty()) + return Evaluate(0); + if (stmts.size() == 1) + return stmts[0]; + + return SeqStmt(stmts); + } + +private: + // Check if a statement body is Evaluate(0) — the empty placeholder + static bool IsEmptyBody(const Stmt &stmt) { + if (const auto *eval = stmt.as()) { + if (const auto *imm = eval->value.as()) { + return imm->value == 0; + } + } + return false; + } + + // Collect stmts[start .. end) into a single Stmt + static Stmt CollectRemaining(const Array &stmts, int start) { + int n = static_cast(stmts.size()); + if (start >= n) { + return Evaluate(0); + } + if (start == n - 1) { + return stmts[start]; + } + Array remaining; + for (int j = start; j < n; ++j) { + remaining.push_back(stmts[j]); + } + return SeqStmt(remaining); + } + + // Keep stmts[0..index), replace stmts[index] with new_stmt, + // discard everything after (already absorbed into new_stmt body) + static Array TruncateAndReplace(const Array &stmts, int index, + Stmt new_stmt) { + Array result; + for (int j = 0; j < index; ++j) { + result.push_back(stmts[j]); + } + result.push_back(new_stmt); + return result; + } +}; + +Stmt ReNestLetStmts(const Stmt &stmt) { + LetStmtNester nester; + return nester(stmt); +} + // StmtMutator to rewrite alloc_buffers in Block nodes class AllocBufferRewriter : public StmtMutator { public: diff --git a/src/transform/auto_schedule/ir_structure.cc b/src/transform/auto_schedule/ir_structure.cc index 4ba3ce9269..394e776307 100644 --- a/src/transform/auto_schedule/ir_structure.cc +++ b/src/transform/auto_schedule/ir_structure.cc @@ -117,6 +117,18 @@ std::vector SequenceNode::GetReadRegions() const { return deduplicated; } +std::vector SequenceNode::GetReadVars() const { + std::vector all_vars; + for (const auto &child : children) { + if (child) { + auto child_read_vars = child->GetReadVars(); + all_vars.insert(all_vars.end(), child_read_vars.begin(), + child_read_vars.end()); + } + } + return all_vars; +} + std::vector SequenceNode::GetWriteRegions() const { std::vector all_write_regions; for (const auto &child : children) { @@ -144,6 +156,18 @@ std::vector SequenceNode::GetWriteRegions() const { return deduplicated; } +std::vector SequenceNode::GetWriteVars() const { + std::vector all_vars; + for (const auto &child : children) { + if (child) { + auto child_write_vars = child->GetWriteVars(); + all_vars.insert(all_vars.end(), child_write_vars.begin(), + child_write_vars.end()); + } + } + return all_vars; +} + int64_t SequenceNode::GetLatency() const { return latency_; } int64_t SequenceNode::GetII() const { return ii_; } @@ -186,20 +210,8 @@ void SequenceNode::SetLatency(int64_t latency) { latency_ = latency; } void SequenceNode::SetII(int64_t ii) { ii_ = ii; } -void SequenceNode::AddReadRegion(const BufferRegion ®ion) { - if (!children.empty() && children[0]) { - children[0]->AddReadRegion(region); - } -} - -void SequenceNode::AddWriteRegion(const BufferRegion ®ion) { - if (!children.empty() && children[0]) { - children[0]->AddWriteRegion(region); - } -} - -std::unique_ptr SequenceNode::Clone() const { - auto new_seq = std::make_unique(); +std::shared_ptr SequenceNode::Clone() const { + auto new_seq = std::make_shared(); new_seq->children.reserve(children.size()); for (const auto &child : children) { if (child) { @@ -214,8 +226,8 @@ std::unique_ptr SequenceNode::Clone() const { return new_seq; } -std::unique_ptr TaskNode::Clone() const { - auto new_task = std::make_unique(); +std::shared_ptr TaskNode::Clone() const { + auto new_task = std::make_shared(); // Copy statements new_task->stmts = stmts; // Copy resource usage flags @@ -293,8 +305,8 @@ bool TaskNode::ContainsLoopBreak() const { return found_loop_break; } -std::unique_ptr ControlNode::Clone() const { - auto new_ctrl = std::make_unique(); +std::shared_ptr ControlNode::Clone() const { + auto new_ctrl = std::make_shared(); // Copy For control (For is a TVM object with reference counting) new_ctrl->control = control; // Clone child if exists @@ -302,15 +314,14 @@ std::unique_ptr ControlNode::Clone() const { new_ctrl->child = child->Clone(); } // Copy latency and II - new_ctrl->wrappers = wrappers; new_ctrl->SetLatency(GetLatency()); new_ctrl->SetII(GetII()); new_ctrl->SetPromote(hasPromote()); return new_ctrl; } -std::unique_ptr WrapperNode::Clone() const { - auto new_wrapper = std::make_unique(); +std::shared_ptr WrapperNode::Clone() const { + auto new_wrapper = std::make_shared(); // Copy var and value (TVM objects with reference counting) new_wrapper->wrapper = wrapper; // Clone child if exists @@ -323,8 +334,8 @@ std::unique_ptr WrapperNode::Clone() const { return new_wrapper; } -std::unique_ptr ScheduleUnit::Clone() const { - auto new_unit = std::make_unique(); +std::shared_ptr ScheduleUnit::Clone() const { + auto new_unit = std::make_shared(); // Copy var and value (TVM objects with reference counting) new_unit->stage = stage; // Clone child if exists diff --git a/src/transform/auto_schedule/ir_structure.h b/src/transform/auto_schedule/ir_structure.h index 8341e27397..6df85be596 100644 --- a/src/transform/auto_schedule/ir_structure.h +++ b/src/transform/auto_schedule/ir_structure.h @@ -28,6 +28,7 @@ class IRStructure; class TaskNode; class ControlNode; class SequenceNode; +class WrapperNode; // Structure to store region access information with warpgroup id struct RegionAccessInfo { @@ -65,7 +66,7 @@ class IRStructure { virtual ~IRStructure() = default; virtual Kind GetKind() const = 0; - virtual std::unique_ptr Clone() const = 0; + virtual std::shared_ptr Clone() const = 0; // Helper methods for safe casting bool IsTask() const { return GetKind() == Kind::kTask; } @@ -83,6 +84,10 @@ class IRStructure { virtual std::vector GetReadRegions() const = 0; virtual std::vector GetWriteRegions() const = 0; + // Variable access (used for dependency analysis) + virtual std::vector GetReadVars() const = 0; + virtual std::vector GetWriteVars() const = 0; + // Latency estimation virtual int64_t GetLatency() const = 0; // Estimated latency in cycles virtual int64_t GetII() const = 0; // Initiation interval in cycles @@ -96,10 +101,6 @@ class IRStructure { virtual void SetLatency(int64_t latency) = 0; virtual void SetII(int64_t ii) = 0; - // Helper methods to add regions (for incremental analysis) - virtual void AddReadRegion(const BufferRegion ®ion) = 0; - virtual void AddWriteRegion(const BufferRegion ®ion) = 0; - // Recursive region collection method virtual void CollectRegions( std::vector &result, @@ -142,6 +143,8 @@ class TaskNode : public IRStructure { std::vector GetWriteRegions() const override { return write_regions_; } + std::vector GetReadVars() const override { return read_vars_; } + std::vector GetWriteVars() const override { return write_vars_; } // Latency estimation int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -156,6 +159,20 @@ class TaskNode : public IRStructure { void SetWriteRegions(const std::vector ®ions) override { write_regions_ = regions; } + void SetReadVars(const std::vector &vars) { read_vars_ = vars; } + void SetWriteVars(const std::vector &vars) { write_vars_ = vars; } + void SubstituteVar(const Var &old_var, const Var &new_var) { + for (auto &var : read_vars_) { + if (var.same_as(old_var)) { + var = new_var; + } + } + for (auto &var : write_vars_) { + if (var.same_as(old_var)) { + var = new_var; + } + } + } void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } @@ -245,10 +262,10 @@ class TaskNode : public IRStructure { } // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; // Helper methods to add regions (for incremental analysis) - void AddReadRegion(const BufferRegion ®ion) override { + void AddReadRegion(const BufferRegion ®ion) { // Check for duplicate regions for (const auto &existing : read_regions_) { if (existing->buffer.same_as(region->buffer) && @@ -259,7 +276,7 @@ class TaskNode : public IRStructure { read_regions_.push_back(region); } - void AddWriteRegion(const BufferRegion ®ion) override { + void AddWriteRegion(const BufferRegion ®ion) { // Check for duplicate regions for (const auto &existing : write_regions_) { if (existing->buffer.same_as(region->buffer) && @@ -270,11 +287,16 @@ class TaskNode : public IRStructure { write_regions_.push_back(region); } + void AddReadVar(const Var &var) { read_vars_.push_back(var); } + void AddWriteVar(const Var &var) { write_vars_.push_back(var); } + void CollectRegions( std::vector &result, std::set>> &visited) const override; - bool containWarpgroupId(int id) const override { return ContainsLoopBreak() || warpgroup_id_ == id; } + bool containWarpgroupId(int id) const override { + return ContainsLoopBreak() || warpgroup_id_ == id; + } // Check if this task contains loop_break call bool ContainsLoopBreak() const; @@ -289,6 +311,9 @@ class TaskNode : public IRStructure { std::vector read_regions_; std::vector write_regions_; + std::vector read_vars_; + std::vector write_vars_; + // Latency estimation int64_t latency_{0}; // Estimated latency in cycles int64_t ii_{0}; // Initiation interval in cycles @@ -315,11 +340,7 @@ class TaskNode : public IRStructure { class ControlNode : public IRStructure { public: For control; // The For operation - std::unique_ptr child; - - // Wrapper Stmts (LetStmt/AttrStmt) that originally wrapped the For loop body. - // Stored outermost-first. Re-applied in reverse order during IR conversion. - std::vector wrappers; + std::shared_ptr child; Kind GetKind() const override { return Kind::kControl; } @@ -342,6 +363,14 @@ class ControlNode : public IRStructure { return child ? child->GetWriteRegions() : std::vector{}; } + // Variable access (aggregate from child) + std::vector GetReadVars() const override { + return child ? child->GetReadVars() : std::vector{}; + } + std::vector GetWriteVars() const override { + return child ? child->GetWriteVars() : std::vector{}; + } + // Latency estimation (aggregate from child) int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -370,16 +399,6 @@ class ControlNode : public IRStructure { void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } - // Helper methods to add regions (delegate to child) - void AddReadRegion(const BufferRegion ®ion) override { - if (child) - child->AddReadRegion(region); - } - void AddWriteRegion(const BufferRegion ®ion) override { - if (child) - child->AddWriteRegion(region); - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -389,7 +408,7 @@ class ControlNode : public IRStructure { void SetPromote(bool promote) { has_promote_ = promote; } // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { return child->containWarpgroupId(id); @@ -407,7 +426,8 @@ class ControlNode : public IRStructure { class WrapperNode : public IRStructure { public: Stmt wrapper; - std::unique_ptr child; + std::shared_ptr child; + std::shared_ptr task; Kind GetKind() const override { return Kind::kWrapper; } @@ -430,6 +450,14 @@ class WrapperNode : public IRStructure { return child ? child->GetWriteRegions() : std::vector{}; } + // Variable access (aggregate from child) + std::vector GetReadVars() const override { + return child ? child->GetReadVars() : std::vector{}; + } + std::vector GetWriteVars() const override { + return child ? child->GetWriteVars() : std::vector{}; + } + // Latency estimation (aggregate from child) int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -458,22 +486,12 @@ class WrapperNode : public IRStructure { void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } - // Helper methods to add regions (delegate to child) - void AddReadRegion(const BufferRegion ®ion) override { - if (child) - child->AddReadRegion(region); - } - void AddWriteRegion(const BufferRegion ®ion) override { - if (child) - child->AddWriteRegion(region); - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { return child->containWarpgroupId(id); @@ -489,7 +507,7 @@ class ScheduleUnit : public IRStructure { public: int stage; std::vector> before, after; - std::unique_ptr child; + std::shared_ptr child; ScheduleUnit() { for (unsigned idx = 0; idx != 2; ++idx) { @@ -519,6 +537,13 @@ class ScheduleUnit : public IRStructure { return child ? child->GetWriteRegions() : std::vector{}; } + std::vector GetReadVars() const override { + return child ? child->GetReadVars() : std::vector{}; + } + std::vector GetWriteVars() const override { + return child ? child->GetWriteVars() : std::vector{}; + } + // Latency estimation (aggregate from child) int64_t GetLatency() const override { return latency_; } int64_t GetII() const override { return ii_; } @@ -547,16 +572,6 @@ class ScheduleUnit : public IRStructure { void SetLatency(int64_t latency) override { latency_ = latency; } void SetII(int64_t ii) override { ii_ = ii; } - // Helper methods to add regions (delegate to child) - void AddReadRegion(const BufferRegion ®ion) override { - if (child) - child->AddReadRegion(region); - } - void AddWriteRegion(const BufferRegion ®ion) override { - if (child) - child->AddWriteRegion(region); - } - void CollectRegions( std::vector &result, std::set>> &visited) const override; @@ -570,7 +585,7 @@ class ScheduleUnit : public IRStructure { } // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { return child->containWarpgroupId(id); @@ -585,7 +600,7 @@ class ScheduleUnit : public IRStructure { // Sequence node: contains a vector of child IRStructures class SequenceNode : public IRStructure { public: - std::vector> children; + std::vector> children; Kind GetKind() const override { return Kind::kSequence; } @@ -598,6 +613,9 @@ class SequenceNode : public IRStructure { std::vector GetReadRegions() const override; std::vector GetWriteRegions() const override; + std::vector GetReadVars() const override; + std::vector GetWriteVars() const override; + // Latency estimation (aggregate from all children) int64_t GetLatency() const override; int64_t GetII() const override; @@ -611,16 +629,12 @@ class SequenceNode : public IRStructure { void SetLatency(int64_t latency) override; void SetII(int64_t ii) override; - // Helper methods to add regions (delegate to first child if exists) - void AddReadRegion(const BufferRegion ®ion) override; - void AddWriteRegion(const BufferRegion ®ion) override; - void CollectRegions( std::vector &result, std::set>> &visited) const override; // Clone method - std::unique_ptr Clone() const override; + std::shared_ptr Clone() const override; bool containWarpgroupId(int id) const override { for (auto &child : children) { @@ -800,13 +814,6 @@ inline void PrintAllStmts(const IRStructure *node, int indent = 0) { LOG(INFO) << indent_str << " For statement:"; LOG(INFO) << indent_str + " " << control->control; - if (!control->wrappers.empty()) { - LOG(INFO) << indent_str << " Wrappers: " << control->wrappers.size(); - for (size_t w = 0; w < control->wrappers.size(); w++) { - LOG(INFO) << indent_str << " Wrapper " << w << ": " - << control->wrappers[w]; - } - } // Recursively print child statements if (control->child) { LOG(INFO) << indent_str << " Loop body:"; @@ -851,6 +858,21 @@ inline void PrintAllStmts(const IRStructure *node, int indent = 0) { } } +inline Stmt GetLetDecl(const LetStmtNode *let_node) { + if (let_node == nullptr) { + return Stmt(); + } + return LetStmt(let_node->var, let_node->value, Evaluate(0)); +} + +inline Stmt GetAttrDecl(const AttrStmtNode *attr_node) { + if (attr_node == nullptr) { + return Stmt(); + } + return AttrStmt(attr_node->node, attr_node->attr_key, attr_node->value, + Evaluate(0)); +} + // Original helper function to print IRStructure (kept for backward // compatibility) inline void PrintIRStructure(const IRStructure *node, int indent = 0) { @@ -873,6 +895,12 @@ inline void PrintIRStructure(const IRStructure *node, int indent = 0) { for (auto ®ion : task->GetWriteRegions()) { LOG(INFO) << indent_str << " Write Region: " << region; } + for (auto &var : task->GetReadVars()) { + LOG(INFO) << indent_str << " Read Var: " << var; + } + for (auto &var : task->GetWriteVars()) { + LOG(INFO) << indent_str << " Write Var: " << var; + } LOG(INFO) << indent_str << " uses_cuda_core: " << task->UsesCUDACore(); LOG(INFO) << indent_str << " uses_tma_core: " << task->UsesTMACore(); LOG(INFO) << indent_str << " uses_tensor_core: " << task->UsesTensorCore(); @@ -882,13 +910,6 @@ inline void PrintIRStructure(const IRStructure *node, int indent = 0) { } else if (node->IsControl()) { const ControlNode *control = static_cast(node); LOG(INFO) << indent_str << "ControlNode (For loop):"; - if (!control->wrappers.empty()) { - LOG(INFO) << indent_str << " Wrappers: " << control->wrappers.size(); - for (size_t w = 0; w < control->wrappers.size(); w++) { - LOG(INFO) << indent_str << " Wrapper " << w << ": " - << control->wrappers[w]; - } - } if (control->child) { LOG(INFO) << indent_str << " Child:"; PrintIRStructure(control->child.get(), indent + 2); @@ -905,6 +926,8 @@ inline void PrintIRStructure(const IRStructure *node, int indent = 0) { const WrapperNode *wrapper = static_cast(node); LOG(INFO) << indent_str << "WrapperNode:"; LOG(INFO) << indent_str << " Wrapper: " << wrapper->wrapper; + LOG(INFO) << indent_str << " Task:"; + PrintIRStructure(wrapper->task.get(), indent + 4); if (wrapper->child) { LOG(INFO) << indent_str << " Child:"; PrintIRStructure(wrapper->child.get(), indent + 2); diff --git a/src/transform/auto_schedule/memory_detector.h b/src/transform/auto_schedule/memory_detector.h index c5e5465d81..8a37fc6eea 100644 --- a/src/transform/auto_schedule/memory_detector.h +++ b/src/transform/auto_schedule/memory_detector.h @@ -38,6 +38,8 @@ class MemoryAccessDetector : public StmtExprVisitor { hint_map_.clear(); pending_conditions_.clear(); let_bindings_.clear(); + read_vars_.clear(); + write_vars_.clear(); operator()(stmt); } @@ -50,6 +52,8 @@ class MemoryAccessDetector : public StmtExprVisitor { hint_map_.clear(); pending_conditions_.clear(); let_bindings_.clear(); + read_vars_.clear(); + write_vars_.clear(); operator()(expr); } @@ -63,6 +67,11 @@ class MemoryAccessDetector : public StmtExprVisitor { return CollectRegions(write_buffers_, write_regions_); } + // Return all variables that are read from + std::vector GetReadVars() const { return read_vars_; } + // Return all variables that are written to + std::vector GetWriteVars() const { return write_vars_; } + private: /*! \brief Iteration range for loop_vars */ std::unordered_map dom_map_; @@ -83,6 +92,11 @@ class MemoryAccessDetector : public StmtExprVisitor { /*! \brief let bindings inside the block */ std::unordered_map let_bindings_; + /*! \brief The set of variables that are read in the current block. */ + std::vector read_vars_; + /*! \brief The set of variables that are written in the current block. */ + std::vector write_vars_; + /*! * \brief Update read/write buffers and regions with provided buffer and * region @@ -106,6 +120,30 @@ class MemoryAccessDetector : public StmtExprVisitor { regions->push_back(region); } + /*! + * \brief Update the set of read variables with the given variable + * \param var The variable to add to the set of read variables + */ + void UpdateReadVar(const Var &var) { + for (const auto &v : read_vars_) { + if (v.same_as(var)) + return; + } + read_vars_.push_back(var); + } + + /*! + * \brief Update the set of write variables with the given variable + * \param var The variable to add to the set of write variables + */ + void UpdateWriteVar(const Var &var) { + for (const auto &v : write_vars_) { + if (v.same_as(var)) + return; + } + write_vars_.push_back(var); + } + /*! * \brief Process a buffer region argument from reduce operation * \param arg The argument which could be BufferRegion, BufferLoad, or @@ -227,10 +265,16 @@ class MemoryAccessDetector : public StmtExprVisitor { void VisitStmt_(const LetStmtNode *op) override { let_bindings_[op->var.get()] = op->value; - StmtExprVisitor::VisitStmt_(op); + UpdateWriteVar(op->var); + StmtExprVisitor::VisitStmt(op->body); let_bindings_.erase(op->var.get()); } + void VisitExpr_(const VarNode *op) override { + UpdateReadVar(tvm::ffi::GetRef(op)); + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const BufferLoadNode *op) override { std::vector relaxed_region; size_t num_indices = op->indices.size(); diff --git a/src/transform/if_condition_extract.cc b/src/transform/if_condition_extract.cc new file mode 100644 index 0000000000..bca573b348 --- /dev/null +++ b/src/transform/if_condition_extract.cc @@ -0,0 +1,131 @@ +/*! + * \file if_condition_extract.cc + * \brief Extract if conditions into temporary LetStmt variables, then expand if + * statements to all branches. + */ + +#include +#include +#include +#include +#include +#include + +#include "../op/builtin.h" + +namespace tvm { +namespace tl { + +using namespace tir; + +class IfConditionExtractor : public StmtExprMutator { +public: + static PrimFunc Substitute(PrimFunc &f) { + auto rewriter = IfConditionExtractor(); + f.CopyOnWrite()->body = rewriter(f->body); + return f; + } + +private: + IfConditionExtractor() = default; + + // counter to generate unique name for each IfStmt + int counter_ = 0; + + //! \brief Check if the expression is a simple variable. + bool IsSimpleVar(const PrimExpr &expr) { + return expr.as() != nullptr; + } + + Stmt VisitStmt_(const IfThenElseNode *op) final { + PrimExpr condition = VisitExpr(op->condition); + Stmt then_case = VisitStmt(op->then_case); + Optional else_case = op->else_case; + if (else_case.defined()) { + else_case = VisitStmt(else_case.value()); + } + + std::string var_name = "__cond_" + std::to_string(counter_++); + Var cond_var(var_name, DataType::Bool()); + bool is_simple = true; + + if (IsSimpleVar(condition)) { + // If the condition is already a simple variable, no need to extract it. + cond_var = Downcast(condition); + } else { + is_simple = false; + } + + auto bind_cond_var = [](const Stmt &sentence, const Var &cond) -> Stmt { + if (auto if_sentence = sentence.as()) { + PrimExpr new_cond = cond & if_sentence->condition; + return IfThenElse(new_cond, if_sentence->then_case, + if_sentence->else_case); + } else { + return IfThenElse(cond, sentence); + } + }; + + auto bind_cond_var_body = [&](const Optional &body, + const Var &cond) -> Stmt { + if (!body.defined()) { + return Stmt(); + } + if (auto seq = body.as()) { + Array new_seq; + for (auto sentence : seq->seq) { + new_seq.push_back(bind_cond_var(sentence, cond)); + } + return SeqStmt(std::move(new_seq)); + } else { + return bind_cond_var(body.value(), cond); + } + }; + + Array new_seq; + new_seq.insert(new_seq.end(), bind_cond_var_body(then_case, cond_var)); + if (else_case.defined()) + new_seq.insert(new_seq.end(), bind_cond_var_body(else_case, cond_var)); + + Stmt body = + new_seq.empty() + ? Stmt() + : (new_seq.size() == 1 ? new_seq[0] : SeqStmt(std::move(new_seq))); + if (is_simple) { + return body; + } else { + return LetStmt(cond_var, condition, body); + } + } + + Stmt VisitStmt_(const SeqStmtNode *op) final { + Array seq; + for (auto stmt : op->seq) { + auto new_stmt = VisitStmt(stmt); + if (!new_stmt.defined()) + continue; + if (auto seq_node = new_stmt.as()) { + seq.insert(seq.end(), seq_node->seq.begin(), seq_node->seq.end()); + } else { + seq.push_back(new_stmt); + } + } + return SeqStmt(std::move(seq)); + } +}; + +using namespace tir::transform; +tvm::transform::Pass IfConditionExtract() { + auto pass_func = [=](PrimFunc f, const IRModule &m, const PassContext &ctx) { + return IfConditionExtractor::Substitute(f); + }; + return CreatePrimFuncPass(pass_func, 0, "tl.IfConditionExtract", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.IfConditionExtract", IfConditionExtract); +} + +} // namespace tl +} // namespace tvm diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py index 8a1364e94f..9ce8b28187 100644 --- a/tilelang/engine/phase.py +++ b/tilelang/engine/phase.py @@ -182,7 +182,7 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule: mod = tilelang.transform.Simplify()(mod) if allow_autoschedule(): # Auto schedule for high-level operations - mod = tilelang.transform.IfStmtBinding()(mod) + mod = tilelang.transform.IfConditionExtract()(mod) mod = tilelang.transform.AutoSchedule(False)(mod) mod = tilelang.transform.Simplify()(mod) # Set layouts for reducers diff --git a/tilelang/transform/__init__.py b/tilelang/transform/__init__.py index 773b7425cf..689258cca0 100644 --- a/tilelang/transform/__init__.py +++ b/tilelang/transform/__init__.py @@ -562,6 +562,11 @@ def LowerSharedTmem(): return _ffi_api.LowerSharedTmem() # type: ignore +def IfConditionExtract(): + """Extract if condition to LetStmt variables.""" + return _ffi_api.IfConditionExtract() # type: ignore + + def AutoSchedule(enable_epi: bool): """Auto schedule for high-level operations""" return _ffi_api.AutoSchedule(enable_epi) # type: ignore