## 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:
```python
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`.
Hongyi Jin
05c487d69a