skip cuda compile on old gpu

apache · May 18, 2022 · 401a820 · 401a820
1 parent 90e01fd
commit 401a820
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/python/unittest/test_tir_schedule_tensorize_ldmatrix_mma.py b/tests/python/unittest/test_tir_schedule_tensorize_ldmatrix_mma.py
@@ -77,8 +77,8 @@ def maybe_swap(i, j):
 
 def is_ampere_or_newer():
     arch = tvm.contrib.nvcc.get_target_compute_version()
-    major, minor = tvm.contrib.nvcc.parse_compute_version(arch)
-    return major * 10 + minor >= 80
+    major, _ = tvm.contrib.nvcc.parse_compute_version(arch)
+    return major >= 8
 
 
 def run_test(
@@ -187,11 +187,11 @@ def tile_wmma_fragment(block_read, height, width):
     sch.tensorize(sch.get_loops(block_init_c)[-2], mma_fill_intrin)
     sch.tensorize(sch.get_loops(C_warp)[-2], mma_store_intrin)
 
-    f = tvm.build(sch.mod["main"], target="cuda", name="dense")
-
     if not is_ampere_or_newer():
         return None
 
+    f = tvm.build(sch.mod["main"], target="cuda", name="dense")
+
     dev = tvm.device("cuda", 0)
 
     if in_dtype == "float16":