[FIX][TIRx] Preserve pointer expression types

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`):
 

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.

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}")

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)
 

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.
 

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,15 +686,25 @@ 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)
         if not isinstance(obj.ty, tvm.ir.PrimType):
             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"

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):

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():

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!")

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)
 # ===========================================================================
 

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."""