[FIX][TIRx] Preserve pointer expression types

[jinhongyii]

Motivation and context

A TIRx pointer carries two pieces of information that later lowering needs: the pointee element type and the storage scope. Both must survive when a pointer-producing expression is assigned to a Python name and then used as the backing storage of a buffer.

A concrete example is accessing an mbarrier in another CTA through distributed shared memory:

ptr_ty = PointerType(PrimType("uint64"), "shared")
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"
)

map_shared_rank returns the raw uint64 address produced by PTX mapa, and reinterpret gives that address the intended PointerType(uint64, shared). Because decl_buffer(data=...) requires a pointer Var, assigning the expression to remote_ptr should create an immutable typed pointer binding.

Before this PR, an unannotated assignment such as remote_ptr = <pointer expression> followed the same parser path as a numeric assignment. That path allocates a mutable local scalar and therefore cannot represent a PointerType. The pointer expression could not be carried as a correctly typed Var into decl_buffer and CUDA lowering.

This PR makes an unannotated pointer-valued assignment emit a TIRx Bind. The bound Var has exactly the type of the right-hand side, including its element type and storage scope. Pointer bindings are immutable, so reassignment in the same scope is diagnosed; shadowing a name supplied through extra_vars remains valid. Numeric assignments keep their existing mutable-local behavior.

Type propagation fixes

The parser fix exposed several other boundaries where pointer type information must remain consistent:

Boundary	Previous behavior	Behavior after this PR
Unannotated pointer assignment	Tried to materialize the value as a local scalar	Emits an immutable Bind with the RHS PointerType
address_of(buffer) / buffer.ptr_to(...)	Reused the raw backing pointer type	Returns a pointer to buffer.dtype while preserving the backing pointer storage scope
tvm_access_ptr / ptr_byte_offset	Accepted strings or annotation expressions, but not a PrimType object directly	Accepts PrimType and produces the corresponding typed pointer
Printed T.ptx.mapa call	The printer emits all intrinsic attributes positionally, but the Python helper required keyword-only arguments	Accepts the canonical printed form so pointer code round-trips through TVMScript

The address_of distinction matters for typed views over byte-addressed storage. For example, if a float32 buffer is backed by a uint8* allocation in shared.dyn, the address of a buffer element must be PointerType(float32, shared.dyn), not PointerType(uint8, shared.dyn).

With these changes, the DSMEM example above round-trips through TVMScript and CUDA codegen declares the remote buffer pointer as uint64_t*.

TMA dtype normalization

This PR also contains a small, separate type-representation fix in TMA descriptor construction. TmaPlan.elem_dtype is a string consumed by the host-side runtime.cuTensorMapEncodeTiled packed call, but _assemble_plan stored g_buf.dtype, which is a PrimType. Converting it with str(g_buf.dtype) ensures that the generated packed-call argument is StringImm("float16") rather than an IR type object. This does not change the pointer-binding semantics described above.

Testing

• Verify that an unannotated pointer expression creates a Bind whose Var type matches the RHS type.
• Verify that pointer reassignment is rejected while shadowing an extra_vars name is allowed.
• Verify parser/printer structural round-tripping for the pointer binding and canonical T.ptx.mapa call.
• Verify that address_of uses the logical buffer element type and preserves the storage scope for byte-backed buffer views.
• Verify that tvm_access_ptr and ptr_byte_offset accept PrimType inputs.
• Compile the DSMEM map_shared_rank example through the CUDA TIRx pipeline and check for a typed uint64_t* remote buffer pointer.
• Verify that the TMA host initialization passes the descriptor dtype as a StringImm.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.