[Task] More descriptive kernel names for nsys/ncu (#315)

Co-authored-by: Allan Lin <allan.lin@centml.ai>
hidet-org · Jul 16, 2023 · 0c18446 · 0c18446
1 parent 15426c8
commit 0c18446
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/python/hidet/ir/schedulers/cpu/scheduler.py b/python/hidet/ir/schedulers/cpu/scheduler.py
@@ -26,7 +26,12 @@ def schedule_grid_compute(self, node: GridCompute, tensor_map: Dict[TensorNode,
 
         params, param_map, call_args = self.grid_compute_params_and_args(node, tensor_map)
 
-        with FunctionBuilder(name=f'compute_{node.name}', kind='cpu_kernel') as fb:
+        if self.task is not None:
+            name = f'{self.task.name}_compute_{node.name}'
+        else:
+            name = f'compute_{node.name}'
+
+        with FunctionBuilder(name=name, kind='cpu_kernel') as fb:
             # set function parameters
             fb.extend_params(params)
 

diff --git a/python/hidet/ir/schedulers/cuda/scheduler.py b/python/hidet/ir/schedulers/cuda/scheduler.py
@@ -34,9 +34,12 @@ def schedule_grid_compute(self, node: GridCompute, tensor_map: Dict[TensorNode,
         block_dim = 512
         grid_dim: Expr = (prod(node.shape) + block_dim - 1) // block_dim
 
-        with FunctionBuilder(
-            name=f'compute_{node.name}', kind='cuda_kernel', grid_dim=grid_dim, block_dim=block_dim
-        ) as fb:
+        if self.task is not None:
+            name = f'{self.task.name}_compute_{node.name}'
+        else:
+            name = f'compute_{node.name}'
+
+        with FunctionBuilder(name=name, kind='cuda_kernel', grid_dim=grid_dim, block_dim=block_dim) as fb:
             # set function parameters
             fb.extend_params(params)