diff --git a/src/transform/auto_schedule.cc b/src/transform/auto_schedule.cc index 664407f17b..225fcc19bd 100644 --- a/src/transform/auto_schedule.cc +++ b/src/transform/auto_schedule.cc @@ -74,7 +74,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( IRStructure *root, IterVar thread_var, std::vector &barrier_buffers, Map &barrier_map, const bool enable_epi, PrimExpr thread_count[2], bool producer_consumer, - const WarpSpecializeConfig &config); + const WarpSpecializeConfig &config, Buffer neutral_sync_shared_barrier); Stmt ConvertIRStructureToStmt(IRStructure *root, const bool outer_enable_epi); @@ -1294,6 +1294,8 @@ class IRStructureBuilder : public StmtVisitor { bool found_tensor{false}; bool found_cuda{false}; + bool found_tma_load{false}; + // Tensor Core shape information (multiple shapes possible) struct TensorCoreShape { int64_t m; @@ -1344,6 +1346,9 @@ class IRStructureBuilder : public StmtVisitor { MemoryType mem_type = GetMemoryTypeFromScope(scope); if (mem_type == MemoryType::kGlobal) { found_global = true; + if (idx == 0) { + found_tma_load = true; + } } } } @@ -1426,6 +1431,9 @@ class IRStructureBuilder : public StmtVisitor { // 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); @@ -1610,9 +1618,15 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { } LoopNestingInfo loop_info; std::vector buffer_infos; - AnalyzeAndInsertBarriers(ir_structure.get(), next_barrier_id, - barrier_buffers, barrier_map, thread_count, - loop_info, buffer_infos); + PrimExpr barrier_count = config.enable_thread_extend + ? thread_count[0] + thread_count[1] + : thread_var->dom->extent; + Buffer neutral_sync_shared_barrier = + makeBarrierBuffer(barrier_count, "neutral_sync_shared_barrier", 1, + barrier_buffers, barrier_map); + AnalyzeAndInsertBarriers( + ir_structure.get(), next_barrier_id, barrier_buffers, barrier_map, + thread_count, loop_info, buffer_infos, neutral_sync_shared_barrier); // Print the modified summary view // PrintIRStructure(ir_structure.get()); @@ -1620,7 +1634,8 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) { // Apply warpgroup partition to entire IRStructure Stmt new_body = ApplyWarpgroupPartitionToIRStructure( ir_structure.get(), thread_var, barrier_buffers, barrier_map, - enable_epi, thread_count, double_thread, config); + enable_epi, thread_count, double_thread, config, + neutral_sync_shared_barrier); if (config.enable_thread_extend) { // sm_90: may need to update thread extent @@ -2053,7 +2068,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( IRStructure *root, IterVar thread_var, std::vector &barrier_buffers, Map &barrier_map, const bool outer_enable_epi, PrimExpr thread_count[2], bool producer_consumer, - const WarpSpecializeConfig &config) { + const WarpSpecializeConfig &config, Buffer neutral_sync_shared_barrier) { if (!root) return Evaluate(0); @@ -2063,7 +2078,8 @@ Stmt ApplyWarpgroupPartitionToIRStructure( if (wrapper->child) { body = ApplyWarpgroupPartitionToIRStructure( wrapper->child.get(), thread_var, barrier_buffers, barrier_map, - outer_enable_epi, thread_count, producer_consumer, config); + outer_enable_epi, thread_count, producer_consumer, config, + neutral_sync_shared_barrier); } if (const auto *let = wrapper->wrapper.as()) { return LetStmt(let->var, let->value, body); @@ -2601,7 +2617,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( // synchronization pro_and_warpgroup_stmt = InsertBarriersForNeutralSync( pro_neutral_body, if_then_else, barrier_buffers, barrier_map, - barrier_count); + barrier_count, neutral_sync_shared_barrier); } else if (!IsEvaluateZero(if_then_else) || !IsEvaluateZero(pro_neutral_body)) { // Only one has actual statements @@ -2649,7 +2665,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure( combined_stmt = InsertBarriersForNeutralSyncWithDependency( pro_and_warpgroup_stmt, epi_neutral_body, barrier_buffers, barrier_map, barrier_count, need_shared_barrier_for_epi, need_tmem_barrier_for_epi, - thread_var->var, 0, thread_count[0]); + Buffer(), thread_var->var, 0, thread_count[0]); } else if (!IsEvaluateZero(epi_neutral_body)) { combined_stmt = epi_neutral_body; } else { diff --git a/src/transform/auto_schedule/barrier.h b/src/transform/auto_schedule/barrier.h index cf41710f63..97c873dee8 100644 --- a/src/transform/auto_schedule/barrier.h +++ b/src/transform/auto_schedule/barrier.h @@ -151,20 +151,23 @@ AnalyzeAndInsertBarriers(IRStructure *node, int &next_barrier_id, std::vector &barrier_buffers, Map &barrier_map, PrimExpr thread_count[2], LoopNestingInfo &loop_info, - std::vector &buffer_infos); + std::vector &buffer_infos, + Buffer neutral_sync_shared_barrier); static void AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id, std::vector &barrier_buffers, Map &barrier_map, PrimExpr thread_count[2], LoopNestingInfo &loop_info, - std::vector &buffer_infos); + std::vector &buffer_infos, + Buffer neutral_sync_shared_barrier); static void AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id, std::vector &barrier_buffers, Map &barrier_map, PrimExpr thread_count[2], LoopNestingInfo &loop_info, - std::vector &buffer_infos); + std::vector &buffer_infos, + Buffer neutral_sync_shared_barrier); // Create a barrier_arrive statement for the given barrier expression // Equivalent to T.barrier_arrive(barrier_expr) in Python @@ -251,7 +254,8 @@ static Stmt InsertBarriersForNeutralSyncWithDependency( Stmt producer_body, Stmt consumer_body, std::vector &barrier_buffers, Map &barrier_map, PrimExpr total_thread_count, - bool need_regular_barrier, bool need_tmem_barrier, Var thread_var = Var(), + bool need_regular_barrier, bool need_tmem_barrier, + Buffer neutral_sync_shared_barrier = Buffer(), Var thread_var = Var(), PrimExpr tensor_core_wg_start = PrimExpr(), PrimExpr tensor_core_wg_end = PrimExpr()) { if (IsEvaluateZero(producer_body) || IsEvaluateZero(consumer_body)) { @@ -267,8 +271,11 @@ static Stmt InsertBarriersForNeutralSyncWithDependency( if (need_regular_barrier) { Buffer barrier_buffer = - makeBarrierBuffer(total_thread_count, "neutral_sync_shared_barrier", 1, - barrier_buffers, barrier_map); + neutral_sync_shared_barrier.defined() + ? neutral_sync_shared_barrier + : makeBarrierBuffer(total_thread_count, + "neutral_sync_shared_barrier", 1, + barrier_buffers, barrier_map); PrimExpr barrier_load = BufferLoad(barrier_buffer, {0}); arrive_stmts.push_back(makeBarrierArrive(barrier_load)); wait_stmts.push_back(makeBarrierWait(barrier_load, 0)); @@ -307,10 +314,11 @@ static Stmt InsertBarriersForNeutralSyncWithDependency( static Stmt InsertBarriersForNeutralSync(Stmt neutral_body, Stmt warpgroup_body, std::vector &barrier_buffers, Map &barrier_map, - PrimExpr total_thread_count) { + PrimExpr total_thread_count, + Buffer neutral_sync_shared_barrier) { return InsertBarriersForNeutralSyncWithDependency( neutral_body, warpgroup_body, barrier_buffers, barrier_map, - total_thread_count, true, false); + total_thread_count, true, false, neutral_sync_shared_barrier); } // StmtExprMutator to rewrite BufferLoad/BufferStore for multi-version buffers @@ -473,6 +481,33 @@ static void RewriteGemmMbar(TaskNode *task, PrimExpr mbar_expr) { } } +static void RewriteCopyMbar(TaskNode *task, PrimExpr mbar_expr) { + static const auto copy_op = Op::Get("tl.tileop.copy"); + static const auto tma_copy_op = Op::Get("tl.tileop.tma_copy"); + + class CopyMbarRewriter : public StmtExprMutator { + public: + CopyMbarRewriter(PrimExpr mbar_expr) : mbar_expr_(std::move(mbar_expr)) {} + + private: + PrimExpr VisitExpr_(const CallNode *op) override { + if (op->op.same_as(copy_op)) { + auto new_ann = op->annotations; + new_ann.Set("barrier", mbar_expr_); + return Call(op->dtype, tma_copy_op, op->args, new_ann, op->span); + } + return StmtExprMutator::VisitExpr_(op); + } + + PrimExpr mbar_expr_; + }; + + CopyMbarRewriter rewriter(mbar_expr); + for (auto &stmt : task->stmts) { + stmt = rewriter(stmt); + } +} + // Helper function to insert a statement into ScheduleUnit's stmts static void InsertStatementIntoScheduleUnit(ScheduleUnit *task, const Stmt &stmt, bool at_beginning, @@ -490,23 +525,26 @@ AnalyzeAndInsertBarriers(IRStructure *node, int &next_barrier_id, std::vector &barrier_buffers, Map &barrier_map, PrimExpr thread_count[2], LoopNestingInfo &loop_info, - std::vector &buffer_infos) { + std::vector &buffer_infos, + Buffer neutral_sync_shared_barrier) { if (!node) return; if (node->IsSequence()) { AnalyzeSequenceNodeBarriers(static_cast(node), next_barrier_id, barrier_buffers, barrier_map, - thread_count, loop_info, buffer_infos); + thread_count, loop_info, buffer_infos, + neutral_sync_shared_barrier); } else if (node->IsControl()) { AnalyzeControlNodeBarriers(static_cast(node), next_barrier_id, barrier_buffers, barrier_map, - thread_count, loop_info, buffer_infos); + thread_count, loop_info, buffer_infos, + neutral_sync_shared_barrier); } else if (node->IsWrapper()) { auto wrapper = static_cast(node); - AnalyzeAndInsertBarriers(wrapper->child.get(), next_barrier_id, - barrier_buffers, barrier_map, thread_count, - loop_info, buffer_infos); + AnalyzeAndInsertBarriers( + wrapper->child.get(), next_barrier_id, barrier_buffers, barrier_map, + thread_count, loop_info, buffer_infos, neutral_sync_shared_barrier); } else if (node->IsTask()) { // For TaskNode, nothing to do at this level } else { @@ -520,7 +558,8 @@ AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id, Map &barrier_map, PrimExpr thread_count[2], LoopNestingInfo &loop_info, - std::vector &buffer_infos) { + std::vector &buffer_infos, + Buffer neutral_sync_shared_barrier) { if (!seq) return; @@ -542,9 +581,9 @@ AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id, auto task = static_cast(promote_child.get()); if (task->child->IsSequence() || task->child->IsControl()) { // If child is SequenceNode or ControlNode, recursively analyze it - AnalyzeAndInsertBarriers(task->child.get(), next_barrier_id, - barrier_buffers, barrier_map, thread_count, - loop_info, buffer_infos); + AnalyzeAndInsertBarriers( + task->child.get(), next_barrier_id, barrier_buffers, barrier_map, + thread_count, loop_info, buffer_infos, neutral_sync_shared_barrier); } // Allocate barrier for TCGEN05MMA and rewrite gemm mbar argument @@ -563,6 +602,29 @@ AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id, } } + // Allocate barrier for TMA + if (task->isInnerTask() && task->UsesTMACore()) { + auto child = static_cast(task->child.get()); + if (child->HasTMALoad()) { + int wg_id = child->GetWarpgroupId(); + if (wg_id != -1) { + int barrier_id = next_barrier_id++; + Buffer barrier_buffer = makeBarrierBuffer( + thread_count[wg_id], "tma_barrier_" + std::to_string(barrier_id), + 1, barrier_buffers, barrier_map); + barrier_unit_map[task] = barrier_buffer; + + PrimExpr barrier_load = BufferLoad(barrier_buffer, {0}); + RewriteCopyMbar(child, barrier_load); + Stmt arrive_stmt = makeBarrierArrive(barrier_load); + InsertStatementIntoScheduleUnit(task, arrive_stmt, false, wg_id); + } else { + PrimExpr barrier_load = BufferLoad(neutral_sync_shared_barrier, {0}); + RewriteCopyMbar(child, barrier_load); + } + } + } + // Check regions for dependencies for (const auto ®ion_access : task->GetReadWriteRegions()) { int wg_id = region_access.warpgroup_id; @@ -706,7 +768,8 @@ AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id, std::vector &barrier_buffers, Map &barrier_map, PrimExpr thread_count[2], LoopNestingInfo &loop_info, - std::vector &buffer_infos) { + std::vector &buffer_infos, + Buffer neutral_sync_shared_barrier) { if (!ctrl || !ctrl->child) return; @@ -744,9 +807,9 @@ AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id, auto task = static_cast(child.get()); if (task->child->IsSequence() || task->child->IsControl()) { // If child is SequenceNode or ControlNode, recursively analyze it - AnalyzeAndInsertBarriers(task->child.get(), next_barrier_id, - barrier_buffers, barrier_map, thread_count, - loop_info, buffer_infos); + AnalyzeAndInsertBarriers( + task->child.get(), next_barrier_id, barrier_buffers, barrier_map, + thread_count, loop_info, buffer_infos, neutral_sync_shared_barrier); } all_tasks.push_back(task); } @@ -842,6 +905,28 @@ AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id, } } + // Allocate barrier for TMA + if (iter == 0 && task->isInnerTask() && task->UsesTMACore()) { + auto child = static_cast(task->child.get()); + if (child->HasTMALoad()) { + int wg_id = child->GetWarpgroupId(); + ICHECK(wg_id != -1) << "TMA loads must have valid warpgroup id"; + int barrier_id = next_barrier_id++; + Buffer barrier_buffer = + makeBarrierBuffer(thread_count[wg_id], + "tma_barrier_" + std::to_string(barrier_id), + num_stages, barrier_buffers, barrier_map); + barrier_unit_map[task] = barrier_buffer; + + PrimExpr version_index = + indexmod(loop_info.CalculateIterationCount(), num_stages); + PrimExpr barrier_load = BufferLoad(barrier_buffer, {version_index}); + RewriteCopyMbar(child, barrier_load); + Stmt arrive_stmt = makeBarrierArrive(barrier_load); + InsertStatementIntoScheduleUnit(task, arrive_stmt, false, wg_id); + } + } + // Check regions for dependencies for (const auto ®ion_access : task->GetReadWriteRegions()) { int wg_id = region_access.warpgroup_id; @@ -1032,9 +1117,9 @@ AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id, } } } else { - AnalyzeAndInsertBarriers(ctrl->child.get(), next_barrier_id, - barrier_buffers, barrier_map, thread_count, - loop_info, buffer_infos); + AnalyzeAndInsertBarriers( + ctrl->child.get(), next_barrier_id, barrier_buffers, barrier_map, + thread_count, loop_info, buffer_infos, neutral_sync_shared_barrier); } // Remove this loop from nesting info when exiting diff --git a/src/transform/auto_schedule/ir_structure.h b/src/transform/auto_schedule/ir_structure.h index c94fa66a5c..2fce54bb78 100644 --- a/src/transform/auto_schedule/ir_structure.h +++ b/src/transform/auto_schedule/ir_structure.h @@ -163,6 +163,10 @@ class TaskNode : public IRStructure { void SetWarpgroupId(int warpgroup_id) { warpgroup_id_ = warpgroup_id; } int GetWarpgroupId() const override { return warpgroup_id_; } + // TMA load flag + void SetHasTMALoad(bool value) { has_tma_load_ = value; } + bool HasTMALoad() const { return has_tma_load_; } + // Tensor Core shape information structure struct TensorCoreShape { int64_t m; @@ -291,6 +295,9 @@ class TaskNode : public IRStructure { int warpgroup_id_{ -1}; // Warpgroup id for warpgroup specialization (-1 means unassigned) + // TMA information + bool has_tma_load_{false}; + // Tensor Core shape information (M, N, K dimensions) // A task may contain multiple Tensor Core operations with different shapes std::vector tensor_core_shapes_;