mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Lei Wang
85218bd973
* [Enhancement] Introduce LetWrapper for handling loop variable substitutions in pipeline rewriting * Added LetWrapper struct to encapsulate variable and value pairs for loop variable substitutions. * Updated PipelineRewriter to accept a vector of LetWrapper instances, allowing for proper handling of Let statements that depend on the pipeline loop variable. * Enhanced the BuildPipeline method to incorporate LetWrapper instances into rewritten blocks, ensuring correct substitutions during pipeline execution. * Refactored logic for processing Let statements to differentiate between those that use the loop variable and those that do not, improving the flexibility of the pipeline transformation. * Refactor lambda expression for clarity in loop variable usage check in inject_pipeline.cc * [Test] Add regression test for loop variable handling in kernel compilation * Introduced a new test case to verify correct handling of loop variables in the kernel compilation process, addressing a regression issue with InjectSoftwarePipeline. * The test ensures that the loop variable is not left as a free variable, which previously caused failures in MakePackedAPI. * Configurations are set to disable warp specialization and TMA lowering to align with the original issue reproduction. * Remove unused import in regression test for loop variable handling in kernel compilation
37 lines
1.0 KiB
Python
37 lines
1.0 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def _make_kernel(M, N):
|
|
dtype = "bfloat16"
|
|
|
|
@T.prim_func
|
|
def fwd_main(KV: T.Tensor((M, N), dtype), ids: T.Tensor((4,), "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 test_make_packed_api_no_free_loop_var():
|
|
func = _make_kernel(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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|