mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Jiawei Lin
af30ac2152
* [BugFix] Fix loop-dependent conditions in IfThenElse within T.Pipelined This commit applies the same strategy used for LetStmt to IfThenElse conditions: 1.Introduced IfWrapper struct to track if conditions that depend on the loop variable 2.Added dependency detection that checks if an if condition uses: - The pipeline loop variable directly, OR - Any variable transitively dependent on the loop variable 3.Loop-dependent conditions are pushed inside each pipeline stage with the loop variable properly substituted for that iteration * Add test for loop-dependent condtions within T.Pipelined * Fix code format
58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def _make_kernel(M, N):
|
|
dtype = T.bfloat16
|
|
|
|
@T.prim_func
|
|
def fwd_main(KV: T.Tensor((M, N), dtype), ids: T.Tensor((4,), T.int32)):
|
|
with T.Kernel(4, threads=1):
|
|
A = T.alloc_shared([N], dtype)
|
|
B = T.alloc_shared([N], dtype)
|
|
|
|
# Regression for a bug where InjectSoftwarePipeline left the loop
|
|
# variable as a free var, causing MakePackedAPI to fail
|
|
for i in T.Pipelined(4, num_stages=1):
|
|
_id = ids[i]
|
|
T.copy(KV[_id, :], A)
|
|
T.clear(B)
|
|
|
|
return fwd_main
|
|
|
|
|
|
def _make_kernel_if_cond(M, N):
|
|
dtype = T.bfloat16
|
|
|
|
@T.prim_func
|
|
def fwd_main(KV: T.Tensor((M, N), dtype), ids: T.Tensor((4,), T.int32)):
|
|
with T.Kernel(4, threads=1):
|
|
A = T.alloc_shared([N], dtype)
|
|
B = T.alloc_shared([N], dtype)
|
|
|
|
# Regression for a bug where InjectSoftwarePipeline left the loop
|
|
# variable as a free var, causing MakePackedAPI to fail
|
|
for i in T.Pipelined(4, num_stages=1):
|
|
if i > 1:
|
|
_id = ids[i]
|
|
T.copy(KV[_id, :], A)
|
|
T.clear(B)
|
|
|
|
return fwd_main
|
|
|
|
|
|
def test_make_packed_api_no_free_loop_var():
|
|
func, func_if_cond = _make_kernel(4, 4), _make_kernel_if_cond(4, 4)
|
|
# Keep warp-specialization/TMA disabled to match the original repro
|
|
cfg = {
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
}
|
|
tilelang.compile(func, pass_configs=cfg)
|
|
tilelang.compile(func_if_cond, pass_configs=cfg)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|