Commit dabe6e0a authored by Lei Wang's avatar Lei Wang Committed by LeiWang1999
Browse files

[Refactor] Update KernelLaunch to clarify CPU and GPU kernel launch logic (#441)

* Added comments to distinguish between CPU and GPU kernel launch sections for better code readability.
* Changed the creation of empty blocks to use a consistent "root" identifier, enhancing clarity in frame management.
parent 72fea0a2
......@@ -159,6 +159,7 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
attrs.defined() && attrs.count(tilelang_is_cpu_kernel_frame);
if (is_cpu_kernel_frame) {
// Launch CPU Kernel
ICHECK(grid_size.size() >= 0);
ICHECK(block_size.size() == 0) << "CPU kernel cannot have block size";
ICHECK(attrs.defined());
......@@ -167,7 +168,6 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
n->frames.push_back(
MakeIterVarFrame("block_var_" + std::to_string(i), grid_size[i]));
}
// Launch CPU Kernel
} else {
// Launch GPU Kernel
ICHECK(grid_size.size() <= 3);
......@@ -200,17 +200,15 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
CreateEnvThread("tz", "threadIdx.z", block_size[2].dtype()),
block_size[2]));
}
} else {
n->frames.push_back(Block(""));
}
}
if (attrs.defined()) {
auto empty_block = Block("");
auto empty_block = Block("root");
empty_block->annotations = attrs;
n->frames.push_back(empty_block);
} else {
n->frames.push_back(Block(""));
n->frames.push_back(Block("root"));
}
return KernelLaunchFrame(n);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment