[Tests][TIRx] Localize hardware test gates

tests/python/tirx/conftest.py → tests/python/tirx/codegen/conftest.py

@@ -14,16 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Suite-level hardware gate for the tirx tests.
-
-The tirx kernels and codegen paths target Blackwell (sm_100a) — they emit
-PTX/SASS (tcgen05, tmem, cp.async ``.async`` modifiers, fp8 conversions, ...)
-that ptxas/NVRTC reject for older targets, and many tests execute on the
-device. Running the suite on a CPU-only node or a pre-sm_100 GPU therefore
-fails at compile/run time rather than skipping. Gate the whole directory on a
-real sm_100a device so it skips cleanly where the hardware is absent and runs
-in full where it is present.
-"""
+"""Hardware requirements for TIRx codegen tests."""
 
 from pathlib import Path
 
@@ -32,20 +23,14 @@
 from tvm.testing import env
 
 
-def pytest_collection_modifyitems(config, items):
+def pytest_collection_modifyitems(items):
     if env.has_cuda_compute(10):
         return
     suite_root = Path(__file__).resolve().parent
-    skip = pytest.mark.skip(
-        reason="tirx suite requires a CUDA compute capability 10.0 (sm_100a) device"
-    )
+    skip = pytest.mark.skip(reason="requires a CUDA compute capability 10.0 device")
     for item in items:
-        path = getattr(item, "path", None)
-        if path is None:
-            continue
-        try:
-            path = Path(path).resolve()
-        except TypeError:
-            continue
-        if path.is_relative_to(suite_root):
+        if (
+            Path(item.path).resolve().is_relative_to(suite_root)
+            and item.get_closest_marker("gpu") is not None
+        ):
             item.add_marker(skip)

tests/python/tirx/codegen/test_codegen_cuda.py

@@ -24,10 +24,13 @@
 from tvm.testing import env
 
 
-def _get_source(func: tvm.tirx.PrimFunc) -> str:
-    target = tvm.target.Target("cuda")
+def _get_source(func: tvm.tirx.PrimFunc, target=None) -> tuple[str, tvm.IRModule]:
+    if target is None:
+        target = {"kind": "cuda", "arch": "sm_100a"}
+    target = tvm.target.Target(target)
     mod = tvm.IRModule({"main": func})
-    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    with target:
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
     src = mod.mod.imports[0].inspect_source()
     return src, mod
 
@@ -102,17 +105,18 @@ def main(A: T.Buffer((1,), "int32")):
     assert "cooperative_groups::cluster_group::block_index" not in src
 
 
+@pytest.mark.gpu
 def test_cuda_handle_uint64_reinterpret_codegen():
     @T.prim_func
     def main(A: T.Buffer((1,), "uint64")):
         T.device_entry()
         tx = T.thread_id([32])
         if tx == 0:
-            ptr = T.reinterpret("handle", A[0])
+            ptr: T.let = T.reinterpret("handle", A[0])
             A[0] = T.reinterpret("uint64", ptr)
 
     src, _ = _get_source(main)
-    assert "reinterpret_cast<void*>" in src
+    assert "(void*)A_ptr[0]" in src
     assert "reinterpret_cast<uint64_t>" in src
     assert "*(void* *)" not in src
 
@@ -166,6 +170,7 @@ def main(A: T.Buffer((1,), "uint64"), B: T.Buffer((1,), "int32"), C: T.Buffer((1
     assert "ld.volatile.global.u64" in src
 
 
+@pytest.mark.gpu
 def test_megamoe_extracted_intrinsics_codegen():
     @T.prim_func
     def main(
@@ -388,7 +393,7 @@ def main(Cache: T.Buffer((1,), "uint64")):
                     cache_policy=Cache[0],
                 )
 
-    src, _ = _get_source(main)
+    src, _ = _get_source(main, {"kind": "cuda", "arch": "sm_100a"})
     assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_cache_hint" in src
     assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_multicast_cache_hint" in src
     assert "g2cluster_unicast" not in src

tests/python/tirx/codegen/test_codegen_dsmem.py

@@ -23,9 +23,10 @@
 
 
 def _get_source(func: tvm.tirx.PrimFunc) -> str:
-    target = tvm.target.Target("cuda")
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
     mod = tvm.IRModule({"main": func})
-    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    with target:
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
     src = mod.mod.imports[0].inspect_source()
     return src
 

tests/python/tirx/codegen/test_codegen_hopper.py

@@ -148,6 +148,7 @@ def run_and_check():
 
 @pytest.mark.parametrize("trans", [False, True])
 @pytest.mark.parametrize("num", [1, 2, 4])
+@pytest.mark.gpu
 def test_ptx_stmatrix(trans, num):
     # fmt: off
     @T.prim_func

tests/python/tirx/operator/tile_primitive/cuda/conftest.py

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Hardware requirements for CUDA tile-primitive tests."""
+
+from pathlib import Path
+
+import pytest
+
+from tvm.testing import env
+
+
+def pytest_collection_modifyitems(items):
+    if env.has_cuda_compute(10):
+        return
+    suite_root = Path(__file__).resolve().parent
+    skip = pytest.mark.skip(reason="requires a CUDA compute capability 10.0 device")
+    for item in items:
+        if Path(item.path).resolve().is_relative_to(suite_root):
+            item.add_marker(skip)

tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py

@@ -469,6 +469,7 @@ def run_and_check():
     tvm.testing.run_with_gpu_lock(run_and_check)
 
 
+@pytest.mark.gpu
 def test_ldstmatrix_tcgen05_warpgroup_atom_emits_ldmatrix():
     """Regression: a warpgroup ``.16x256b`` tcgen05 register atom loaded from a
     128B-swizzled SMEM tile must dispatch to ``ldmatrix.x4``.

tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py

@@ -338,6 +338,7 @@ def run_and_check():
         tvm.testing.run_with_gpu_lock(run_and_check)
 
 
+@pytest.mark.gpu
 def test_reg_copy_wg_local_to_swizzled_shared_uses_swizzle_fastpath():
     """Regression: R→S copy where R has a ``wg_local_layout`` (thread iter
     ``1 @ tid_in_wg``) must pick the widest vec PTX ``st.shared.v4`` AND use the

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py

@@ -447,6 +447,7 @@ def test_dispatch_rejects_bad_inputs(bad):
             tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
 
 
+@pytest.mark.gpu
 def test_multi_cp_encodes_descriptor_once_and_patches_addr():
     """Compile-only regression for the shared-descriptor cp path.
 

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py

@@ -1003,6 +1003,7 @@ def _tma_case(
 # fmt: on
 
 
+@pytest.mark.gpu
 @pytest.mark.parametrize("case", TMA_CASES)
 def test_copy_tma_codegen(case):
     """Unified structural-golden driver for every TMA unit test case.
@@ -1576,6 +1577,7 @@ def copy_async_dynamic_mask(A_ptr: T.handle) -> None:
     assert "multicast" in src, "Expected multicast TMA instruction in generated code"
 
 
+@pytest.mark.gpu
 def test_copy_tma_uint32_shape_extent():
     BK = 64
     A_layout = mma_shared_layout("float16", 3, (128, BK))
@@ -1609,6 +1611,7 @@ def tma_load(n: T.uint32, a_ptr: T.handle, o_ptr: T.handle) -> None:
         tvm.compile(tvm.IRModule({"main": tma_load}), target=target, tir_pipeline="tirx")
 
 
+@pytest.mark.gpu
 def test_copy_tma_uint32_slice_base():
     BK = 64
     A_layout = mma_shared_layout("float16", 3, (128, BK))

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py

@@ -826,6 +826,7 @@ def run_and_check():
         ("16x256b", 64, 64),  # .16x256b.x8 fp32
     ],
 )
