mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Zheng YuandZheng Yu
bd5ca2f05a
Preserve buffer offsets in access pointers Co-authored-by: Zheng Yu <17133261+cla7aye15I4nd@users.noreply.github.com>
110 lines
4.1 KiB
Python
110 lines
4.1 KiB
Python
import tilelang
|
|
import tvm
|
|
from tvm import arith
|
|
from tvm import tirx
|
|
from tvm.tirx import op
|
|
import tilelang.language as T
|
|
import tilelang.transform
|
|
|
|
|
|
def test_access_ptr_builds_tl_access_ptr_from_bufferload_1d():
|
|
buf = tirx.decl_buffer((64,), "uint8", name="A")
|
|
load = tirx.BufferLoad(buf, [tirx.IntImm("int32", 16)])
|
|
|
|
ptr = T.access_ptr(load, "r", 16)
|
|
|
|
assert isinstance(ptr, tirx.Call)
|
|
assert ptr.op.same_as(op.Op.get("tl.access_ptr"))
|
|
assert len(ptr.args) == 3
|
|
# args: (base_load, extent, rw_mask)
|
|
assert isinstance(ptr.args[0], tirx.BufferLoad)
|
|
assert isinstance(ptr.args[1], tirx.IntImm)
|
|
assert int(ptr.args[1].value) == 16
|
|
assert isinstance(ptr.args[2], tirx.IntImm)
|
|
assert int(ptr.args[2].value) == 1
|
|
|
|
|
|
def test_access_ptr_defaults_to_element_extent_for_bufferload():
|
|
buf = tirx.decl_buffer((64,), "float16", name="A")
|
|
load = tirx.BufferLoad(buf, [tirx.IntImm("int32", 7)])
|
|
|
|
ptr = T.access_ptr(load, "rw")
|
|
|
|
assert isinstance(ptr, tirx.Call)
|
|
assert ptr.op.same_as(op.Op.get("tl.access_ptr"))
|
|
assert isinstance(ptr.args[0], tirx.BufferLoad)
|
|
assert isinstance(ptr.args[1], tirx.IntImm)
|
|
assert int(ptr.args[1].value) == 1
|
|
assert isinstance(ptr.args[2], tirx.IntImm)
|
|
assert int(ptr.args[2].value) == 3
|
|
|
|
|
|
def test_access_ptr_multiplies_extents_for_2d_load():
|
|
buf = tirx.decl_buffer((8, 8), "float16", name="A")
|
|
load = tirx.BufferLoad(buf, [tirx.IntImm("int32", 2), tirx.IntImm("int32", 3)])
|
|
|
|
ptr = T.access_ptr(load, "w", 4, 5)
|
|
|
|
assert isinstance(ptr, tirx.Call)
|
|
assert ptr.op.same_as(op.Op.get("tl.access_ptr"))
|
|
assert isinstance(ptr.args[0], tirx.BufferLoad)
|
|
# extent = 4*5 = 20
|
|
assert isinstance(ptr.args[1], tirx.IntImm)
|
|
assert int(ptr.args[1].value) == 20
|
|
assert isinstance(ptr.args[2], tirx.IntImm)
|
|
assert int(ptr.args[2].value) == 2
|
|
|
|
|
|
def test_lower_access_ptr_rewrites_to_tvm_access_ptr():
|
|
buf = tirx.decl_buffer((8, 8), "float16", name="A")
|
|
load = tirx.BufferLoad(buf, [tirx.IntImm("int32", 2), tirx.IntImm("int32", 3)])
|
|
ptr = T.access_ptr(load, "w", 4, 5)
|
|
|
|
func = tirx.PrimFunc([buf.data], tirx.Evaluate(ptr), buffer_map={buf.data: buf})
|
|
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
|
|
lowered = tilelang.transform.LowerAccessPtr()(mod)
|
|
|
|
calls: list[tirx.Call] = []
|
|
|
|
def _collect(e):
|
|
if isinstance(e, tirx.Call):
|
|
calls.append(e)
|
|
|
|
tirx.stmt_functor.post_order_visit(lowered["main"].body, _collect)
|
|
assert any(c.op.same_as(op.Op.get("tirx.tvm_access_ptr")) for c in calls)
|
|
assert not any(c.op.same_as(op.Op.get("tl.access_ptr")) for c in calls)
|
|
|
|
# Check the lowered tvm_access_ptr carries the expected linear offset/extents.
|
|
acc = [c for c in calls if c.op.same_as(op.Op.get("tirx.tvm_access_ptr"))][0]
|
|
assert len(acc.args) == 5
|
|
analyzer = arith.Analyzer()
|
|
offset = analyzer.simplify(acc.args[2])
|
|
extent = analyzer.simplify(acc.args[3])
|
|
assert isinstance(offset, tirx.IntImm)
|
|
assert int(offset.value) == 19
|
|
assert isinstance(extent, tirx.IntImm)
|
|
assert int(extent.value) == 20
|
|
assert isinstance(acc.args[4], tirx.IntImm)
|
|
assert int(acc.args[4].value) == 2
|
|
|
|
|
|
def test_lower_access_ptr_preserves_buffer_element_offset():
|
|
elem_offset = tirx.IntImm("int32", 5)
|
|
buf = tirx.decl_buffer((8, 8), "float16", name="A", elem_offset=elem_offset)
|
|
load = tirx.BufferLoad(buf, [tirx.IntImm("int32", 2), tirx.IntImm("int32", 3)])
|
|
ptr = T.access_ptr(load, "r", 1)
|
|
|
|
func = tirx.PrimFunc([buf.data], tirx.Evaluate(ptr), buffer_map={buf.data: buf})
|
|
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
|
|
lowered = tilelang.transform.LowerAccessPtr()(mod)
|
|
|
|
calls = []
|
|
tirx.stmt_functor.post_order_visit(
|
|
lowered["main"].body,
|
|
lambda expr: calls.append(expr) if isinstance(expr, tirx.Call) else None,
|
|
)
|
|
acc = [call for call in calls if call.op.same_as(op.Op.get("tirx.tvm_access_ptr"))][0]
|
|
offset = arith.Analyzer().simplify(acc.args[2])
|
|
assert isinstance(offset, tirx.IntImm)
|
|
assert int(offset.value) == 24
|