mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
Lei Wang
4ff81c7d40
* [Refactor] Update atomic operations to use access_ptr for buffer access * Replaced instances of address_of with access_ptr in atomic operations to improve clarity and maintainability. * Enhanced atomic operation implementations in various files, including atomic_add, atomic_max, and atomic_min. * Updated related logic in loop vectorization and tile language to support new access patterns. * Added a utility function to extract destination scalar dtype for atomic operations, supporting both BufferLoad and tvm_access_ptr. * Improved kernel cache logic to include a build-stamp for the TileLang runtime library, ensuring cache validity during development. * lint fix * Fix access_ptr layout rewrite and atomic vectorization * lint fix * Refactor atomic operations to remove const qualifiers from pointer parameters in atomicAddx2 and ToFloat functions. This change enhances type flexibility and consistency across atomic operations. * Refactor buffer access logic in lower_tile_op and vectorize_loop to enhance clarity and maintainability. Removed unused code and improved offset calculations for tile access. Updated atomic operation handling to streamline dtype extraction and ensure correct vectorization behavior. * Add logging for loop layout and lowered loop in atomic_add; implement tvm_access_ptr handling in loop_vectorize and lower_tile_op This commit introduces logging statements to track the loop layout and the result of the lowered loop in the AtomicAddNode. Additionally, it implements a new method to handle tvm_access_ptr calls in the loop_vectorize and lower_tile_op files, enhancing the handling of memory access patterns and ensuring correct vectorization behavior. * fix * lint fix * fix * Update ToFloat2 and ToFloat4 functions to use const qualifiers for pointer parameters in atomic.h; modify test script to call test_atomic_addx2_float instead of main. * Refactor tvm_access_ptr handling in loop_vectorize to improve argument validation and buffer resolution logic. Ensure at least three arguments are provided and streamline the search for matching buffers in layout_map_. Update vector size calculations to enhance clarity and maintainability. * Add handling for tvm_access_ptr in vectorize_loop to support offset substitution and buffer creation This commit introduces logic to handle tvm_access_ptr calls, allowing for the substitution of loop variables in offsets and the creation of dummy buffers. The new method ensures correct base address calculation and enhances the overall handling of memory access patterns in vectorization. * Add handling for tvm_access_ptr offset validation in loop_vectorize This commit enhances the loop_vectorize functionality by adding checks for the element offset in tvm_access_ptr calls. It ensures that offsets are invariant within the vector boundary, improving the robustness of vectorization logic and memory access patterns. * fix * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update improves the robustness of vectorization logic and ensures offsets remain invariant within vector boundaries. * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update ensures offsets remain invariant within vector boundaries, improving memory access robustness and vectorization logic. * fix * Add tl.access_ptr support in TileLang This commit introduces the `tl.access_ptr` operation, which carries pointer access metadata in the frontend. It allows for better handling of buffer loads by retaining necessary information for downstream analysis. The new operation is integrated into various components, including the loop vectorization and memory access legalizations, ensuring compatibility with existing functionality while enhancing the robustness of memory access patterns. * lint fix * fix * Add support for tl.access_ptr in atomic operations This commit updates the AtomicAddNode and AtomicOpBaseNode implementations to utilize the new tl.access_ptr operation for building access pointers to destination elements. It enhances the handling of buffer loads by incorporating type checks for index data types, improving the robustness of memory access patterns. Additionally, it removes unnecessary logging statements in the lowering process, streamlining the codebase. * lint fix * Refactor buffer load extraction in vectorize_loop to check for non-empty arguments in access_ptr calls. This change improves the robustness of buffer load handling by ensuring valid argument presence before processing. * Remove unused fallback logic in KernelCache and clean up import statements in test_tilelang_language_access_ptr.py for improved code clarity. * Refactor memory access in BitNet and dequantization examples to use tl.access_ptr This update replaces T.address_of calls with T.access_ptr in multiple files, enhancing memory access handling by specifying read and write permissions. This change improves code clarity and aligns with recent enhancements in pointer access management.
89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
import tilelang
|
|
import tvm
|
|
from tvm import arith
|
|
from tvm import tir
|
|
from tvm.tir import op
|
|
import tilelang.language as T
|
|
import tilelang.transform
|
|
|
|
|
|
def test_access_ptr_builds_tl_access_ptr_from_bufferload_1d():
|
|
buf = tir.decl_buffer((64,), "uint8", name="A")
|
|
load = tir.BufferLoad(buf, [tir.IntImm("int32", 16)])
|
|
|
|
ptr = T.access_ptr(load, "r", 16)
|
|
|
|
assert isinstance(ptr, tir.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], tir.BufferLoad)
|
|
assert isinstance(ptr.args[1], tir.IntImm)
|
|
assert int(ptr.args[1].value) == 16
|
|
assert isinstance(ptr.args[2], tir.IntImm)
|
|
assert int(ptr.args[2].value) == 1
|
|
|
|
|
|
def test_access_ptr_defaults_to_element_extent_for_bufferload():
|
|
buf = tir.decl_buffer((64,), "float16", name="A")
|
|
load = tir.BufferLoad(buf, [tir.IntImm("int32", 7)])
|
|
|
|
ptr = T.access_ptr(load, "rw")
|
|
|
|
assert isinstance(ptr, tir.Call)
|
|
assert ptr.op.same_as(op.Op.get("tl.access_ptr"))
|
|
assert isinstance(ptr.args[0], tir.BufferLoad)
|
|
assert isinstance(ptr.args[1], tir.IntImm)
|
|
assert int(ptr.args[1].value) == 1
|
|
assert isinstance(ptr.args[2], tir.IntImm)
|
|
assert int(ptr.args[2].value) == 3
|
|
|
|
|
|
def test_access_ptr_multiplies_extents_for_2d_load():
|
|
buf = tir.decl_buffer((8, 8), "float16", name="A")
|
|
load = tir.BufferLoad(buf, [tir.IntImm("int32", 2), tir.IntImm("int32", 3)])
|
|
|
|
ptr = T.access_ptr(load, "w", 4, 5)
|
|
|
|
assert isinstance(ptr, tir.Call)
|
|
assert ptr.op.same_as(op.Op.get("tl.access_ptr"))
|
|
assert isinstance(ptr.args[0], tir.BufferLoad)
|
|
# extent = 4*5 = 20
|
|
assert isinstance(ptr.args[1], tir.IntImm)
|
|
assert int(ptr.args[1].value) == 20
|
|
assert isinstance(ptr.args[2], tir.IntImm)
|
|
assert int(ptr.args[2].value) == 2
|
|
|
|
|
|
def test_lower_access_ptr_rewrites_to_tvm_access_ptr():
|
|
buf = tir.decl_buffer((8, 8), "float16", name="A")
|
|
load = tir.BufferLoad(buf, [tir.IntImm("int32", 2), tir.IntImm("int32", 3)])
|
|
ptr = T.access_ptr(load, "w", 4, 5)
|
|
|
|
func = tir.PrimFunc([buf.data], tir.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[tir.Call] = []
|
|
|
|
def _collect(e):
|
|
if isinstance(e, tir.Call):
|
|
calls.append(e)
|
|
|
|
tir.stmt_functor.post_order_visit(lowered["main"].body, _collect)
|
|
assert any(c.op.same_as(op.Op.get("tir.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("tir.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, tir.IntImm)
|
|
assert int(offset.value) == 19
|
|
assert isinstance(extent, tir.IntImm)
|
|
assert int(extent.value) == 20
|
|
assert isinstance(acc.args[4], tir.IntImm)
|
|
assert int(acc.args[4].value) == 2
|