mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tilelang.layout import Layout
|
|
|
|
|
|
def _assert_indices_equal(lhs, rhs):
|
|
assert tvm.arith.Analyzer().can_prove_equal(lhs, rhs), f"{lhs} != {rhs}"
|
|
|
|
|
|
def test_linearized_forward_index_for_multi_dimensional_output():
|
|
layout = Layout([2, 3, 4], lambda i, j, k: [i, j, k])
|
|
linear = Layout([2, 3, 4], lambda i, j, k: i * 12 + j * 4 + k)
|
|
|
|
_assert_indices_equal(
|
|
layout.get_linearized_forward_index(),
|
|
linear.get_linearized_forward_index(),
|
|
)
|
|
|
|
|
|
def test_linearized_forward_index_for_split_output():
|
|
layout = Layout([64], lambda i: [i // 32, i % 32])
|
|
linear = Layout([64], lambda i: i)
|
|
|
|
_assert_indices_equal(
|
|
layout.get_linearized_forward_index(),
|
|
linear.get_linearized_forward_index(),
|
|
)
|
|
|
|
|
|
def test_linearized_forward_index_uses_output_shape():
|
|
transposed = Layout([2, 32], lambda i, j: [j, i])
|
|
expected = Layout([2, 32], lambda i, j: j * 2 + i)
|
|
|
|
_assert_indices_equal(
|
|
transposed.get_linearized_forward_index(),
|
|
expected.get_linearized_forward_index(),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|