fused scheduler task name

fishingguy456 · Mar 9, 2024 · 8f73d4c · 8f73d4c
1 parent 703f49e
commit 8f73d4c
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/python/hidet/ir/schedulers/cpu/scheduler.py b/python/hidet/ir/schedulers/cpu/scheduler.py
@@ -26,7 +26,13 @@ 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)
 
         if self.task is not None:
-            name = f'{self.task.name}_compute_{node.name}'
+            from hidet.graph.ops.fusion.fused_operator import FusedTask
+
+            if isinstance(self.task, FusedTask):
+                fused_name = self.task.attrs['fused_ops'].replace(' ', '_')
+                name = f'fused_{fused_name}_{node.name}'
+            else:
+                name = f'{self.task.name}_{node.name}'
         else:
             name = f'compute_{node.name}'