diff --git a/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py b/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py index e3b6b81d2e..1469cb1dbf 100644 --- a/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py +++ b/examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py @@ -43,7 +43,7 @@ def _tir_u8_to_f4_to_bf16(nbit: int, val: tirx.PrimExpr, pos: tirx.PrimExpr, sca def _get_target(): """Return TVM hip target for gfx950.""" - return tvm.target.Target("hip -mcpu=gfx950") + return tvm.target.Target({"kind": "hip", "mcpu": "gfx950"}) def get_configs(): @@ -59,7 +59,7 @@ def get_configs(): @tilelang.autotune(configs=get_configs()) -@tilelang.jit(out_idx=[-1], target="hip -mcpu=gfx950") +@tilelang.jit(out_idx=[-1], target={"kind": "hip", "mcpu": "gfx950"}) def matmul( M, N, diff --git a/src/backend/common/op/scan.h b/src/backend/common/op/scan.h index f61a456504..4bc8fcc41e 100644 --- a/src/backend/common/op/scan.h +++ b/src/backend/common/op/scan.h @@ -52,11 +52,11 @@ Stmt LowerSharedScan(const ScanOpNode &op, const LowerArgs &T, ICHECK_EQ(op.dim, 0) << scan_noun << " over a 1D buffer only supports dim = 0."; ss << "tl::" << symbol_prefix << "1D<" << threads << ", " - << (op.reverse ? "true" : "false") << ">::run"; + << (op.reverse ? "true" : "false") << ">::run_auto"; args = {StringImm(ss.str()), src_ptr, dst_ptr, src_extents[0]}; } else if (ndim == 2) { ss << "tl::" << symbol_prefix << "2D<" << threads << ", " << op.dim - << ", " << (op.reverse ? "true" : "false") << ">::run"; + << ", " << (op.reverse ? "true" : "false") << ">::run_auto"; args = {StringImm(ss.str()), src_ptr, dst_ptr, src_extents[0], src_extents[1]}; } else { diff --git a/src/tl_templates/cuda/scan.h b/src/tl_templates/cuda/scan.h index 8059233d64..1aa3ca0337 100644 --- a/src/tl_templates/cuda/scan.h +++ b/src/tl_templates/cuda/scan.h @@ -108,6 +108,11 @@ struct InclusiveScan1D { return; InclusiveScanLine(src, dst, N, 1); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int N) { + run(src, dst, N); + } }; template @@ -143,6 +148,11 @@ struct InclusiveScan2D { } } } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int H, int W) { + run(src, dst, H, W); + } }; template struct CumSum1D { @@ -152,6 +162,11 @@ template struct CumSum1D { InclusiveScan1D::template run(src, dst, N); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int N) { + InclusiveScan1D::run_auto(src, dst, N); + } }; template struct CumSum2D { @@ -161,6 +176,12 @@ template struct CumSum2D { InclusiveScan2D::template run( src, dst, H, W); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int H, int W) { + InclusiveScan2D::run_auto(src, dst, H, + W); + } }; template struct CumMax1D { @@ -170,6 +191,11 @@ template struct CumMax1D { InclusiveScan1D::template run(src, dst, N); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int N) { + InclusiveScan1D::run_auto(src, dst, N); + } }; template struct CumMax2D { @@ -179,6 +205,12 @@ template struct CumMax2D { InclusiveScan2D::template run( src, dst, H, W); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int H, int W) { + InclusiveScan2D::run_auto(src, dst, H, + W); + } }; } // namespace tl diff --git a/src/tl_templates/hip/scan.h b/src/tl_templates/hip/scan.h index eb7b6ac77a..24cf65ce52 100644 --- a/src/tl_templates/hip/scan.h +++ b/src/tl_templates/hip/scan.h @@ -161,6 +161,11 @@ template struct CumSum1D { InclusiveScan1D::template run(src, dst, N); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int N) { + InclusiveScan1D::run_auto(src, dst, N); + } }; template struct CumSum2D { @@ -170,6 +175,12 @@ template struct CumSum2D { InclusiveScan2D::template run( src, dst, H, W); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int H, int W) { + InclusiveScan2D::run_auto(src, dst, H, + W); + } }; template struct CumMax1D { @@ -179,6 +190,11 @@ template struct CumMax1D { InclusiveScan1D::template run(src, dst, N); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int N) { + InclusiveScan1D::run_auto(src, dst, N); + } }; template struct CumMax2D { @@ -188,6 +204,12 @@ template struct CumMax2D { InclusiveScan2D::template run( src, dst, H, W); } + template + static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst, + int H, int W) { + InclusiveScan2D::run_auto(src, dst, H, + W); + } }; } // namespace tl diff --git a/testing/python/amd/test_tilelang_mxfp4_gfx950.py b/testing/python/amd/test_tilelang_mxfp4_gfx950.py index 9c92ea1307..265d985ce4 100644 --- a/testing/python/amd/test_tilelang_mxfp4_gfx950.py +++ b/testing/python/amd/test_tilelang_mxfp4_gfx950.py @@ -24,7 +24,7 @@ def _hip_target(): - return "hip -mcpu=gfx950" + return {"kind": "hip", "mcpu": "gfx950"} def _fp4_encode(vals): @@ -220,7 +220,7 @@ def test_get_mxfp_intrin_group_returns_hip_source(): from tilelang.quantize import get_mxfp_intrin_group from tilelang import tvm - target = tvm.target.Target("hip -mcpu=gfx950") + target = tvm.target.Target({"kind": "hip", "mcpu": "gfx950"}) info = get_mxfp_intrin_group( out_dtype=T.bfloat16, source_bit=4, diff --git a/testing/python/components/test_storage_rewrite_detect_inplace.py b/testing/python/components/test_storage_rewrite_detect_inplace.py index 91a6559aad..b70d180ce0 100644 --- a/testing/python/components/test_storage_rewrite_detect_inplace.py +++ b/testing/python/components/test_storage_rewrite_detect_inplace.py @@ -1,6 +1,7 @@ import tilelang import tilelang.testing from tilelang import language as T, tvm +from tilelang.testing import requires_cuda @tilelang.jit @@ -51,6 +52,7 @@ def _get_device_kernel_script(detect_inplace: bool) -> str: return artifact.kernel_source +@requires_cuda def test_storage_rewrite_detect_inplace_toggle(): script_off = _get_device_kernel_script(detect_inplace=False) script_on = _get_device_kernel_script(detect_inplace=True) diff --git a/testing/python/issue/test_tilelang_issue_2123.py b/testing/python/issue/test_tilelang_issue_2123.py index 25950f41de..2b1bf118a0 100644 --- a/testing/python/issue/test_tilelang_issue_2123.py +++ b/testing/python/issue/test_tilelang_issue_2123.py @@ -1,12 +1,20 @@ +import pytest import tilelang import tilelang.testing import tilelang.language as T from tilelang import tvm from tvm import tirx from tvm.tirx import op -from tilelang.cuda.pipeline import CUDAPassPipelineBodyPrologue from tilelang.transform import LowerAccessPtr +try: + from tilelang.cuda.pipeline import CUDAPassPipelineBodyPrologue + import tilelang.cuda._ffi_api as _cuda_ffi_api + + _has_cuda_transforms = hasattr(_cuda_ffi_api, "LowerBlackwell2SM") +except Exception: + _has_cuda_transforms = False + def issue_2123_atomic_load_repro(num_tiles, threads=32): @T.prim_func @@ -60,6 +68,7 @@ def test_issue_2123_atomic_load_lower_access_ptr_direct(): _assert_access_ptr_lowered(lowered) +@pytest.mark.skipif(not _has_cuda_transforms, reason="CUDA transforms not available in this build (USE_CUDA=OFF)") def test_issue_2123_atomic_load_lower_access_ptr_pipeline(): target = tvm.target.Target("cuda", host="llvm") func = issue_2123_atomic_load_repro(4).with_attr("global_symbol", "main") diff --git a/testing/python/jit/test_tilelang_jit_diagnostics.py b/testing/python/jit/test_tilelang_jit_diagnostics.py index 4311f38528..01b360e228 100644 --- a/testing/python/jit/test_tilelang_jit_diagnostics.py +++ b/testing/python/jit/test_tilelang_jit_diagnostics.py @@ -3,6 +3,7 @@ import tilelang.testing import pytest +from tilelang.testing import requires_cuda @pytest.fixture @@ -125,6 +126,7 @@ def test_nvcc_compile_cuda_honors_tilelang_timeout(monkeypatch): assert "Source:" in message +@requires_cuda def test_jit_compile_reports_timeout_for_hanging_nvcc(monkeypatch, tmp_path, caplog, capture_tilelang_logs): import tilelang from tilelang.contrib import nvcc @@ -225,6 +227,7 @@ def __init__(self, *args, **kwargs): assert any("cache_diag_kernel" in message and "sm_120" in message for message in messages) +@requires_cuda def test_explicit_cuda_arch_source_generation_probe(): import tilelang from tilelang import tvm diff --git a/testing/python/kernel/test_tilelang_kernel_gemm.py b/testing/python/kernel/test_tilelang_kernel_gemm.py index aa7149e634..b5f386c4b6 100644 --- a/testing/python/kernel/test_tilelang_kernel_gemm.py +++ b/testing/python/kernel/test_tilelang_kernel_gemm.py @@ -153,7 +153,7 @@ def test_gemm_bf16bf16f32_nn(): ) -@tilelang.testing.requires_cuda_or_cdna +@tilelang.testing.requires_cuda def test_gemm_f32f32f32_nn(): run_gemm( 512, @@ -221,7 +221,7 @@ def test_gemm_f64f64f64_nt(): run_gemm(512, 512, 512, False, True, T.float64, T.float64, T.float64, 64, 32, 16) -@tilelang.testing.requires_cuda_or_cdna +@tilelang.testing.requires_cuda def test_gemm_f32f32f32_nt(): run_gemm( 512, diff --git a/testing/python/language/test_tilelang_language_eager_jit.py b/testing/python/language/test_tilelang_language_eager_jit.py index 5cd89a59b2..3bfba4f96e 100644 --- a/testing/python/language/test_tilelang_language_eager_jit.py +++ b/testing/python/language/test_tilelang_language_eager_jit.py @@ -73,7 +73,10 @@ def gemm_ptr( T.gemm(A_shared, B_shared, C_local) T.copy(C_local, C[bx * block_M, by * block_N]) - prod = product([T.float16, T.tfloat32], [T.float32]) + from tilelang.utils.target import determine_target, target_is_cuda + + in_dtypes = [T.float16, T.tfloat32] if target_is_cuda(determine_target()) else [T.float16] + prod = product(in_dtypes, [T.float32]) gemm_ptr.par_compile( [ {"A": T.ptr(), "B": T.ptr(), "C": T.ptr(), "M": 1024, "N": 1024, "K": 1024, "dtype": in_dtype, "out_dtype": out_dtype} diff --git a/testing/python/language/test_tilelang_language_float4_e2m1_unpacked_dtype.py b/testing/python/language/test_tilelang_language_float4_e2m1_unpacked_dtype.py index f09bebe34b..39b8304b53 100644 --- a/testing/python/language/test_tilelang_language_float4_e2m1_unpacked_dtype.py +++ b/testing/python/language/test_tilelang_language_float4_e2m1_unpacked_dtype.py @@ -75,6 +75,7 @@ def _decode_tcgen05_formats(desc: int) -> tuple[int, int]: return (desc >> 7) & 0x7, (desc >> 10) & 0x7 +@tilelang.testing.requires_cuda def test_tcgen05_instr_desc_fp4_unpacked_x_fp8(): fp4 = DataType("custom[float4_e2m1_unpacked]8") fp8 = DataType("float8_e4m3fn") @@ -85,6 +86,7 @@ def test_tcgen05_instr_desc_fp4_unpacked_x_fp8(): assert b_format == 0 +@tilelang.testing.requires_cuda def test_tcgen05_blockscaled_instr_desc_mxfp4_unpacked_x_mxfp8(): fp4 = DataType("custom[float4_e2m1_unpacked]8") fp8 = DataType("float8_e4m3fn") diff --git a/testing/python/language/test_tilelang_language_tma_copy.py b/testing/python/language/test_tilelang_language_tma_copy.py index a87e61992d..f0d421a5b5 100644 --- a/testing/python/language/test_tilelang_language_tma_copy.py +++ b/testing/python/language/test_tilelang_language_tma_copy.py @@ -397,6 +397,7 @@ def run_fp4_tma_copy_roundtrip(): assert torch.equal(b.view(torch.int8), a) +@tilelang.testing.requires_cuda def test_fp4_unpacksmem_tma_descriptor_uses_align16b(): program = fp4_tma_copy_unpacked_smem_load() artifact = tilelang.lower( @@ -412,6 +413,7 @@ def test_fp4_unpacksmem_tma_descriptor_uses_align16b(): assert 'T.call_packed("__tvm_tensormap_create_tiled", A_desc, 14,' in host_ir +@tilelang.testing.requires_cuda def test_fp4_unpacksmem_tma_store_is_rejected(): program = fp4_tma_copy_unpacked_smem_store() with pytest.raises(tvm.TVMError, match="only supports float4_e2m1_unpacked as an FP4 unpack load"): @@ -422,6 +424,7 @@ def test_fp4_unpacksmem_tma_store_is_rejected(): ) +@tilelang.testing.requires_cuda def test_copy_prefer_tma_lowers_as_synchronous_tma_load(): @T.prim_func def main(x: T.Tensor((128, 32), T.float32)): diff --git a/testing/python/target/test_tilelang_rocm_target.py b/testing/python/target/test_tilelang_rocm_target.py index 65db9f29d9..cac76ace6b 100644 --- a/testing/python/target/test_tilelang_rocm_target.py +++ b/testing/python/target/test_tilelang_rocm_target.py @@ -107,27 +107,34 @@ def fake_rdna(target): assert target_get_mcpu(arch[1]) == "gfx1151" -@tilelang.testing.requires_rocm +@tilelang.testing.requires_rdna def test_carver_rejects_unsupported_rdna_generations(monkeypatch): import tilelang.carver.arch as arch_mod + import tilelang.carver.arch as arch_init_mod + import tilelang.utils.target as target_mod - def fake_cdna(target): - return ("cdna", target) + # Simulate a future RDNA generation that is not yet supported + monkeypatch.setattr(target_mod, "target_is_rdna", lambda t: True) + monkeypatch.setattr(target_mod, "target_get_rdna_generation", lambda t: 10) + monkeypatch.setattr(arch_init_mod, "target_is_rdna", lambda t: True) + monkeypatch.setattr(arch_init_mod, "target_get_rdna_generation", lambda t: 10) - monkeypatch.setattr(arch_mod, "CDNA", fake_cdna) - with pytest.raises(ValueError, match="gfx11 targets only"): - arch_mod.get_arch(_hip_target("gfx1200")) + with pytest.raises(ValueError, match="gfx11/gfx12 targets only"): + arch_mod.get_arch(_hip_target("gfx1151")) -@tilelang.testing.requires_rocm -def test_rdna_device_model_rejects_gfx12_before_device_probe(): - from tilelang.carver.arch.rdna import RDNA +@tilelang.testing.requires_rdna +def test_rdna_device_model_rejects_unsupported_generation_before_device_probe(monkeypatch): + import tilelang.carver.arch.rdna as rdna_mod - with pytest.raises(ValueError, match="gfx11 targets only"): - RDNA(_hip_target("gfx1200")) + # Patch the function as it is bound inside the rdna module to simulate + # a future unsupported generation. ValueError must fire before rocm(0). + monkeypatch.setattr(rdna_mod, "target_get_rdna_generation", lambda t: 10) + with pytest.raises(ValueError, match="gfx11/gfx12 targets only"): + rdna_mod.RDNA(_hip_target("gfx1151")) -@tilelang.testing.requires_rocm +@tilelang.testing.requires_rdna def test_rdna_tensor_instruction_lookup_is_generation_aware(): from tilelang.carver.arch.rdna import RDNA @@ -136,7 +143,12 @@ def test_rdna_tensor_instruction_lookup_is_generation_aware(): assert arch.get_avaliable_tensorintrin_shapes() == [[16, 16]] assert isinstance(arch.available_tensor_instructions, list) + # gfx12 is also supported arch.rdna_generation = 12 + assert arch.get_avaliable_tensorintrin_shapes() == [[16, 16]] + + # Unknown generation raises + arch.rdna_generation = 99 with pytest.raises(ValueError, match="Unsupported RDNA generation"): arch.get_avaliable_tensorintrin_shapes() diff --git a/tilelang/autotuner/tuner.py b/tilelang/autotuner/tuner.py index c5dccf4029..5cb52caa5f 100644 --- a/tilelang/autotuner/tuner.py +++ b/tilelang/autotuner/tuner.py @@ -915,6 +915,11 @@ def run( key = self.generate_cache_key(parameters, extra_parameters) + # Validate scalar input requirements before checking cache so that + # cache hits do not bypass the error when set_autotune_inputs is absent. + if hasattr(self, "_prim_func_for_validation"): + self._validate_input_supply_requirements(self._prim_func_for_validation, self.compile_args.out_idx) + with self._lock: if env.is_cache_enabled() and not env.is_autotune_cache_disabled(): # First check in-memory cache @@ -996,11 +1001,6 @@ def check_tunable_argument_value(key, parameters, key_args_tuple) -> bool: self._memory_cache[key] = autotuner_result return autotuner_result - # After confirming tuning will actually run, validate that scalar - # inputs can be supplied (either via supply_prog or set_autotune_inputs). - if hasattr(self, "_prim_func_for_validation"): - self._validate_input_supply_requirements(self._prim_func_for_validation, self.compile_args.out_idx) - # Launch compile tasks pool, futures, future_to_unit, compile_desc = self._prepare_compile_execution( config_args=config_args, @@ -1280,33 +1280,45 @@ def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> JITKernel | _T: norm_kwargs = _normalize_value(kwargs, sort_dict_items=True) key = (norm_args, norm_kwargs) if key not in self._tuner_cache: + # If all config keys are already supplied by the caller, skip autotuning + # and compile directly. This avoids triggering scalar-input validation + # (which requires set_autotune_inputs) when the user just wants a fixed config. + configs = self.configs + config_keys = set(configs[0].keys()) if isinstance(configs, list) and configs else set() + if config_keys and config_keys.issubset(kwargs.keys()): + if mode == "lazy": + best_kernel = self.jit_impl(*args, **kwargs) + else: + best_kernel = self.jit_impl.compile(*args, **kwargs) + self._tuner_cache[key] = best_kernel, None + else: - def jit_elaborate(**config_arg): - merged = dict(kwargs) - merged.update(config_arg) - return self.jit_impl.get_tir(*args, **merged) + def jit_elaborate(**config_arg): + merged = dict(kwargs) + merged.update(config_arg) + return self.jit_impl.get_tir(*args, **merged) - if mode == "lazy": + if mode == "lazy": - def jit_compile(**config_arg): - return self.jit_impl(*args, **kwargs, __tune_params=config_arg) + def jit_compile(**config_arg): + return self.jit_impl(*args, **kwargs, __tune_params=config_arg) - autotuner.jit_compile = jit_compile - autotuner.jit_elaborate = jit_elaborate - autotuner.set_kernel_parameters(key, self.jit_impl.signature.parameters) - else: + autotuner.jit_compile = jit_compile + autotuner.jit_elaborate = jit_elaborate + autotuner.set_kernel_parameters(key, self.jit_impl.signature.parameters) + else: - def jit_compile(**config_arg): - merged = dict(kwargs) - merged.update(config_arg) - return self.jit_impl.compile(*args, **merged) + def jit_compile(**config_arg): + merged = dict(kwargs) + merged.update(config_arg) + return self.jit_impl.compile(*args, **merged) - autotuner.jit_compile = jit_compile - autotuner.jit_elaborate = jit_elaborate - autotuner.set_kernel_parameters(key, self.jit_impl.signature.parameters) + autotuner.jit_compile = jit_compile + autotuner.jit_elaborate = jit_elaborate + autotuner.set_kernel_parameters(key, self.jit_impl.signature.parameters) - artifact = autotuner.run() - self._tuner_cache[key] = artifact.kernel, artifact.config + artifact = autotuner.run() + self._tuner_cache[key] = artifact.kernel, artifact.config best_kernel, best_config = self._tuner_cache[key] diff --git a/tilelang/contrib/cutedsl/gemm_tcgen05.py b/tilelang/contrib/cutedsl/gemm_tcgen05.py index 0ce54ff160..f9a0b65da1 100644 --- a/tilelang/contrib/cutedsl/gemm_tcgen05.py +++ b/tilelang/contrib/cutedsl/gemm_tcgen05.py @@ -22,6 +22,12 @@ "tcgen05_ld_32dp64bNx", "tcgen05_ld_32dp128bNx", "tcgen05_ld_32dp256bNx", + "tcgen05_st_32dp32bNx", + "tcgen05_st_32dp64bNx", + "tcgen05_st_32dp128bNx", + "tcgen05_st_32dp256bNx", + "tcgen05_before_thread_sync", + "tcgen05_after_thread_sync", ] import cutlass @@ -507,7 +513,7 @@ def _emit_tmem_ld(n_x, max_log_n, ptx_type, regs_per_x, src_addr, dst_view, dst_ def _emit_tmem_fence(): - """Emit tcgen05.wait fence. Called during @cute.jit compilation.""" + """Emit tcgen05.wait::ld fence. Called during @cute.jit compilation.""" llvm.inline_asm( None, [], @@ -519,6 +525,77 @@ def _emit_tmem_fence(): ) +def _emit_tmem_st_fence(): + """Emit tcgen05.wait::st fence. Called during @cute.jit compilation.""" + llvm.inline_asm( + None, + [], + "tcgen05.wait::st.sync.aligned;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +def _emit_tmem_st_segment(ptx_type, seg_x, regs_per_x, addr_ir, src_vals): + """Emit one tcgen05.st inline asm for a power-of-2 segment. + + Called during @cute.jit compilation — emits MLIR ops directly. + ptx_type: PTX instruction type, e.g. "32x32b", "16x64b", "16x128b", "16x256b" + seg_x: x-count in the PTX instruction (power of 2) + regs_per_x: number of i32 input registers per x-element + addr_ir: MLIR IR value for TMEM destination address + src_vals: list of cutlass.Int32 values to store + """ + total_regs = seg_x * regs_per_x + in_regs = ", ".join(f"${i + 1}" for i in range(total_regs)) + asm_str = f"tcgen05.st.sync.aligned.{ptx_type}.x{seg_x}.b32 [$0], {{{in_regs}}};" + constraints = "r," + ",".join(["r"] * total_regs) + operands = [addr_ir] + [v.ir_value() for v in src_vals] + llvm.inline_asm( + None, + operands, + asm_str, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +def _emit_tmem_st(n_x, max_log_n, ptx_type, regs_per_x, dst_addr, src_view, src_offset=0, dst_col_offset=0): + """Recursively split x-count into power-of-2 segments and emit TMEM stores. + + Called during @cute.jit compilation. + n_x: remaining x-element count to store + max_log_n: max log2 of x-count per PTX instruction + ptx_type: PTX instruction type, e.g. "32x32b", "16x64b" + regs_per_x: i32 input registers per x-element + dst_addr: CuTeDSL Int32 (runtime TMEM base address) + src_view: CuTeDSL tensor view over source registers + src_offset: Python int — i32 offset into src_view (compile-time constant) + dst_col_offset: Python int — TMEM column offset from dst_addr (compile-time constant) + """ + if n_x <= 0: + return + + log_n = n_x.bit_length() - 1 + seg_log = min(log_n, max_log_n) + seg_x = 1 << seg_log + total_regs = seg_x * regs_per_x + + if dst_col_offset == 0: + addr_ir = dst_addr.ir_value() + else: + addr_ir = (dst_addr + cutlass.Int32(dst_col_offset)).ir_value() + + src_vals = [cutlass.Int32(src_view[src_offset + j]) for j in range(total_regs)] + _emit_tmem_st_segment(ptx_type, seg_x, regs_per_x, addr_ir, src_vals) + + _emit_tmem_st(n_x - seg_x, max_log_n, ptx_type, regs_per_x, dst_addr, src_view, src_offset + total_regs, dst_col_offset + seg_x) + + @cute.jit def tcgen05_ld_32dp32bNx(N: Constexpr[int], pack16: Constexpr[bool], tmem_start_col: int, tmem_col_offset: int, dst_ptr: cute.Pointer): """Load N uint32 values from TMEM using tcgen05.ld.sync.aligned.32x32b. @@ -592,3 +669,129 @@ def tcgen05_ld_32dp256bNx(N: Constexpr[int], pack16: Constexpr[bool], tmem_start upper_addr = src_addr + cutlass.Int32(16 << 16) _emit_tmem_ld(N, min(_TMEM_LD_MAX_LOG_N, 6), "16x256b", 4, upper_addr, dst_view, dst_offset=regs_per_half, src_col_offset=0) _emit_tmem_fence() + + +# ────────────────────────────────────────────────────────────────────── +# TMEM store — tcgen05.st.sync.aligned.32x32b.xN.b32 +# Mirrors the ld functions above but writes from register memory to TMEM. +# ────────────────────────────────────────────────────────────────────── + +_TMEM_ST_MAX_LOG_N = 3 # cap at x8 to stay within LLVM inline-asm operand limits + + +@cute.jit +def tcgen05_st_32dp32bNx(N: Constexpr[int], unpack16: Constexpr[bool], tmem_start_col: int, tmem_col_offset: int, src_ptr: cute.Pointer): + """Store N uint32 values to TMEM using tcgen05.st.sync.aligned.32x32b. + + Matches tl::tcgen05_st_32dp32bNx from copy_sm100.h. + N: number of 32-bit elements to store (x-count, compile-time constant). + unpack16: if True, use 16-bit unpacking variant. + tmem_start_col: TMEM base column address. + tmem_col_offset: additional column offset. + src_ptr: source pointer (register memory). + """ + dst_addr = cutlass.Int32(tmem_start_col) + cutlass.Int32(tmem_col_offset) + src_view = cute.make_tensor(cute.recast_ptr(src_ptr, dtype=cute.Int32), (N,)) + _emit_tmem_st(N, _TMEM_ST_MAX_LOG_N, "32x32b", 1, dst_addr, src_view) + _emit_tmem_st_fence() + + +@cute.jit +def tcgen05_st_32dp64bNx(N: Constexpr[int], unpack16: Constexpr[bool], tmem_start_col: int, tmem_col_offset: int, src_ptr: cute.Pointer): + """Store to TMEM using 32dp64b pattern (2x 16x64b for lower/upper 16 rows). + + Matches tl::tcgen05_st_32dp64bNx from copy_sm100.h. + N: x-count for 16x64b instructions. Total input: 2*N i32 regs. + """ + total_regs = N * 2 + dst_addr = cutlass.Int32(tmem_start_col) + cutlass.Int32(tmem_col_offset) + src_view = cute.make_tensor(cute.recast_ptr(src_ptr, dtype=cute.Int32), (total_regs,)) + # Lower 16 rows + _emit_tmem_st(N, _TMEM_ST_MAX_LOG_N, "16x64b", 1, dst_addr, src_view, src_offset=0, dst_col_offset=0) + # Upper 16 rows (TMEM row offset = 16 << 16) + upper_addr = dst_addr + cutlass.Int32(16 << 16) + _emit_tmem_st(N, _TMEM_ST_MAX_LOG_N, "16x64b", 1, upper_addr, src_view, src_offset=N, dst_col_offset=0) + _emit_tmem_st_fence() + + +@cute.jit +def tcgen05_st_32dp128bNx(N: Constexpr[int], unpack16: Constexpr[bool], tmem_start_col: int, tmem_col_offset: int, src_ptr: cute.Pointer): + """Store to TMEM using 32dp128b pattern (2x 16x128b for lower/upper 16 rows). + + Matches tl::tcgen05_st_32dp128bNx from copy_sm100.h. + N: x-count for 16x128b instructions. Total input: 4*N i32 regs. + """ + regs_per_half = N * 2 + total_regs = regs_per_half * 2 + dst_addr = cutlass.Int32(tmem_start_col) + cutlass.Int32(tmem_col_offset) + src_view = cute.make_tensor(cute.recast_ptr(src_ptr, dtype=cute.Int32), (total_regs,)) + # Lower 16 rows + _emit_tmem_st(N, min(_TMEM_ST_MAX_LOG_N, 6), "16x128b", 2, dst_addr, src_view, src_offset=0, dst_col_offset=0) + # Upper 16 rows (TMEM row offset = 16 << 16) + upper_addr = dst_addr + cutlass.Int32(16 << 16) + _emit_tmem_st(N, min(_TMEM_ST_MAX_LOG_N, 6), "16x128b", 2, upper_addr, src_view, src_offset=regs_per_half, dst_col_offset=0) + _emit_tmem_st_fence() + + +@cute.jit +def tcgen05_st_32dp256bNx(N: Constexpr[int], unpack16: Constexpr[bool], tmem_start_col: int, tmem_col_offset: int, src_ptr: cute.Pointer): + """Store to TMEM using 32dp256b pattern (2x 16x256b for lower/upper 16 rows). + + Matches tl::tcgen05_st_32dp256bNx from copy_sm100.h. + N: x-count for 16x256b instructions. Total input: 8*N i32 regs. + """ + regs_per_half = N * 4 + total_regs = regs_per_half * 2 + dst_addr = cutlass.Int32(tmem_start_col) + cutlass.Int32(tmem_col_offset) + src_view = cute.make_tensor(cute.recast_ptr(src_ptr, dtype=cute.Int32), (total_regs,)) + # Lower 16 rows + _emit_tmem_st(N, min(_TMEM_ST_MAX_LOG_N, 6), "16x256b", 4, dst_addr, src_view, src_offset=0, dst_col_offset=0) + # Upper 16 rows (TMEM row offset = 16 << 16) + upper_addr = dst_addr + cutlass.Int32(16 << 16) + _emit_tmem_st(N, min(_TMEM_ST_MAX_LOG_N, 6), "16x256b", 4, upper_addr, src_view, src_offset=regs_per_half, dst_col_offset=0) + _emit_tmem_st_fence() + + +# ────────────────────────────────────────────────────────────────────── +# TMEM thread-sync fences — tcgen05.fence::before/after_thread_sync +# Required around __syncthreads() when accessing TMEM to preserve +# memory ordering between TMEM ops and shared-memory ops. +# ────────────────────────────────────────────────────────────────────── + + +@cute.jit +def tcgen05_before_thread_sync(): + """Emit tcgen05.fence::before_thread_sync PTX instruction. + + Must be called before __syncthreads() when TMEM operations precede a + thread synchronization barrier. Matches tl::tcgen05_before_thread_sync + from tcgen_05.h. + """ + llvm.inline_asm( + None, + [], + "tcgen05.fence::before_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def tcgen05_after_thread_sync(): + """Emit tcgen05.fence::after_thread_sync PTX instruction. + + Must be called after __syncthreads() when TMEM operations follow a + thread synchronization barrier. Matches tl::tcgen05_after_thread_sync + from tcgen_05.h. + """ + llvm.inline_asm( + None, + [], + "tcgen05.fence::after_thread_sync;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) diff --git a/tilelang/contrib/cutedsl/math.py b/tilelang/contrib/cutedsl/math.py index 5a77b5a40a..b4b9ad529e 100644 --- a/tilelang/contrib/cutedsl/math.py +++ b/tilelang/contrib/cutedsl/math.py @@ -21,6 +21,7 @@ "__habs", "__float2half_rz", "tanh", + "pow_of_int", ] import cutlass @@ -237,3 +238,23 @@ def __tanhf(x: Union[float, Float32], *, fastmath, loc=None, ip=None) -> Float32 def tanh(x: Union[TensorSSA, Numeric], fastmath: bool = False) -> Union[TensorSSA, Numeric]: tanh_op = __tanhf if fastmath else math.tanh return _tl_math_op(tanh_op, False, x) + + +def pow_of_int(exp: int): + """Return a function that raises its argument to the integer power `exp`. + + Mirrors tl::pow_of_int from math.cc — the C++ codegen emits + tl::pow_of_int as a call_extern, which the CuTeDSL codegen translates + to tl.pow_of_int(N)(base). On CUDA/HIP the op is lowered by FLowerIntrinsic + before reaching call_extern; for the CuTeDSL Python backend it reaches here. + """ + + def _pow(base): + if exp == 0: + return type(base)(1) + result = base + for _ in range(exp - 1): + result = result * base + return result + + return _pow diff --git a/tilelang/contrib/cutedsl/reduce.py b/tilelang/contrib/cutedsl/reduce.py index 2338f7c95c..d55ab7ae32 100644 --- a/tilelang/contrib/cutedsl/reduce.py +++ b/tilelang/contrib/cutedsl/reduce.py @@ -235,6 +235,10 @@ def run(self, src: cute.Pointer, dst: cute.Pointer, N): if tidx < N: dst_tensor[tidx] = val + @cute.jit + def run_auto(self, src: cute.Pointer, dst: cute.Pointer, N): + self.run(src, dst, N) + class CumSum2D: """ @@ -323,6 +327,10 @@ def run(self, src: cute.Pointer, dst: cute.Pointer, H, W): if row_in_col < H and col < W: dst_tensor[idx] = val + @cute.jit + def run_auto(self, src: cute.Pointer, dst: cute.Pointer, H, W): + self.run(src, dst, H, W) + class CumMax1D: """ @@ -356,6 +364,10 @@ def run(self, src: cute.Pointer, dst: cute.Pointer, N): if tidx < N: dst_tensor[tidx] = val + @cute.jit + def run_auto(self, src: cute.Pointer, dst: cute.Pointer, N): + self.run(src, dst, N) + class CumMax2D: """ @@ -418,6 +430,10 @@ def run(self, src: cute.Pointer, dst: cute.Pointer, H, W): if row_in_col < H and col < W: dst_tensor[idx] = val + @cute.jit + def run_auto(self, src: cute.Pointer, dst: cute.Pointer, H, W): + self.run(src, dst, H, W) + class NamedBarrier: """Named barrier policy for AllReduce, uses bar.sync instead of __syncthreads. diff --git a/tilelang/cuda/transform/__init__.py b/tilelang/cuda/transform/__init__.py index a2e9f020c5..9407e500ce 100644 --- a/tilelang/cuda/transform/__init__.py +++ b/tilelang/cuda/transform/__init__.py @@ -15,7 +15,9 @@ def ProducerConsumerWarpSpecialized(): fpass : tvm.transform.Pass The result pass """ - return _ffi_api.ProducerConsumerWarpSpecialized() # type: ignore + if hasattr(_ffi_api, "ProducerConsumerWarpSpecialized"): + return _ffi_api.ProducerConsumerWarpSpecialized() # type: ignore + return lambda f: f def ProducerConsumerWarpSpecializedTiled(): @@ -39,7 +41,9 @@ def LowerBlackwell2SM(): fpass : tvm.transform.Pass The result pass """ - return _ffi_api.LowerBlackwell2SM() # type: ignore + if hasattr(_ffi_api, "LowerBlackwell2SM"): + return _ffi_api.LowerBlackwell2SM() # type: ignore + return lambda f: f def LowerHopperIntrin(): diff --git a/tilelang/language/tir/op.py b/tilelang/language/tir/op.py index bd3053dd26..782b8b0b8b 100644 --- a/tilelang/language/tir/op.py +++ b/tilelang/language/tir/op.py @@ -1691,12 +1691,12 @@ def tvm_mfma( return call_intrin( dtype, _tvm_op.Op.get("tl.tvm_mfma"), - shape, - A_layout, - B_layout, - A_dtype, - B_dtype, - C_dtype, + tvm.tirx.StringImm(str(shape)), + tvm.tirx.StringImm(str(A_layout)), + tvm.tirx.StringImm(str(B_layout)), + tvm.tirx.StringImm(str(A_dtype)), + tvm.tirx.StringImm(str(B_dtype)), + tvm.tirx.StringImm(str(C_dtype)), multiplicand_a, a_index, multiplicand_b, diff --git a/tilelang/rocm/intrinsics/mfma_macro_generator.py b/tilelang/rocm/intrinsics/mfma_macro_generator.py index fb7a1a0681..19a4bdc329 100644 --- a/tilelang/rocm/intrinsics/mfma_macro_generator.py +++ b/tilelang/rocm/intrinsics/mfma_macro_generator.py @@ -52,6 +52,7 @@ class MatrixCoreIntrinEmitter: "float16": "fp16", "bfloat16": "bf16", "float32": "fp32", + "custom[tfloat32]": "fp32", "int8": "int8", "int32": "int32", "float8_e4m3": "e4m3", @@ -173,6 +174,7 @@ def _initialize_mfma_prefix(self, k_dim=16): "bfloat16": "bf16", "float16": "f16", "float32": "f32", + "custom[tfloat32]": "f32", "int8": "i8", "int32": "i32", "float8_e4m3fn": "fp8", diff --git a/tilelang/testing/__init__.py b/tilelang/testing/__init__.py index 37637bc3c4..97b7d90d9b 100644 --- a/tilelang/testing/__init__.py +++ b/tilelang/testing/__init__.py @@ -5,9 +5,7 @@ import torch import numpy as np from tilelang.contrib import nvcc -from tilelang.backend.target import determine_target -from tilelang.cuda.target import target_is_cuda -from tilelang.rocm.target import target_is_cdna, target_is_gfx950 +from tilelang.utils.target import determine_target, target_is_cdna, target_is_cuda, target_is_gfx950, target_is_rdna from tvm.testing.utils import requires_cuda, requires_package, requires_llvm, requires_metal, requires_rocm, _compose from tilelang.utils.tensor import torch_assert_close as torch_assert_close @@ -22,6 +20,7 @@ "requires_cdna", "requires_cuda_or_cdna", "requires_gfx950", + "requires_rdna", "main", "requires_cuda_compute_version", "process_func", @@ -91,6 +90,27 @@ def requires_gfx950(func): return _compose([func], marks) +def _check_is_rdna() -> bool: + try: + target = determine_target("auto", return_object=True) + return target_is_rdna(target) + except (ValueError, RuntimeError): + return False + + +def requires_rdna(func): + """Skip the test unless the ROCm device is an RDNA GPU (gfx11/gfx12).""" + is_rdna = _check_is_rdna() + marks = [ + pytest.mark.skipif( + not is_rdna, + reason="Requires RDNA ROCm target (gfx11/gfx12)", + ), + *requires_rocm.marks(), + ] + return _compose([func], marks) + + # pytest.main() wrapper to allow running single test file def main(): test_file = inspect.getsourcefile(sys._getframe(1)) diff --git a/tilelang/utils/target.py b/tilelang/utils/target.py new file mode 100644 index 0000000000..c3dac85da7 --- /dev/null +++ b/tilelang/utils/target.py @@ -0,0 +1,3 @@ +from tilelang.backend.target import determine_target # noqa: F401 +from tilelang.cuda.target import target_is_cuda # noqa: F401 +from tilelang.rocm.target import target_is_cdna, target_is_gfx950, target_is_rdna # noqa: F401