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, ...@@ -159,6 +159,7 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
attrs.defined() && attrs.count(tilelang_is_cpu_kernel_frame); attrs.defined() && attrs.count(tilelang_is_cpu_kernel_frame);
if (is_cpu_kernel_frame) { if (is_cpu_kernel_frame) {
// Launch CPU Kernel
ICHECK(grid_size.size() >= 0); ICHECK(grid_size.size() >= 0);
ICHECK(block_size.size() == 0) << "CPU kernel cannot have block size"; ICHECK(block_size.size() == 0) << "CPU kernel cannot have block size";
ICHECK(attrs.defined()); ICHECK(attrs.defined());
...@@ -167,7 +168,6 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size, ...@@ -167,7 +168,6 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
n->frames.push_back( n->frames.push_back(
MakeIterVarFrame("block_var_" + std::to_string(i), grid_size[i])); MakeIterVarFrame("block_var_" + std::to_string(i), grid_size[i]));
} }
// Launch CPU Kernel
} else { } else {
// Launch GPU Kernel // Launch GPU Kernel
ICHECK(grid_size.size() <= 3); ICHECK(grid_size.size() <= 3);
...@@ -200,17 +200,15 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size, ...@@ -200,17 +200,15 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
CreateEnvThread("tz", "threadIdx.z", block_size[2].dtype()), CreateEnvThread("tz", "threadIdx.z", block_size[2].dtype()),
block_size[2])); block_size[2]));
} }
} else {
n->frames.push_back(Block(""));
} }
} }
if (attrs.defined()) { if (attrs.defined()) {
auto empty_block = Block(""); auto empty_block = Block("root");
empty_block->annotations = attrs; empty_block->annotations = attrs;
n->frames.push_back(empty_block); n->frames.push_back(empty_block);
} else { } else {
n->frames.push_back(Block("")); n->frames.push_back(Block("root"));
} }
return KernelLaunchFrame(n); 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