+@pytest.mark.gpu
 def test_alloc_tcgen05_frag_wrapper_compiles(shape, frag_rows, K_cols):
     """Ensure T.alloc_tcgen05_ldst_frag yields a buffer that ``T.copy_async`` accepts
     and lowers to the correct tcgen05 atom for each supported instr_shape."""

tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py

@@ -369,6 +369,7 @@ def test_cuda_gemm_mma_variant_is_registered():
 
 
 @pytest.mark.parametrize("dtype", ["bfloat16", "float16"])
+@pytest.mark.gpu
 def test_cuda_gemm_mma_lowers_to_mma_sync(dtype):
     """beta=0: the dispatch clears D, then issues a single accumulating mma with
     the registers laid out in the fixed PTX fragment order."""
@@ -389,6 +390,7 @@ def test_cuda_gemm_mma_lowers_to_mma_sync(dtype):
         assert f"b_local[{r}]" in script
 
 
+@pytest.mark.gpu
 def test_cuda_gemm_mma_accumulates_c_when_beta_one():
     """beta=1: the accumulator is initialized by copying C instead of zeroing."""
     script = _lower(_build_gemm(alpha=1.0, beta=1.0))["main"].script()
@@ -599,6 +601,7 @@ def run_and_check():
         (2, 2, 3, 8),  # k8, every dim tiled
     ],
 )
+@pytest.mark.gpu
 def test_cuda_gemm_mma_lowers_tiled(Mt, Nt, Kt, kinst):
     """Every tiling we expect to dispatch must lower, selecting the right mma.
 
@@ -644,6 +647,7 @@ def test_cuda_gemm_mma_codegen_issue_count(Mt, Nt, Kt, kinst):
     "transpose_A, transpose_B",
     [(False, False), (True, False), (False, True), (True, True)],
 )
+@pytest.mark.gpu
 def test_cuda_gemm_mma_lowers_transpose(transpose_A, transpose_B):
     """All four A/B orientations dispatch to the same m16n8k16. transpose only
     describes the input's logical orientation; the .row.col mma is unchanged."""

tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py

@@ -2285,6 +2285,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
 
 
 @pytest.mark.parametrize("smem_desc", ["hoist", "recompute"])
+@pytest.mark.gpu
 def test_gemm_smem_desc_hoist_vs_recompute(smem_desc):
     """Compile-only: the SMEM matrix descriptor is built per-MMA from the buffer
     base address, selected by the ``smem_desc`` config.

tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py

@@ -421,6 +421,7 @@ def f(A: T.handle, B: T.handle):
 
 
 @pytest.mark.parametrize("dtype", ["uint32", "float32"])
+@pytest.mark.gpu
 def test_shared_to_shared_uses_direct_ldst(dtype):
     """Compile-only: a shared->shared 32b transpose must take the direct
     base-ptr + byte-offset ``ld.shared`` / ``st.shared`` path.

tests/python/tirx/test_parser_printer.py

@@ -1230,7 +1230,7 @@ def test_annotation_syntax_comprehensive():
     def test_let_var():
         T.device_entry()
         smem = T.alloc_shared([128], "float16")
-        ptr: T.let[T.Var(name="ptr", dtype=PointerType(PrimType("uint64")))] = T.reinterpret(
+        ptr: T.let[T.Var(name="ptr", dtype=PointerType(PrimType("void")))] = T.reinterpret(
             "handle", smem.access_ptr("rw")
         )
         T.evaluate(ptr)

tests/python/tirx/transform/test_transform_lower_tirx.py

@@ -583,7 +583,7 @@ def before():
         T.thread_id([128])
         buf = T.alloc_buffer([1024], "uint8", scope="shared.dyn")
         A = T.decl_buffer([128], "float16", buf.data, elem_offset=32)
-        T.evaluate(A.access_ptr("rw", offset=A.elem_offset_of([64])))
+        T.evaluate(A.access_ptr("rw", ptr_type="float16", offset=A.elem_offset_of([64])))
 
     @T.prim_func(private=True)
     def after():