workaround for invalid split_k_slice

masahi · Feb 4, 2022 · b141271 · b141271
1 parent 6e4c7e1
commit b141271
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -363,6 +363,9 @@ def evaluate(self, op, args):
         try:
             sp = subprocess.run(cmd, capture_output=True, check=True)
             rt = float(sp.stdout)
+            if rt == 0.0:
+                # This seems to happen with split-k using invalid split-k-slices
+                rt = float("inf")
             logger.info("%s, %f", op_name, rt)
         except subprocess.CalledProcessError:
             rt = float("inf")