implement cpu issue softmax

fishingguy456 · Mar 9, 2024 · eafa10f · eafa10f
1 parent 713b016
commit eafa10f
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/python/hidet/graph/ops/softmax.py b/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: