[Fixbug] Fix for softmmax cpu causing issues

python/hidet/graph/ops/softmax.py

@@ -157,11 +157,6 @@ def softmax_kernel(xs: xdtype[shape], ys: xdtype[shape]):
 
         return ir_module
 
-    def implement_cpu(self, working_dir: str) -> Union[IRModule, List[IRModule]]:
-        if self.inputs[0].type.dtype != float32:
-            return NotImplemented  # use auto-scheduler
-        return tune.extract_ir_modules(self.schedule_softmax_cpu)
-
 
 class CPUSoftmaxTask(SoftmaxTask):
     def allow_epilogue(self) -> bool:
@@ -170,6 +165,11 @@ def allow_epilogue(self) -> bool:
     def allow_prologue(self) -> bool:
         return False
 
+    def implement_cpu(self, working_dir: str) -> Union[IRModule, List[IRModule]]:
+        if self.inputs[0].type.dtype != float32:
+            return NotImplemented  # use auto-scheduler
+        return tune.extract_ir_modules(self.schedule_softmax_cpu)
+
     @tune.space(2, nthreads=['', 4, 8, 16, 32, 64, 96])
     @tune.space(1, nthreads=['', 8, 16])
     def schedule_softmax_cpu(self, nthreads='') -> IRModule:

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}'