Skip to content

rewrite copy to tma_copy #2

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 2 commits into from
Mar 20, 2026
Merged

rewrite copy to tma_copy #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ef48e3f
rewrite copy to tma_copy
AutumnKite Mar 20, 2026
5ec5426
move TMA load detecting into IRStructureBuilder
AutumnKite Mar 20, 2026
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: 25 additions & 9 deletions src/transform/auto_schedule.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure(
IRStructure *root, IterVar thread_var, std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &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);

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
}
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -1610,17 +1618,24 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) {
}
LoopNestingInfo loop_info;
std::vector<MultiVersionBufferInfo> 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());

// 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
Expand Down Expand Up @@ -2053,7 +2068,7 @@ Stmt ApplyWarpgroupPartitionToIRStructure(
IRStructure *root, IterVar thread_var, std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &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);

Expand All @@ -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<LetStmtNode>()) {
return LetStmt(let->var, let->value, body);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
135 changes: 110 additions & 25 deletions src/transform/auto_schedule/barrier.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,20 +151,23 @@ AnalyzeAndInsertBarriers(IRStructure *node, int &next_barrier_id,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &barrier_map,
PrimExpr thread_count[2], LoopNestingInfo &loop_info,
std::vector<MultiVersionBufferInfo> &buffer_infos);
std::vector<MultiVersionBufferInfo> &buffer_infos,
Buffer neutral_sync_shared_barrier);
static void
AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &barrier_map,
PrimExpr thread_count[2],
LoopNestingInfo &loop_info,
std::vector<MultiVersionBufferInfo> &buffer_infos);
std::vector<MultiVersionBufferInfo> &buffer_infos,
Buffer neutral_sync_shared_barrier);
static void
AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &barrier_map,
PrimExpr thread_count[2], LoopNestingInfo &loop_info,
std::vector<MultiVersionBufferInfo> &buffer_infos);
std::vector<MultiVersionBufferInfo> &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
Expand Down Expand Up @@ -251,7 +254,8 @@ static Stmt InsertBarriersForNeutralSyncWithDependency(
Stmt producer_body, Stmt consumer_body,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &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)) {
Expand All @@ -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));
Expand Down Expand Up @@ -307,10 +314,11 @@ static Stmt InsertBarriersForNeutralSyncWithDependency(
static Stmt InsertBarriersForNeutralSync(Stmt neutral_body, Stmt warpgroup_body,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &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
Expand Down Expand Up @@ -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,
Expand All @@ -490,23 +525,26 @@ AnalyzeAndInsertBarriers(IRStructure *node, int &next_barrier_id,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &barrier_map,
PrimExpr thread_count[2], LoopNestingInfo &loop_info,
std::vector<MultiVersionBufferInfo> &buffer_infos) {
std::vector<MultiVersionBufferInfo> &buffer_infos,
Buffer neutral_sync_shared_barrier) {
if (!node)
return;

if (node->IsSequence()) {
AnalyzeSequenceNodeBarriers(static_cast<SequenceNode *>(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<ControlNode *>(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<WrapperNode *>(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 {
Expand All @@ -520,7 +558,8 @@ AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id,
Map<ObjectRef, ObjectRef> &barrier_map,
PrimExpr thread_count[2],
LoopNestingInfo &loop_info,
std::vector<MultiVersionBufferInfo> &buffer_infos) {
std::vector<MultiVersionBufferInfo> &buffer_infos,
Buffer neutral_sync_shared_barrier) {
if (!seq)
return;

Expand All @@ -542,9 +581,9 @@ AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id,
auto task = static_cast<ScheduleUnit *>(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
Expand All @@ -563,6 +602,29 @@ AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id,
}
}

// Allocate barrier for TMA
if (task->isInnerTask() && task->UsesTMACore()) {
auto child = static_cast<TaskNode *>(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 &region_access : task->GetReadWriteRegions()) {
int wg_id = region_access.warpgroup_id;
Expand Down Expand Up @@ -706,7 +768,8 @@ AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id,
std::vector<Buffer> &barrier_buffers,
Map<ObjectRef, ObjectRef> &barrier_map,
PrimExpr thread_count[2], LoopNestingInfo &loop_info,
std::vector<MultiVersionBufferInfo> &buffer_infos) {
std::vector<MultiVersionBufferInfo> &buffer_infos,
Buffer neutral_sync_shared_barrier) {
if (!ctrl || !ctrl->child)
return;

Expand Down Expand Up @@ -744,9 +807,9 @@ AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id,
auto task = static_cast<ScheduleUnit *>(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);
}
Expand Down Expand Up @@ -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<TaskNode *>(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 &region_access : task->GetReadWriteRegions()) {
int wg_id = region_access.warpgroup_id;
Expand Down Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions src/transform/auto_schedule/ir_structure.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<TensorCoreShape> tensor_core_shapes_;
Expand Down