diff --git a/docs/tirx/native_basics/cuda/buffers.rst b/docs/tirx/native_basics/cuda/buffers.rst index feb2763f6252..18de83ea033b 100644 --- a/docs/tirx/native_basics/cuda/buffers.rst +++ b/docs/tirx/native_basics/cuda/buffers.rst @@ -38,8 +38,9 @@ Two fundamental APIs create a buffer: it allocates, like ``alloc_buffer``. A buffer's ``data`` pointer is an immutable ``Var`` (``alloc_buffer`` defines it; -``decl_buffer`` takes one). To back a buffer with a pointer *expression*, bind it -first — see :doc:`data_types`. +``decl_buffer`` takes one). To back a buffer with a pointer *expression*, assign +it to a name first; the parser creates an immutable pointer binding. See +:doc:`data_types`. Both share one descriptor; the parameters that matter most: @@ -439,6 +440,11 @@ to an intrinsic or inline function; ``data`` is the base pointer: B_ptr[tx] = ld(&A_ptr[tx]); // ptr_to([tx]) -> &A_ptr[tx]; A.data -> A_ptr +The pointer returned by ``ptr_to`` has the buffer's element type and storage +scope. This remains true when the buffer is a typed view over a byte-addressed +allocation pool; the pool's raw backing-pointer type does not leak through the +element address. + **Vectorized access — ``vload`` / ``vstore``.** Move several elements as one wide transfer (see also :doc:`data_types`): diff --git a/docs/tirx/native_basics/cuda/data_types.rst b/docs/tirx/native_basics/cuda/data_types.rst index 4b17543566a4..884cfdb07dd1 100644 --- a/docs/tirx/native_basics/cuda/data_types.rst +++ b/docs/tirx/native_basics/cuda/data_types.rst @@ -103,14 +103,19 @@ A buffer's ``data`` — its pointer — is a ``Var`` of pointer type, and it is - ``T.decl_buffer(..., data=ptr)`` declares a buffer over an existing pointer ``Var`` ``ptr``. - To back a buffer with a pointer **expression** — e.g. ``T.ptx.map_shared_rank`` - (PTX ``mapa``) giving another cluster CTA's shared address — you must first bind - that expression to a pointer ``Var`` (``data`` must be a ``Var``, not an - expression), using a ``T.let`` of ``PointerType``: + (PTX ``mapa``) giving another cluster CTA's shared address — convert the raw + ``uint64`` address returned by ``map_shared_rank`` to a pointer with the + intended element type and storage scope. Assigning that pointer expression + to an unannotated name automatically creates an immutable ``Bind`` (and + ``data`` receives the resulting pointer ``Var``): .. code-block:: python from tvm.ir.type import PointerType, PrimType - ptr: T.let[T.Var(name="ptr", ty=PointerType(PrimType("uint64")))] = \ - T.reinterpret("handle", T.ptx.map_shared_rank(mbar.ptr_to([0]), 0)) + ptr_ty = PointerType(PrimType("uint64"), "shared") + ptr = T.reinterpret(ptr_ty, T.ptx.map_shared_rank(mbar.ptr_to([0]), 0)) remote_mbar = T.decl_buffer([1], "uint64", data=ptr, scope="shared") + + Pointer bindings cannot be reassigned; use a new name for a different + pointer value. diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py index 13bf481543c3..988e48f4d01d 100644 --- a/python/tvm/backend/cuda/op.py +++ b/python/tvm/backend/cuda/op.py @@ -4209,7 +4209,7 @@ def ptx_map_shared_rank(ptr, rank): return ptx_mapa(ptr, rank, space="", ptx_type="u64", return_type="uint64") -def ptx_mapa(ptr, rank, *, space="", ptx_type="u64", return_type="uint64"): +def ptx_mapa(ptr, rank, space="", ptx_type="u64", return_type="uint64"): """TVM intrinsic for PTX ``mapa{.space}.type d, a, b``.""" if space not in ("", "shared::cluster"): raise ValueError(f"Unsupported mapa space {space!r}") diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py index 8f78679de203..e83cd827c2c1 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py @@ -812,7 +812,7 @@ def _assemble_plan( issue_axes=issue_axes, tensor_ptr=g_buf.data, elem_bytes=elem_bytes, - elem_dtype=g_buf.dtype, + elem_dtype=str(g_buf.dtype), ) return _merge_contig_full_box_dims(plan, analyzer) diff --git a/python/tvm/script/parser/core/parser.py b/python/tvm/script/parser/core/parser.py index 89ccdcddaf7f..c94c8fe97fea 100644 --- a/python/tvm/script/parser/core/parser.py +++ b/python/tvm/script/parser/core/parser.py @@ -288,6 +288,10 @@ def get(self) -> dict[str, Any]: """ return {key: values[-1] for key, values in self.name2value.items() if values} + def contains_in_current_frame(self, name: str) -> bool: + """Check whether a variable name exists in the current frame.""" + return bool(self.frames) and name in self.frames[-1].vars + def get_at_depth(self, depth: int) -> dict[str, Any]: """Get variables visible at the given frame depth, using current values. diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 782fbcfe58be..35fc990720d1 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -59,6 +59,8 @@ def _canonical_device_intrin_name(func_name: str) -> str: def _primexpr_ty(expr): """Return the runtime primitive type of an expression.""" + if isinstance(expr, tvm.ir.PrimType): + return expr ty = getattr(expr, "ty", None) if isinstance(ty, tvm.ir.PrimType): return ty @@ -658,6 +660,13 @@ def _is_tensormap_var(obj: Var) -> bool: return isinstance(obj.ty, PointerType) and isinstance(obj.ty.element_type, TensorMapType) +def _buffer_element_pointer_type(buffer: Buffer) -> PointerType: + """Return a pointer to ``buffer`` elements in the buffer's storage scope.""" + data_ty = buffer.data.ty + storage_scope = data_ty.storage_scope if isinstance(data_ty, PointerType) else "global" + return PointerType(buffer.dtype, storage_scope) + + def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr: """Returns the address of a buffer element or addressable variable. @@ -677,7 +686,12 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr if isinstance(obj, Buffer): n_dim = len(obj.shape) buffer_load = BufferLoad(obj, [0] * n_dim) - return Call("tirx.address_of", [buffer_load], span=span, ret_ty=obj.data.ty) + return Call( + "tirx.address_of", + [buffer_load], + span=span, + ret_ty=_buffer_element_pointer_type(obj), + ) elif isinstance(obj, Var): if _is_tensormap_var(obj): return call_intrin("uint64", "tirx.address_of", obj, span=span) @@ -685,7 +699,12 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> Expr raise TypeError(f"address_of expects a scalar or TensorMap Var, but got {obj.ty}") return Call("tirx.address_of", [obj], span=span, ret_ty=PointerType(obj.ty)) elif isinstance(obj, BufferLoad): - return Call("tirx.address_of", [obj], span=span, ret_ty=obj.buffer.data.ty) + return Call( + "tirx.address_of", + [obj], + span=span, + ret_ty=_buffer_element_pointer_type(obj.buffer), + ) else: raise ValueError(f"Invalid object type: {type(obj)}") @@ -918,11 +937,11 @@ def tvm_access_ptr(ptype, data, offset, extent, rw_mask): Parameters ---------- - ptype : Expr or str - The data type of pointer. If a ``str``, it is wrapped via - :func:`type_annotation` so that the lowering rule (which reads - ``args[0].dtype()`` for the cast type) sees the intended dtype - instead of ``void`` from a raw StringImm. + ptype : Expr, PrimType, or str + The data type of pointer. If a ``PrimType`` or ``str``, it is wrapped + via :func:`type_annotation` so that the lowering rule (which reads + ``args[0].dtype()`` for the cast type) sees the intended dtype instead + of ``void`` from a raw StringImm. data : DType* The data of pointer. @@ -941,7 +960,7 @@ def tvm_access_ptr(ptype, data, offset, extent, rw_mask): call : Expr The call expression. """ - if isinstance(ptype, str): + if isinstance(ptype, str | tvm.ir.PrimType): ptype = type_annotation(ptype) data_type = getattr(data, "ty", None) storage_scope = data_type.storage_scope if isinstance(data_type, PointerType) else "global" @@ -962,7 +981,7 @@ def ptr_byte_offset(data, byte_offset, dtype): ``byte_offset`` is always in bytes. Use this when the source CUDA shape needs an explicitly typed local pointer derived from a byte-addressed base. """ - if isinstance(dtype, str): + if isinstance(dtype, str | tvm.ir.PrimType): dtype = type_annotation(dtype) data_type = getattr(data, "ty", None) storage_scope = data_type.storage_scope if isinstance(data_type, PointerType) else "global" diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index ddbfac29ab18..40520bdb639a 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -23,7 +23,7 @@ from typing import Any import tvm -from tvm.ir import Expr, GlobalVar, PrimType +from tvm.ir import Expr, GlobalVar, PointerType, PrimType from tvm.script.ir_builder import ir as I from tvm.script.ir_builder.base import IRBuilder from tvm.script.ir_builder.base import IRBuilderFrame as Frame @@ -217,6 +217,15 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) - IRBuilder.name(var_name, value) return value else: + is_pointer_expr = isinstance(value, tvm.ir.Expr) and isinstance( + getattr(value, "ty", None), PointerType + ) + if is_pointer_expr: + if self.var_table.contains_in_current_frame(var_name): + self.report_error(node, f"Pointer variable '{var_name}' cannot be reassigned") + ann_var = T.Bind(value) + IRBuilder.name(var_name, ann_var) + return ann_var if not tvm.ir.is_prim_expr(value): value = tvm.tirx.const(value) if not isinstance(value, tvm.tirx.StringImm): diff --git a/tests/python/tirx-base/test_tir_op_types.py b/tests/python/tirx-base/test_tir_op_types.py index 78e379b76524..e0e7bf1f4272 100644 --- a/tests/python/tirx-base/test_tir_op_types.py +++ b/tests/python/tirx-base/test_tir_op_types.py @@ -47,6 +47,11 @@ def test_tir_op_address_of(): buffer = tirx.decl_buffer((128), "float32") expr = tirx.address_of(buffer[0]) assert expr.op.name == "tirx.address_of" + storage = tirx.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"), "shared.dyn")) + pooled_buffer = tirx.decl_buffer((128), "float32", data=storage, scope="shared.dyn") + expected_ty = tvm.ir.PointerType(tvm.ir.PrimType("float32"), "shared.dyn") + assert tirx.address_of(pooled_buffer).ty == expected_ty + assert tirx.address_of(pooled_buffer[0]).ty == expected_ty scalar_address = tirx.address_of(tirx.Var("value", "uint32")) assert scalar_address.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint32")) @@ -116,6 +121,11 @@ def test_tir_op_tvm_access_ptr(): assert expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("float32")) offset_expr = tirx.ptr_byte_offset(buffer.data, 16, "uint8") assert offset_expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8")) + prim_type = tvm.ir.PrimType("uint8") + typed_access_expr = tirx.tvm_access_ptr(prim_type, buffer.data, 0, 1, 2) + assert typed_access_expr.ty == tvm.ir.PointerType(prim_type) + typed_offset_expr = tirx.ptr_byte_offset(buffer.data, 16, prim_type) + assert typed_offset_expr.ty == tvm.ir.PointerType(prim_type) def test_tir_op_tvm_throw_last_error(): diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py b/tests/python/tirx/codegen/test_codegen_dsmem.py index 052034dd0d20..1a960d724ee3 100644 --- a/tests/python/tirx/codegen/test_codegen_dsmem.py +++ b/tests/python/tirx/codegen/test_codegen_dsmem.py @@ -19,6 +19,7 @@ import tvm import tvm.testing +from tvm.ir import PointerType, PrimType, assert_structural_equal from tvm.script import tirx as T @@ -87,7 +88,48 @@ def main(A: T.Buffer((64,), "float32")): assert "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" in src +def test_ptx_map_shared_rank_pointer_bind_codegen(): + ptr_ty = PointerType(PrimType("uint64"), "shared") + + # fmt: off + @T.prim_func + def main(A: T.Buffer((1,), "uint64")): + T.device_entry() + cta_id = T.cta_id([1]) + tid = T.thread_id([1]) + mbar = T.alloc_shared([2], "uint64") + remote_ptr = T.reinterpret( + ptr_ty, T.ptx.map_shared_rank(mbar.ptr_to([0]), T.int32(0)) + ) + remote_mbar = T.decl_buffer([1], "uint64", data=remote_ptr, scope="shared") + A[0] = remote_mbar[0] + # fmt: on + + binds = [] + loads = [] + + def collect(node): + if isinstance(node, tvm.tirx.Bind): + binds.append(node) + elif isinstance(node, tvm.tirx.BufferLoad): + loads.append(node) + + tvm.tirx.stmt_functor.post_order_visit(main.body, collect) + assert len(binds) == 1 + assert isinstance(binds[0].var.ty, PointerType) + assert binds[0].var.ty.storage_scope == "shared" + assert binds[0].value.ty.storage_scope == "shared" + assert_structural_equal(binds[0].var.ty, binds[0].value.ty) + assert any(load.buffer.data.same_as(binds[0].var) for load in loads) + + assert_structural_equal(main, tvm.script.from_source(main.script())) + src = _get_source(main) + assert "uint64_t* remote_mbar_ptr" in src + assert "tvm_builtin_ptx_mapa_u64" in src + + if __name__ == "__main__": test_ptx_cp_async_bulk_s2c_codegen() test_ptx_cp_async_bulk_s2c_codegen_address_conversion() + test_ptx_map_shared_rank_pointer_bind_codegen() print("All codegen tests passed!") diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index 7cbce484ba45..0189493049e3 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py @@ -1046,6 +1046,20 @@ def test_copy_tma_codegen(case): tvm.ir.assert_structural_equal(host_init_stmts[0], expected_host, map_free_vars=True) +def test_copy_tma_host_init_dtype_is_string(): + _, host_init_stmts = _make_tma_call( + g_shape=(8, 256), + g_region=((0, 8), (0, 256)), + s_shape=(8, 256), + s_region=((0, 8), (0, 256)), + gmem_layout=TileLayout(S[8, 256]), + smem_layout=TileLayout(S[8, 256]), + ) + encode_call = host_init_stmts[0].seq[1].value + assert isinstance(encode_call.args[2], StringImm) + assert encode_call.args[2].value == "float16" + + # Section 3: TMA special cases (symbolic dimension, buffer view) # =========================================================================== diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index e570f72200c2..6c5627f36ba1 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -1376,6 +1376,61 @@ def func() -> None: assert from_source(code).script() == code +def test_pointer_expression_assignment_uses_bind(): + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + buf = T.alloc_buffer((4,), "uint32", scope="shared") + ptr = buf.ptr_to([1]) + T.evaluate(T.reinterpret("uint64", ptr)) + # fmt: on + + binds = [] + tvm.tirx.stmt_functor.post_order_visit( + func.body, lambda node: binds.append(node) if isinstance(node, tvm.tirx.Bind) else None + ) + assert len(binds) == 1 + assert isinstance(binds[0].var.ty, PointerType) + assert_structural_equal(binds[0].var.ty, binds[0].value.ty) + + code = func.script() + assert_structural_equal(func, from_source(code)) + + +def test_pointer_expression_assignment_rejects_reassignment(): + with pytest.raises(tvm.error.DiagnosticError, match="cannot be reassigned"): + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + buf = T.alloc_buffer((4,), "uint32", scope="shared") + ptr = buf.ptr_to([0]) + ptr = buf.ptr_to([1]) + T.evaluate(T.reinterpret("uint64", ptr)) + # fmt: on + + +def test_pointer_expression_assignment_can_shadow_extra_var(): + source = """ +@T.prim_func +def func() -> None: + T.device_entry() + buf = T.alloc_buffer((4,), "uint32", scope="shared") + ptr = buf.ptr_to([1]) + view = T.decl_buffer((3,), "uint32", data=ptr, scope="shared") + view[0] = T.uint32(0) +""" + func = tvm.script.from_source(source, extra_vars={"T": T, "ptr": object()}) + + binds = [] + tvm.tirx.stmt_functor.post_order_visit( + func.body, lambda node: binds.append(node) if isinstance(node, tvm.tirx.Bind) else None + ) + assert len(binds) == 1 + assert_structural_equal(func, from_source(func.script())) + + def test_buffer_permute_ir(): """Verify .permute(1, 0): shape swapped, layout permuted, shared data."""