mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Yu ChengandLeiWang1999
8c45417584
* [Feature] Implement Hoist Loop-Invariant If Pass - Added a new pass to hoist loop-invariant if statements out of loops, improving optimization opportunities. - Introduced classes for collecting and checking conditions, as well as for rewriting statements. - Integrated the new pass into the optimization pipeline and provided a corresponding API for usage. * Add tilelang semantics guide to programming guides (#1740) * Add CallNodeChecker to Loop Unswitching Pass and integrate tilelang cache disabling in benchmark script - Introduced CallNodeChecker class to identify CallNode expressions in loop conditions, enhancing loop-invariant checks. - Updated IsLoopInvariant function to reject conditions containing CallNodes, preventing potential side effects. - Added tilelang.disable_cache() in benchmark_mha_sink_fwd.py to optimize performance during benchmarking. * lint fix * lint fix * Enhance Loop Unswitching Pass with Let Binding Support - Added support for Let-bound variables in the WrittenBufferReadChecker to improve buffer read checks. - Introduced UsesLoopVarThroughLetBindings function to check if conditions depend on loop variables through Let bindings. - Updated IsLoopInvariant function to account for Let bindings when determining loop invariance. - Enhanced HoistableIfFinder to track Let bindings for variables bound to BufferLoad expressions. - Added debug print statements in the OptimizeForTarget function to visualize the module state before and after loop unswitching. * fix * lint fix * Enhance Loop Unswitching with CallNode Exclusion Check - Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching. - Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness. - Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching. - Disabled tilelang cache in the benchmark script for improved performance. * Update TVM subproject to latest commit and remove debug print statements in OptimizeForTarget function --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T
|
|
|
|
|
|
def test_issue_1734():
|
|
"""Test that loop-invariant if statements are hoisted out of loops."""
|
|
|
|
@tilelang.jit()
|
|
def kernel():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(2, 512), T.float32],
|
|
B: T.Tensor[(2, 512), T.float32],
|
|
C: T.Tensor[(2,), T.float32],
|
|
):
|
|
with T.Kernel(1, threads=256):
|
|
A_local = T.alloc_fragment((2, 512), T.float32)
|
|
B_local = T.alloc_fragment((2, 512), T.float32)
|
|
C_local = T.alloc_fragment((2,), T.float32)
|
|
|
|
T.copy(A, A_local)
|
|
T.copy(C, C_local)
|
|
|
|
for i, j in T.Parallel(2, 512):
|
|
if C_local[i] >= 0:
|
|
B_local[i, j] = A_local[i, j]
|
|
|
|
T.copy(B_local, B)
|
|
|
|
return main
|
|
|
|
mod = kernel.compile()
|
|
source = mod.get_kernel_source()
|
|
# Verify that the if statement is hoisted outside the for loop
|
|
# After hoisting, we should see "if" before "for" pattern
|
|
if_pos = source.find("if (")
|
|
for_pos = source.find("for (")
|
|
assert if_pos < for_pos, "Loop-invariant if should be hoisted outside the loop"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|