mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
from tilelang import tvm
|
|
import tilelang as tl
|
|
import tilelang.testing
|
|
from tvm.script import tir as T
|
|
|
|
|
|
@T.prim_func
|
|
def negative_index_before(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
|
|
T.func_attr({"tir.noalias": True})
|
|
B[0] = A[T.int32(-1)]
|
|
|
|
|
|
@T.prim_func
|
|
def negative_index_expected(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
|
|
T.func_attr({"tir.noalias": True})
|
|
B[0] = A[T.int32(15)]
|
|
|
|
|
|
@T.prim_func
|
|
def negative_index_loop_before(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32")):
|
|
T.func_attr({"tir.noalias": True})
|
|
for i in T.serial(4):
|
|
B[i] = A[-i - 1]
|
|
|
|
|
|
@T.prim_func
|
|
def negative_index_loop_expected(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32")):
|
|
T.func_attr({"tir.noalias": True})
|
|
for i in T.serial(4):
|
|
B[i] = A[15 - i]
|
|
|
|
|
|
@T.prim_func
|
|
def negative_index_symbolic_before(shift: T.int32, A: T.Buffer((16,), "float32"),
|
|
B: T.Buffer((16,), "float32")):
|
|
T.func_attr({"tir.noalias": True})
|
|
for i in T.serial(16):
|
|
B[i] = A[shift + i]
|
|
|
|
|
|
def test_legalize_negative_index_scalar():
|
|
mod = tvm.IRModule({"main": negative_index_before})
|
|
transformed = tl.transform.LegalizeNegativeIndex()(mod)
|
|
tvm.ir.assert_structural_equal(transformed["main"].body, negative_index_expected.body)
|
|
|
|
|
|
def test_legalize_negative_index_affine_expr():
|
|
mod = tvm.IRModule({"main": negative_index_loop_before})
|
|
transformed = tl.transform.LegalizeNegativeIndex()(mod)
|
|
tvm.ir.assert_structural_equal(transformed["main"].body, negative_index_loop_expected.body)
|
|
|
|
|
|
def test_legalize_negative_index_symbolic_passthrough():
|
|
mod = tvm.IRModule({"main": negative_index_symbolic_before})
|
|
transformed = tl.transform.LegalizeNegativeIndex()(mod)
|
|
tvm.ir.assert_structural_equal(transformed["main"].body, negative_index_symbolic_before.body)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|