mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
b939fa0137
* [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 * Add support for disabling loop unswitching optimization - Introduced a new configuration option `tl.disable_loop_unswitching` to allow users to disable the loop unswitching optimization. - Updated the Loop Unswitching pass to check this configuration and return the original function if the option is enabled. - Added relevant documentation in the PassConfigKey enumeration for clarity. * update tvm * Refactor: Update TileLang to use tirx instead of tir - Removed references to `tvm.tir` and replaced them with `tvm.tirx` across various files, including examples and backend operations. - Updated target configuration in documentation to reflect the new usage of target config dictionaries instead of CLI-style strings. - Cleaned up `pyproject.toml` and `load_tvm.cmake` by removing obsolete paths. - Enhanced examples for dequantization and GEMM operations to utilize the new `tirx` constructs. - Adjusted various backend operations to ensure compatibility with the new `tirx` namespace. * pass fix * Update submodule 'tvm' to latest commit f24c7b19 * Apply formatting fixes * Update CMakeLists.txt to set HIDE_PRIVATE_SYMBOLS for TileLang and modify target_host initialization in lower.py for compatibility with TVM's target handling. * Refactor CMakeLists.txt and kernel_cache.py to update target names and library dependencies for TVM integration. Adjusted output targets and library names for better compatibility across platforms. * Update repair-wheel-command in pyproject.toml to exclude libz3.so with wildcard for better compatibility with dependency packages. * Fix macOS Z3 dylib lookup * Move macOS Z3 install-name fix to install step * Update TVM submodule for runtime import fix * Update ROCm backend for TVM module factory changes * Fix build-tree Z3 runtime lookup * Fix macOS build-tree Z3 lookup * Support legacy target option strings * Normalize CUDA target in transform test * Qualify Metal stream FFI lookup * Fix PyPI Z3 packaging links * Update TVM for Windows DLL lookup * Register Windows dependency DLL directories * Bump tvm-ffi to match tirx * tvm update * Apply lint fixes * Enhance reduction operations with support for nan propagation and vectorization. Introduce new functions for preferred vectorized size and update existing reduction logic to handle packed operations for bfloat16 and float16 types. Add nan-aware min and max operations in CUDA and ROCm backends, and update related tests to validate functionality. * Refactor ReduceLowerer: Remove unnecessary buffer allocation logic for packed operations in reduce.h * Normalize PyPI Z3 dylib before linking * Drop legacy target option string support * Fix type reference for barrier annotations in CUDA copy operations. Updated `BufferLoadNode` to `tirx::BufferLoadNode` in `GetBarrier` and `LowerCluster` methods to ensure compatibility with recent changes in the TIR API. --------- Co-authored-by: Yu Cheng <yu.cheng@pku.edu.cn> Co-authored-by: Yichen Yan <oraluben@outlook.com>
179 lines
8.1 KiB
Python
179 lines
8.1 KiB
Python
"""Tests for Layout and Fragment equality comparison."""
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
from tilelang.layout import Layout
|
|
from tilelang.layout.fragment import Fragment
|
|
|
|
tilelang.testing.set_random_seed()
|
|
|
|
|
|
class TestLayoutEqual:
|
|
"""Test cases for Layout.is_equal()."""
|
|
|
|
def test_same_layout_is_equal(self):
|
|
"""Two layouts with identical mapping should be equal."""
|
|
layout1 = Layout([32, 4], lambda i, j: i * 4 + j)
|
|
layout2 = Layout([32, 4], lambda i, j: i * 4 + j)
|
|
assert layout1.is_equal(layout2)
|
|
|
|
def test_different_index_order_not_equal(self):
|
|
"""Layouts with different index order (i*4+j vs j*4+i) should not be equal."""
|
|
layout1 = Layout([32, 4], lambda i, j: i * 4 + j)
|
|
layout2 = Layout([32, 4], lambda i, j: j * 4 + i)
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
def test_different_coefficient_not_equal(self):
|
|
"""Layouts with different coefficients should not be equal."""
|
|
layout1 = Layout([32, 4], lambda i, j: i * 4 + j)
|
|
layout2 = Layout([32, 4], lambda i, j: i * 8 + j)
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
def test_different_shape_not_equal(self):
|
|
"""Layouts with different shapes should not be equal."""
|
|
layout1 = Layout([32, 4], lambda i, j: i * 4 + j)
|
|
layout2 = Layout([16, 8], lambda i, j: i * 8 + j)
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
def test_same_layout_different_var_names(self):
|
|
"""Layouts with same mapping but created with different variable names should be equal."""
|
|
layout1 = Layout([32, 4], lambda x, y: x * 4 + y)
|
|
layout2 = Layout([32, 4], lambda a, b: a * 4 + b)
|
|
assert layout1.is_equal(layout2)
|
|
|
|
def test_2d_output_layout_equal(self):
|
|
"""Layouts with 2D output should compare correctly."""
|
|
layout1 = Layout([32, 4], lambda i, j: [i, j])
|
|
layout2 = Layout([32, 4], lambda i, j: [i, j])
|
|
assert layout1.is_equal(layout2)
|
|
|
|
def test_2d_output_layout_different_order(self):
|
|
"""Layouts with swapped output dimensions should not be equal."""
|
|
layout1 = Layout([32, 4], lambda i, j: [i, j])
|
|
layout2 = Layout([32, 4], lambda i, j: [j, i])
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
def test_complex_expression_equal(self):
|
|
"""Layouts with complex but equivalent expressions should be equal."""
|
|
layout1 = Layout([16, 8], lambda i, j: i * 8 + j)
|
|
layout2 = Layout([16, 8], lambda i, j: j + i * 8)
|
|
# Note: This tests if the comparison handles commutative operations
|
|
# With StructuralEqual, a*b+c and c+a*b have different AST structure
|
|
# So this may or may not be equal depending on implementation
|
|
# For now we test the actual behavior
|
|
result = layout1.is_equal(layout2)
|
|
# The key point is it should not crash and return a boolean
|
|
assert isinstance(result, bool)
|
|
|
|
|
|
class TestFragmentEqual:
|
|
"""Test cases for Fragment.is_equal()."""
|
|
|
|
def test_same_fragment_is_equal(self):
|
|
"""Two fragments with identical mapping should be equal."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j)
|
|
assert frag1.is_equal(frag2)
|
|
|
|
def test_different_thread_mapping_not_equal(self):
|
|
"""Fragments with different thread mapping (i*4+j vs j*4+i) should not be equal."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda i, j: j * 4 + i)
|
|
assert not frag1.is_equal(frag2)
|
|
|
|
def test_different_forward_index_not_equal(self):
|
|
"""Fragments with different forward_index should not be equal."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j, forward_index_fn=lambda i, j: i)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j, forward_index_fn=lambda i, j: j)
|
|
assert not frag1.is_equal(frag2)
|
|
|
|
def test_same_fragment_different_var_names(self):
|
|
"""Fragments with same mapping but different variable names should be equal."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda x, y: x * 4 + y)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda a, b: a * 4 + b)
|
|
assert frag1.is_equal(frag2)
|
|
|
|
def test_fragment_with_replicate_equal(self):
|
|
"""Fragments with same replicate factor should be equal."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda i, j, rep: i * 4 + rep, replicate=4)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda i, j, rep: i * 4 + rep, replicate=4)
|
|
assert frag1.is_equal(frag2)
|
|
|
|
def test_fragment_different_replicate_not_equal(self):
|
|
"""Fragments with different replicate factors should not be equal."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda i, j, rep: i * 4 + rep, replicate=4)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda i, j, rep: i * 4 + rep, replicate=2)
|
|
assert not frag1.is_equal(frag2)
|
|
|
|
def test_fragment_with_forward_fn(self):
|
|
"""Fragments created with forward_fn should compare correctly."""
|
|
frag1 = Fragment([32, 4], forward_fn=lambda i, j: (i * 4 + j, i * 4 + j))
|
|
frag2 = Fragment([32, 4], forward_fn=lambda i, j: (i * 4 + j, i * 4 + j))
|
|
assert frag1.is_equal(frag2)
|
|
|
|
def test_fragment_forward_fn_different_thread(self):
|
|
"""Fragments with different thread mapping via forward_fn should not be equal."""
|
|
frag1 = Fragment([32, 4], forward_fn=lambda i, j: (i * 4 + j, i))
|
|
frag2 = Fragment([32, 4], forward_fn=lambda i, j: (j * 4 + i, i))
|
|
assert not frag1.is_equal(frag2)
|
|
|
|
def test_fragment_forward_fn_different_index(self):
|
|
"""Fragments with different forward_index via forward_fn should not be equal."""
|
|
frag1 = Fragment([32, 4], forward_fn=lambda i, j: (i * 4 + j, i))
|
|
frag2 = Fragment([32, 4], forward_fn=lambda i, j: (i * 4 + j, j))
|
|
assert not frag1.is_equal(frag2)
|
|
|
|
|
|
class TestLayoutFragmentEdgeCases:
|
|
"""Edge cases and regression tests."""
|
|
|
|
def test_single_dim_layout_equal(self):
|
|
"""Single dimension layouts should compare correctly."""
|
|
layout1 = Layout([128], lambda i: i)
|
|
layout2 = Layout([128], lambda i: i)
|
|
assert layout1.is_equal(layout2)
|
|
|
|
def test_single_dim_layout_not_equal(self):
|
|
"""Single dimension layouts with different mappings should not be equal."""
|
|
layout1 = Layout([128], lambda i: i)
|
|
layout2 = Layout([128], lambda i: i * 2)
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
def test_three_dim_layout_equal(self):
|
|
"""Three dimension layouts should compare correctly."""
|
|
layout1 = Layout([8, 16, 4], lambda i, j, k: i * 64 + j * 4 + k)
|
|
layout2 = Layout([8, 16, 4], lambda i, j, k: i * 64 + j * 4 + k)
|
|
assert layout1.is_equal(layout2)
|
|
|
|
def test_three_dim_layout_different_order(self):
|
|
"""Three dimension layouts with different index order should not be equal."""
|
|
layout1 = Layout([8, 16, 4], lambda i, j, k: i * 64 + j * 4 + k)
|
|
layout2 = Layout([8, 16, 4], lambda i, j, k: k * 64 + j * 4 + i)
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
def test_fragment_empty_forward_index(self):
|
|
"""Fragments with empty forward_index should compare correctly."""
|
|
frag1 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j)
|
|
frag2 = Fragment([32, 4], forward_thread_fn=lambda i, j: i * 4 + j)
|
|
assert frag1.is_equal(frag2)
|
|
|
|
def test_constant_layout_equal(self):
|
|
"""Layouts mapping to constant should be equal."""
|
|
from tvm.tirx import const
|
|
|
|
layout1 = Layout([32, 4], lambda i, j: const(0, "int32"))
|
|
layout2 = Layout([32, 4], lambda i, j: const(0, "int32"))
|
|
assert layout1.is_equal(layout2)
|
|
|
|
def test_constant_vs_variable_layout_not_equal(self):
|
|
"""Layout mapping to constant vs variable should not be equal."""
|
|
from tvm.tirx import const
|
|
|
|
layout1 = Layout([32, 4], lambda i, j: const(0, "int32"))
|
|
layout2 = Layout([32, 4], lambda i, j: i)
|
|
assert not layout1.is_equal(layout2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|