This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fa903cf4f7 [Tests][TIRx] Localize hardware test gates (#19985)
fa903cf4f7 is described below
commit fa903cf4f759b73a776cfb7946ded56c3ae7e425
Author: Shushi Hong <[email protected]>
AuthorDate: Tue Jul 14 17:31:41 2026 -0400
[Tests][TIRx] Localize hardware test gates (#19985)
This PR removes the suite-wide TIRx compute-capability gate and
localizes hardware skips to CUDA codegen and tile-primitive tests. It
keeps the original test parameterization unchanged, allowing parser,
printer, IR, transform, and other non-hardware TIRx tests to run in
regular CI while device-dependent cases are skipped when SM100 hardware
is unavailable. CUDA codegen helpers use explicit target architectures
where needed, and the run-only benchmark utility tests retain a local
SM100 gate.
This intentionally avoids adding separate compile/run parameter cases. A
follow-up PR can audit slow frontend execution tests and define a
focused runtime regression budget for real TIRx kernels.
Local validation:
- `pytest -n auto -m "not gpu" tests/python/tirx`: 486 passed, 79
skipped.
- `pytest -n auto -m gpu tests/python/tirx` without matching hardware:
1509 skipped.
- Pre-commit passed on all changed files.
---
tests/python/tirx/{ => codegen}/conftest.py | 29 ++++++----------------
tests/python/tirx/codegen/test_codegen_cuda.py | 17 ++++++++-----
tests/python/tirx/codegen/test_codegen_dsmem.py | 5 ++--
tests/python/tirx/codegen/test_codegen_hopper.py | 1 +
.../{ => operator/tile_primitive/cuda}/conftest.py | 26 +++----------------
.../tile_primitive/cuda/copy/test_ld_stmatrix.py | 1 +
.../operator/tile_primitive/cuda/copy/test_reg.py | 1 +
.../cuda/copy_async/test_smem_tmem.py | 1 +
.../tile_primitive/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 +
.../cuda/permute_layout/test_permute_layout.py | 1 +
tests/python/tirx/test_parser_printer.py | 2 +-
.../tirx/transform/test_transform_lower_tirx.py | 2 +-
15 files changed, 41 insertions(+), 54 deletions(-)
diff --git a/tests/python/tirx/conftest.py
b/tests/python/tirx/codegen/conftest.py
similarity index 51%
copy from tests/python/tirx/conftest.py
copy to tests/python/tirx/codegen/conftest.py
index 2653a1f05a..69e14a36c6 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 @@ import pytest
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 53e46eb23a..6358867ed2 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.script import tirx as T
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 test_cluster_cta_id_codegen_uses_coordinate_sregs():
assert "cooperative_groups::cluster_group::block_index" not in src
[email protected]
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 test_ptx_ld_acquire_and_volatile_codegen():
assert "ld.volatile.global.u64" in src
[email protected]
def test_megamoe_extracted_intrinsics_codegen():
@T.prim_func
def main(
@@ -388,7 +393,7 @@ def test_tma_cache_policy_operand_codegen():
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 d538be571f..052034dd0d 100644
--- a/tests/python/tirx/codegen/test_codegen_dsmem.py
+++ b/tests/python/tirx/codegen/test_codegen_dsmem.py
@@ -23,9 +23,10 @@ from tvm.script import tirx as T
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 0d17fd0f20..59b93ecf2a 100644
--- a/tests/python/tirx/codegen/test_codegen_hopper.py
+++ b/tests/python/tirx/codegen/test_codegen_hopper.py
@@ -148,6 +148,7 @@ def test_stmatrix_sync_aligned(trans):
@pytest.mark.parametrize("trans", [False, True])
@pytest.mark.parametrize("num", [1, 2, 4])
[email protected]
def test_ptx_stmatrix(trans, num):
# fmt: off
@T.prim_func
diff --git a/tests/python/tirx/conftest.py
b/tests/python/tirx/operator/tile_primitive/cuda/conftest.py
similarity index 51%
rename from tests/python/tirx/conftest.py
rename to tests/python/tirx/operator/tile_primitive/cuda/conftest.py
index 2653a1f05a..f333ccb7bf 100644
--- a/tests/python/tirx/conftest.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/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 CUDA tile-primitive tests."""
from pathlib import Path
@@ -32,20 +23,11 @@ import pytest
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):
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 f19b7e59a9..d5ae8130a7 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 test_ldstmatrix_swizzle_multi_iter_pow2():
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
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 752531f450..a19ca76d12 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 test_copy_g2l_l2g_vec_load(task, dtype):
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
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 a615032e37..e58f4b2423 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")
[email protected]
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 6b6b1ebcdc..7cbce484ba 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 @@ TMA_CASES = [
# fmt: on
[email protected]
@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 test_copy_tma_dynamic_cta_mask(dtype):
assert "multicast" in src, "Expected multicast TMA instruction in
generated code"
[email protected]
def test_copy_tma_uint32_shape_extent():
BK = 64
A_layout = mma_shared_layout("float16", 3, (128, BK))
@@ -1609,6 +1611,7 @@ def test_copy_tma_uint32_shape_extent():
tvm.compile(tvm.IRModule({"main": tma_load}), target=target,
tir_pipeline="tirx")
[email protected]
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 d6d730e4f9..ff9f20c6a6 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 test_tcgen05_st_16xnb_store(shape, rep, dtype):
("16x256b", 64, 64), # .16x256b.x8 fp32
],
)
[email protected]
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 c424e96046..5f63201196 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"])
[email protected]
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
[email protected]
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 test_cuda_gemm_mma_numerical_transpose(transpose_A,
transpose_B, dtype):
(2, 2, 3, 8), # k8, every dim tiled
],
)
[email protected]
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)],
)
[email protected]
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 eceaa50be4..1cbe4421f9 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 _build_smem_desc_kernel(smem_desc):
@pytest.mark.parametrize("smem_desc", ["hoist", "recompute"])
[email protected]
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 e4cdc629cc..4e589a7031 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 test_reject_non_warp_scope():
@pytest.mark.parametrize("dtype", ["uint32", "float32"])
[email protected]
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 34fab19a5a..ad08690a2a 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 037e415fe9..272b4598d4 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 test_lower_decl_buffer_access_ptr():
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():
