Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand All @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/backend/common/op/scan.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
32 changes: 32 additions & 0 deletions src/tl_templates/cuda/scan.h
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,11 @@ struct InclusiveScan1D {
return;
InclusiveScanLine<Reducer, reverse, T, SEG>(src, dst, N, 1);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int N) {
run<T, 32>(src, dst, N);
}
};

template <class Reducer, int threads, int Axis = 0, bool reverse = false>
Expand Down Expand Up @@ -143,6 +148,11 @@ struct InclusiveScan2D {
}
}
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int H, int W) {
run<T, 32>(src, dst, H, W);
}
};

template <int threads, bool reverse = false> struct CumSum1D {
Expand All @@ -152,6 +162,11 @@ template <int threads, bool reverse = false> struct CumSum1D {
InclusiveScan1D<ScanSumOp, threads, reverse>::template run<T, SEG>(src, dst,
N);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int N) {
InclusiveScan1D<ScanSumOp, threads, reverse>::run_auto(src, dst, N);
}
};

template <int threads, int Axis = 0, bool reverse = false> struct CumSum2D {
Expand All @@ -161,6 +176,12 @@ template <int threads, int Axis = 0, bool reverse = false> struct CumSum2D {
InclusiveScan2D<ScanSumOp, threads, Axis, reverse>::template run<T, SEG>(
src, dst, H, W);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int H, int W) {
InclusiveScan2D<ScanSumOp, threads, Axis, reverse>::run_auto(src, dst, H,
W);
}
};

template <int threads, bool reverse = false> struct CumMax1D {
Expand All @@ -170,6 +191,11 @@ template <int threads, bool reverse = false> struct CumMax1D {
InclusiveScan1D<ScanMaxOp, threads, reverse>::template run<T, SEG>(src, dst,
N);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int N) {
InclusiveScan1D<ScanMaxOp, threads, reverse>::run_auto(src, dst, N);
}
};

template <int threads, int Axis = 0, bool reverse = false> struct CumMax2D {
Expand All @@ -179,6 +205,12 @@ template <int threads, int Axis = 0, bool reverse = false> struct CumMax2D {
InclusiveScan2D<ScanMaxOp, threads, Axis, reverse>::template run<T, SEG>(
src, dst, H, W);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int H, int W) {
InclusiveScan2D<ScanMaxOp, threads, Axis, reverse>::run_auto(src, dst, H,
W);
}
};

} // namespace tl
22 changes: 22 additions & 0 deletions src/tl_templates/hip/scan.h
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,11 @@ template <int threads, bool reverse = false> struct CumSum1D {
InclusiveScan1D<ScanSumOp, threads, reverse>::template run<T, SEG>(src, dst,
N);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int N) {
InclusiveScan1D<ScanSumOp, threads, reverse>::run_auto(src, dst, N);
}
};

template <int threads, int Axis = 0, bool reverse = false> struct CumSum2D {
Expand All @@ -170,6 +175,12 @@ template <int threads, int Axis = 0, bool reverse = false> struct CumSum2D {
InclusiveScan2D<ScanSumOp, threads, Axis, reverse>::template run<T, SEG>(
src, dst, H, W);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int H, int W) {
InclusiveScan2D<ScanSumOp, threads, Axis, reverse>::run_auto(src, dst, H,
W);
}
};

template <int threads, bool reverse = false> struct CumMax1D {
Expand All @@ -179,6 +190,11 @@ template <int threads, bool reverse = false> struct CumMax1D {
InclusiveScan1D<ScanMaxOp, threads, reverse>::template run<T, SEG>(src, dst,
N);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int N) {
InclusiveScan1D<ScanMaxOp, threads, reverse>::run_auto(src, dst, N);
}
};

template <int threads, int Axis = 0, bool reverse = false> struct CumMax2D {
Expand All @@ -188,6 +204,12 @@ template <int threads, int Axis = 0, bool reverse = false> struct CumMax2D {
InclusiveScan2D<ScanMaxOp, threads, Axis, reverse>::template run<T, SEG>(
src, dst, H, W);
}
template <typename T>
static TL_DEVICE void run_auto(const T *__restrict__ src, T *__restrict__ dst,
int H, int W) {
InclusiveScan2D<ScanMaxOp, threads, Axis, reverse>::run_auto(src, dst, H,
W);
}
};

} // namespace tl
4 changes: 2 additions & 2 deletions testing/python/amd/test_tilelang_mxfp4_gfx950.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@


def _hip_target():
return "hip -mcpu=gfx950"
return {"kind": "hip", "mcpu": "gfx950"}


def _fp4_encode(vals):
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import tilelang
import tilelang.testing
from tilelang import language as T, tvm
from tilelang.testing import requires_cuda


@tilelang.jit
Expand Down Expand Up @@ -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)
Expand Down
11 changes: 10 additions & 1 deletion testing/python/issue/test_tilelang_issue_2123.py
Original file line number Diff line number Diff line change
@@ -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
Comment on lines +10 to +16

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Narrow the exception handling to catch only import-related exceptions.

The bare except Exception: clause catches all exceptions, which could mask unexpected errors during import such as syntax errors or attribute errors in the imported modules. This makes debugging harder if the CUDA modules have actual defects.

🛡️ Proposed fix to narrow exception handling
 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:
+except (ImportError, AttributeError):
     _has_cuda_transforms = False
🧰 Tools
🪛 Ruff (0.15.15)

[warning] 15-15: Do not catch blind exception: Exception

(BLE001)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@testing/python/issue/test_tilelang_issue_2123.py` around lines 10 - 16, The
broad except Exception around importing CUDAPassPipelineBodyPrologue and
tilelang.cuda._ffi_api should be narrowed to only catch import-related errors;
replace the generic except with an except that catches ImportError and
ModuleNotFoundError so that failures inside the imported modules (e.g.,
syntax/attribute errors) surface instead of being masked. Update the try/except
around the CUDAPassPipelineBodyPrologue and _cuda_ffi_api imports that set
_has_cuda_transforms to False on failure to only catch
ImportError/ModuleNotFoundError while leaving other exceptions to propagate.



def issue_2123_atomic_load_repro(num_tiles, threads=32):
@T.prim_func
Expand Down Expand Up @@ -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")
Expand Down
3 changes: 3 additions & 0 deletions testing/python/jit/test_tilelang_jit_diagnostics.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import tilelang.testing
import pytest
from tilelang.testing import requires_cuda


@pytest.fixture
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions testing/python/kernel/test_tilelang_kernel_gemm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 4 additions & 1 deletion testing/python/language/test_tilelang_language_eager_jit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand Down
3 changes: 3 additions & 0 deletions testing/python/language/test_tilelang_language_tma_copy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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"):
Expand All @@ -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)):
Expand Down
36 changes: 24 additions & 12 deletions testing/python/target/test_tilelang_rocm_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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()

Expand Down
Loading
Loading