mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-29 07:53:03 +08:00
Lei Wang
cdc67fc440
* • Enable configurable StorageRewrite inplace detection - Add kStorageRewriteDetectInplace constant and register the flag with PassContext so C++ code no longer hard-codes the key. - Wire StorageRewrite to include TileLang builtin constants and honor the new config toggle when deciding inplace reuse. - Document the flag across Python surfaces (PassConfigKey, JIT/autotuner docs) with usage guidance and simplified IR examples. * lint fix * add test * lint fix
62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def _compile_kernel_without_inplace():
|
|
num_tokens = T.symbolic("num_tokens")
|
|
|
|
@T.prim_func
|
|
def buggy_kernel(x: T.Tensor[(num_tokens,), "float"]):
|
|
with T.Kernel(num_tokens, threads=32) as pid:
|
|
read = T.alloc_var("int")
|
|
read = x[pid]
|
|
|
|
write = T.alloc_var("int")
|
|
write = read * 2
|
|
x[pid] = write
|
|
|
|
return buggy_kernel
|
|
|
|
|
|
@tilelang.jit(
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_STORAGE_REWRITE_DETECT_INPLACE: True,
|
|
},)
|
|
def _compile_kernel_with_inplace():
|
|
num_tokens = T.symbolic("num_tokens")
|
|
|
|
@T.prim_func
|
|
def buggy_kernel(x: T.Tensor[(num_tokens,), "float"]):
|
|
with T.Kernel(num_tokens, threads=32) as pid:
|
|
read = T.alloc_var("int")
|
|
read = x[pid]
|
|
|
|
write = T.alloc_var("int")
|
|
write = read * 2
|
|
x[pid] = write
|
|
|
|
return buggy_kernel
|
|
|
|
|
|
def _get_device_kernel_script(detect_inplace: bool) -> str:
|
|
if detect_inplace:
|
|
kernel = _compile_kernel_with_inplace()
|
|
else:
|
|
kernel = _compile_kernel_without_inplace()
|
|
source = kernel.get_kernel_source()
|
|
return source
|
|
|
|
|
|
def test_storage_rewrite_detect_inplace_toggle():
|
|
script_off = _get_device_kernel_script(detect_inplace=False)
|
|
script_on = _get_device_kernel_script(detect_inplace=True)
|
|
|
|
assert script_off.count("read = (read * 2);") == 0
|
|
assert script_on.count("read = (read * 2);") > 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|