Skip to content

Commit 6fc627e

Browse files
authored
[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.
1 parent 3340d89 commit 6fc627e

File tree

1 file changed

+3
-5
lines changed

1 file changed

+3
-5
lines changed

src/ir.cc

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
162162
attrs.defined() && attrs.count(tilelang_is_cpu_kernel_frame);
163163

164164
if (is_cpu_kernel_frame) {
165+
// Launch CPU Kernel
165166
ICHECK(grid_size.size() >= 0);
166167
ICHECK(block_size.size() == 0) << "CPU kernel cannot have block size";
167168
ICHECK(attrs.defined());
@@ -170,7 +171,6 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
170171
n->frames.push_back(
171172
MakeIterVarFrame("block_var_" + std::to_string(i), grid_size[i]));
172173
}
173-
// Launch CPU Kernel
174174
} else {
175175
// Launch GPU Kernel
176176
ICHECK(grid_size.size() <= 3);
@@ -203,17 +203,15 @@ KernelLaunchFrame KernelLaunch(Array<PrimExpr> grid_size,
203203
CreateEnvThread("tz", "threadIdx.z", block_size[2].dtype()),
204204
block_size[2]));
205205
}
206-
} else {
207-
n->frames.push_back(Block(""));
208206
}
209207
}
210208

211209
if (attrs.defined()) {
212-
auto empty_block = Block("");
210+
auto empty_block = Block("root");
213211
empty_block->annotations = attrs;
214212
n->frames.push_back(empty_block);
215213
} else {
216-
n->frames.push_back(Block(""));
214+
n->frames.push_back(Block("root"));
217215
}
218216

219217
return KernelLaunchFrame(n);

0 commit comments

Comments
 (0)