From 9ba68e1038bf6ed52e572deee40af2e943f3165d Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Sat, 11 Jul 2026 02:53:42 -0400 Subject: [PATCH] [CI][TIRx] Localize hardware test gates --- tests/python/tirx/{ => codegen}/conftest.py | 29 ++++------------ .../python/tirx/codegen/test_codegen_cuda.py | 17 ++++++---- .../python/tirx/codegen/test_codegen_dsmem.py | 5 +-- .../tirx/codegen/test_codegen_hopper.py | 1 + .../operator/tile_primitive/cuda/conftest.py | 33 +++++++++++++++++++ .../cuda/copy/test_ld_stmatrix.py | 1 + .../tile_primitive/cuda/copy/test_reg.py | 1 + .../cuda/copy_async/test_smem_tmem.py | 1 + .../cuda/copy_async/test_tma.py | 3 ++ .../cuda/copy_async/test_tmem_16xnb.py | 1 + .../cuda/gemm/test_gemm_mma_m16n8k_.py | 4 +++ .../cuda/gemm_async/test_gemm_async.py | 1 + .../permute_layout/test_permute_layout.py | 1 + tests/python/tirx/test_parser_printer.py | 2 +- .../transform/test_transform_lower_tirx.py | 2 +- 15 files changed, 70 insertions(+), 32 deletions(-) rename tests/python/tirx/{ => codegen}/conftest.py (51%) create mode 100644 tests/python/tirx/operator/tile_primitive/cuda/conftest.py diff --git a/tests/python/tirx/conftest.py b/tests/python/tirx/codegen/conftest.py similarity index 51% rename from tests/python/tirx/conftest.py rename to tests/python/tirx/codegen/conftest.py index 2653a1f05ae1..69e14a36c6ef 100644 --- a/tests/python/tirx/conftest.py +++ b/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) diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py index 53e46eb23aa3..6358867ed2a1 100644 --- a/tests/python/tirx/codegen/test_codegen_cuda.py +++ b/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" in src + assert "(void*)A_ptr[0]" in src assert "reinterpret_cast" 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 diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py b/tests/python/tirx/codegen/test_codegen_dsmem.py index d538be571f88..052034dd0d20 100644 --- a/tests/python/tirx/codegen/test_codegen_dsmem.py +++ b/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 diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py index 0d17fd0f20a0..59b93ecf2a32 100644 --- a/tests/python/tirx/codegen/test_codegen_hopper.py +++ b/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 diff --git a/tests/python/tirx/operator/tile_primitive/cuda/conftest.py b/tests/python/tirx/operator/tile_primitive/cuda/conftest.py new file mode 100644 index 000000000000..f333ccb7bfd0 --- /dev/null +++ b/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) diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py index f19b7e59a9af..d5ae8130a764 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py +++ b/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``. diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py index 752531f450de..a19ca76d1272 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py +++ b/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 diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py index a615032e37b5..e58f4b2423b4 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py +++ b/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. diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index 6b6b1ebcdcf8..7cbce484ba45 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py +++ b/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)) diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py index d6d730e4f92a..ff9f20c6a643 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py +++ b/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.""" diff --git a/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py b/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py index c424e960466f..5f632011960a 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py +++ b/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.""" diff --git a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py index eceaa50be4b4..1cbe4421f96d 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py +++ b/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. diff --git a/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py b/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py index e4cdc629cc48..4e589a7031f0 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py +++ b/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. diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 34fab19a5a3a..ad08690a2a33 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/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) diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py index 037e415fe9f6..272b4598d432 100644 --- a/tests/python/tirx/transform/test_transform_lower_tirx.py +++ b/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():