From 401a820ad1006427bb07664e3d45aec5ff8036f1 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 19 May 2022 08:14:52 +0900 Subject: [PATCH] skip cuda compile on old gpu --- .../unittest/test_tir_schedule_tensorize_ldmatrix_mma.py | 8 ++++---- 1 file changed, 4 insertions(+), 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 index 78c615ff06c34..67e8ae0ad8367 100644 --- 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":