mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
andyluo7
dd9c20dbf7
A large number of tests import CUDA-only internals directly (for example
`tilelang.cuda.pipeline.CUDAPassPipelineBodyPrologue`) or exercise
NVIDIA-specific features such as TMA, LDGSTG, mbarrier, warp specialization,
tcgen05 and tfloat32. `src/cuda/CMakeLists.txt` only compiles a small
"ALWAYS" subset when `USE_CUDA=OFF`, so those FFI symbols are never
registered and the tests fail with:
AttributeError: module 'tilelang.cuda._ffi_api' has no attribute
'AnnotateDeviceBoundTmaCopies'
This is not ROCm-specific -- the same tests fail on any build without the
CUDA backend, including Metal and CPU-only.
Add `@tilelang.testing.requires_cuda` to those tests, plus three
portability fixes that keep coverage on non-CUDA backends instead of
skipping it:
* `test_tilelang_transform_layout_inference.py`: CUDA emits
`__launch_bounds__(N, 1)` while HIP emits `__launch_bounds__(N)`.
Match both.
* `test_tilelang_transform_pipeline_planning.py`: PipelinePlanning
deliberately strips pipeline annotations on ROCm targets other than
gfx950 (see the guard in src/transform/pipeline_planning.cc), so assert
the documented sequential fallback there.
* `test_tilelang_kernel_gemm.py`: the two f32f32f32 cases use `T.tfloat32`,
which is NVIDIA-only, so narrow `requires_cuda_or_cdna` to
`requires_cuda`.
Also refresh three stale RDNA tests: gfx12 is now a supported generation
(`get_arch` accepts 11 and 12), so they can no longer assert that gfx1200
is rejected. Force an unsupported generation instead, which is the only
way to reach the guard now that `target_is_rdna` reports only gfx11/gfx12.
Verified on 8x MI300X (gfx942, ROCm 7.2 container, torch nightly rocm7.2):
the ROCm suite goes from 105 failures to 0 with no loss of passing tests.
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import tilelang.language as T
|
|
from tilelang.engine.lower import lower
|
|
from tilelang.jit.adapter.utils import match_declare_kernel
|
|
import tilelang.testing
|
|
|
|
|
|
def _simple_add_kernel():
|
|
@T.prim_func
|
|
def main(
|
|
x: T.Tensor((128,), T.float32),
|
|
y: T.Tensor((128,), T.float32),
|
|
):
|
|
# One-dimensional kernel; writes y from x without modifying x
|
|
with T.Kernel(128, threads=32) as pid:
|
|
y[pid] = x[pid] + 1.0
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_codegen_emits_const_for_readonly_params():
|
|
# Lower without device compilation to retrieve CUDA source reliably
|
|
func = _simple_add_kernel()
|
|
artifact = lower(func, target="cuda", enable_device_compile=False)
|
|
|
|
src = artifact.kernel_source
|
|
print(src)
|
|
assert 'extern "C" __global__' in src
|
|
|
|
# Extract kernel signature and check qualifiers
|
|
lparen = match_declare_kernel(src)
|
|
rparen = src.find(")", lparen)
|
|
assert rparen != -1
|
|
signature = src[lparen:rparen]
|
|
|
|
# x is read-only: should be `const` and `__restrict__`
|
|
assert "const float* __restrict__" in signature
|
|
# y is written: must not be const, but still `__restrict__` due to noalias
|
|
# We ensure there is a non-const float* parameter with __restrict__ as well
|
|
assert "const float* __restrict__ x" in src or "const float *__restrict__ x" in src
|
|
assert " float* __restrict__ y" in src or " float *__restrict__ y" in src
|
|
|
|
# Also validate the function attribute carries read-only param indices
|
|
# Expect only the first handle parameter (x) to be marked read-only
|
|
device_mod = artifact.device_mod
|
|
prim_funcs = [f for f in device_mod.functions.values() if hasattr(f, "attrs")]
|
|
assert prim_funcs, "No PrimFunc found in device module"
|
|
pf = prim_funcs[0]
|
|
ro = pf.attrs.get("tl.readonly_param_indices")
|
|
assert ro is not None, "Expected tl.readonly_param_indices to be present"
|
|
ro_list = [int(i) for i in ro]
|
|
assert 0 in ro_list and 1 not in ro_list
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_codegen_emits_const_for_readonly_params()
|