mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-28 07:43:01 +08:00
andyluo7
02cf41f178
The test added in #2899 lowers with an explicit `tvm.target.Target("cuda")`, which reaches `tilelang/cuda/transform/__init__.py`. In a `USE_CUDA=OFF` build only the "ALWAYS" CUDA sources compile, so `AnnotateDeviceBoundTmaCopies` is never registered and the test fails on ROCm: AttributeError: module 'tilelang.cuda._ffi_api' has no attribute 'AnnotateDeviceBoundTmaCopies' Same pattern as #2863. `requires_cuda` is the weakest sufficient gate: the test is compile-only (`enable_device_compile=False`), so it needs no particular GPU generation and keeps running on the NVIDIA leg. Verified in CI: SKIPPED on the gfx942 leg, PASSED on CUDA-auto.
38 lines
1014 B
Python
38 lines
1014 B
Python
"""Regression test for GitHub issue #2883.
|
|
|
|
Layout inference must accept a large injective ``T.Parallel`` domain even
|
|
when it exceeds the exact-enumeration limit and requires a stronger proof.
|
|
"""
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_large_parallel_layout_is_injective():
|
|
stride = T.dynamic("stride")
|
|
|
|
@T.prim_func
|
|
def main(
|
|
dst: T.StridedTensor((1, 8192), (stride, 1), "uint8"),
|
|
src: T.StridedTensor((1, 8192), (stride, 1), "uint8"),
|
|
):
|
|
T.assume(stride % 2 == 0)
|
|
|
|
with T.Kernel(1):
|
|
for token, feature in T.Parallel(64, 8192):
|
|
if token < 1:
|
|
dst[token, feature] = src[token, feature]
|
|
|
|
target = tvm.target.Target("cuda")
|
|
with target:
|
|
artifact = tilelang.lower(main, target=target, enable_device_compile=False)
|
|
|
|
assert artifact.kernel_source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|