mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
|
|
|
|
def test_cpu_kernel_source_generation_with_while() -> None:
|
|
"""Regression test: CPU `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(..., is_cpu=True)` could leak the synthetic fallback thread
|
|
variable `v_thread` into host/device splitting. That in turn caused CPU
|
|
compilation to fail during packed-API generation or later C wrapper
|
|
generation.
|
|
|
|
This test follows the existing compile-only regression style used in the
|
|
repository: success is defined by `tilelang.compile(..., target="c")`
|
|
completing without raising an exception.
|
|
"""
|
|
|
|
@T.prim_func
|
|
def main(flag: T.Buffer((1,), "int32"), out: T.Buffer((1,), "int32")):
|
|
with T.Kernel(1, is_cpu=True):
|
|
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="c")
|
|
assert compiled is not None
|