mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Yuxi Chi
7fc47066cd
* Port CuTeDSL backend to CUTLASS primitives * Update CuTeDSL dependency to CUTLASS DSL 4.7.0 * Fix CuTeDSL dynamic-shape wrapper generation
34 lines
888 B
Python
34 lines
888 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
|
|
|
|
|
|
@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]
|
|
|
|
kernel = tilelang.compile(main)
|
|
assert kernel.get_kernel_source()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|