mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Witherstrike
a4ab99e466
feat: add end-to-end LLVM backend execution support Enable LLVM as a TileLang execution backend in CMake and route llvm targets through TVM LLVM codegen during lowering. Fix tvm_ffi output tensor device inference by deriving the output device from tensor inputs, and add LLVM execution regression tests for matmul, T.gemm, T.copy, and T.While.
31 lines
1018 B
Python
31 lines
1018 B
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.testing.requires_llvm
|
|
def test_llvm_kernel_source_generation_with_while() -> None:
|
|
"""Regression test: LLVM `T.While` kernels should compile without errors.
|
|
|
|
See: https://github.com/tile-ai/tilelang/issues/2202
|
|
|
|
Historically, a CPU kernel containing `T.While(...)` inside
|
|
`T.Kernel(...)` could leak the synthetic fallback thread variable
|
|
`v_thread` into host/device splitting. LLVM uses the same scalar lowering
|
|
path for this case, so this keeps equivalent coverage for target="llvm".
|
|
"""
|
|
|
|
@T.prim_func
|
|
def main(flag: T.Tensor((1,), "int32"), out: T.Tensor((1,), "int32")):
|
|
with T.Kernel(1):
|
|
state = T.alloc_fragment((1,), "int32")
|
|
state[0] = 0
|
|
|
|
with T.While(state[0] == 0):
|
|
state[0] = flag[0]
|
|
|
|
out[0] = state[0]
|
|
|
|
compiled = tilelang.compile(main, target="llvm", execution_backend="tvm_ffi")
|
|
assert compiled is not None
|