## Motivation and context
Copy regions can retain constant-valued expressions after variable
substitution. Converting those expressions directly to Python integers
rejects valid slices even though their extents are statically known.
## Changes
- Simplify swizzled-region extents before converting them to Python
integers.
- Continue rejecting genuinely symbolic extents.
- Add a hardware-independent regression to the existing layout test
suite.
## Testing
- `python -m pytest tests/python/tirx/test_layout.py`
- Changed-files pre-commit checks
## Motivation and context
`tcgen05.ld/st` with a `.16x*b` atom accesses one 16-row half-slab from
each 32-row TMEM partition owned by a warp. Across a four-warp
warpgroup, the two physical half-slabs are:
| View | Physical TMEM lanes | PTX row immediate |
|---|---|---|
| lower half | `0..15, 32..47, 64..79, 96..111` | `0` |
| upper half | `16..31, 48..63, 80..95, 112..127` | `16` |
Layout D for an M=128 accumulator occupies both halves. Layout F exposes
an M=64 logical tile over one half, which is useful both for a native
M=64 accumulator and for reading either 64-row half of an existing
Layout D accumulator.
Before this PR, `tmem_datapath_layout("F", 64, cols)` could describe
only the lower half. The copy dispatcher classified a TMEM buffer only
as `"D"` or `"F"` and every M=64 `.16x*b` operation started at `row=0`.
As a result, there was no layout-preserving way to create a recognized
64-row view of the upper half of a Layout D accumulator: a normal Layout
F view still addressed the lower half, while a hand-written `+16@TLane`
layout was not recognized by the dispatcher.
The half-slab selection belongs in the buffer layout because layout is
the source of truth for physical placement in TIRx. It should not be an
out-of-band `copy_async` option. This PR therefore records the selection
in Layout F, carries it through datapath classification, and derives the
PTX row immediate from it.
A related invariant is that a default M=128 TMEM allocation must be
structurally identical to named Layout D. The dispatcher recognizes
datapaths structurally, so keeping a separate hand-written default
layout creates an unnecessary drift risk. This PR makes the default call
the public Layout D factory directly.
Finally, the register-side `tcgen05_atom_layout` must agree with the PTX
mapping from a logical `(row, col)` to `(laneid, wid_in_wg, register)`.
A self-consistent load/store round trip is not enough to prove that
mapping: raw PTX can move the same bits back even when the logical
layout label is wrong. Elementwise dispatch does consume that label, so
this PR adds direct mapping and elementwise compilation fences for all
`.16x*b` atom families. These are coverage additions; the production
atom-layout construction itself is unchanged here.
## Changes
- Add `sub_slab={0,1}` to `tmem_datapath_layout("F", ...)`.
- Encode the upper view as a `+16@TLane` offset and reject invalid
selectors, including nonzero selectors for Layout D.
- Classify TMEM layouts as `(datapath, sub_slab)` and emit `.16x*b` with
`row=(sub_slab + slab) * 16`.
- Preserve the existing M=128 behavior: Layout D with a 128-row `.16x*b`
fragment still emits two operations at rows `0` and `16`.
- Build the default M=128 TMEM layout through `tmem_datapath_layout("D",
...)`.
- Document the physical lane mapping and supported datapath/atom
combinations.
- Add direct atom-layout mapping coverage and a warpgroup elementwise
regression.
## Testing
- Static Layout F checks cover every logical row for both `sub_slab=0`
and `sub_slab=1`.
- B200 readback tests populate one Layout D accumulator, then verify
that lower and upper Layout F views reproduce the two corresponding
register halves for `.16x64b`, `.16x128b`, and `.16x256b`.
- Negative tests cover invalid sub-slab values and incompatible
datapath/atom pairings.
- Direct `(row, col) -> (laneid, wid_in_wg, register)` sweeps cover
supported `.16x*b` shapes and repetitions.
- Warpgroup elementwise codegen verifies that an atom-layout fragment
canonicalizes and slices correctly.
- Changed-files pre-commit checks.
## 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`.
## Motivation and context
`MBarrier.remote_view(rank)` represents an mbarrier owned by another CTA
in the same cluster. The existing view kept only a buffer whose pointer
had been mapped to the remote CTA with PTX `mapa`. Calling
`remote_bar.arrive(...)` then followed the inherited local-arrive path
and emitted the local `mbarrier.arrive.shared.b64` form against that
mapped remote address.
A remote arrival must instead use the `shared::cluster` instruction
form. TIRx models that form with the owner-local barrier pointer plus
the destination CTA rank and predicate. Using the local instruction with
a remote address is not equivalent and is reported by synccheck as a
local arrival on a remote mbarrier address.
## Changes
- Keep the owner-local buffer and target CTA rank when constructing a
remote mbarrier view.
- Route `MBarrier` remote arrivals through the cluster helper using the
local barrier pointer and stored CTA rank.
- Apply the same routing to `TMABar`, including
`mbarrier.arrive.expect_tx`.
- Keep a typed `PointerType(uint64, shared)` mapped buffer on the view
so `ptr_to` remains available to operations that explicitly consume a
remote shared-memory pointer.
- Reject operations with ambiguous or invalid ownership semantics:
- initializing or waiting on a remote view,
- supplying another `cta_id` to a view that already fixes its target,
- creating a remote view from another remote view.
- Preserve the existing local-CTA behavior for ordinary barriers.
## Testing
- Verify the typed `mapa` binding and remote buffer in TIRx IR.
- Verify CUDA codegen for plain and counted
`mbarrier.arrive.shared::cluster.b64` forms.
- Verify CUDA codegen for remote
`mbarrier.arrive.expect_tx.shared::cluster.b64`.
- Verify that the corresponding local instruction forms are not emitted
for remote views.
- Verify diagnostics for remote init, wait, nested views, and
conflicting `cta_id` arguments.
- Run changed-files pre-commit checks.
Focused result: 3 tests passed.
## Rationale
TIRx variables use inherited `ExprNode::ty` as their single semantic
type. Retaining a primitive handle surrogate erases the distinction
between scalar values, typed pointers, and true opaque pointers, then
forces later passes and code generators to reconstruct information that
the IR already owns.
## Changes
- Remove the duplicate reflected `Var::type_annotation` state and
preserve exact `PrimType` or `PointerType` through construction,
visitors, transforms, specialization, builders, printers, and code
generation.
- Keep scalar-only boundaries explicit through `PrimExpr`, `PrimVar`,
and `PrimType`; pointer-capable values remain general `Expr` or `Var`.
- Keep helper boundaries no broader than their contracts: TE tensor
variable indices use `PrimVar`, while expression deep equality recurses
through general `Expr` only where pointer-bearing `Call` arguments
require it and does not generalize private arithmetic subclasses.
- Keep core statement reflection typed as `Expr`, name general
reinterpret targets as `target_ty`, and preserve exact pointer calls in
the general vectorization path with explicit scalarization behavior.
- Delete `PrimType::Handle()` and `PrimType::IsHandle()`. True opaque
pointers use `PointerType::VoidPointerTy()`; TVMScript renders the
canonical global type as `T.handle`, standalone values as `T.handle()`,
and scoped void pointers with a keyword-only storage scope.
- Make `CodeGenSourceBase::SSAGetID` a single `Type` boundary across
source backends, without a separate primitive-type or runtime-dtype
variant.
- Keep WebGPU semantic argument classification type-aware: storage
buffers are identified from `PointerType`, POD arguments from
`PrimType`, and only the final `FunctionInfo` launch ABI is serialized
to `DLDataType`.
- Preserve exact pointer semantics at runtime boundaries, including
access pointers, packed calls and returns, external calls, storage
rewrites, and target-specific lowering.
## Migration guide
- **Variable types:** In C++, replace `var->type_annotation` with
`var->ty`; in Python, replace `var.type_annotation` with `var.ty`. The
result is the exact `Type`: scalar variables carry `PrimType`, while
pointer variables carry `PointerType`.
- **Scalar boundaries:** Use `PrimVar` and `PrimExpr` for variables and
expressions that are semantically scalar. When starting from a general
view, narrow explicitly with `var.as_or_throw<PrimVar>()` or
`expr.as_or_throw<PrimExpr>()`. Keep pointer-capable fields and call
arguments as `Var` or `Expr`. A default-constructed `PrimVar` is
nullable, so construct local scalar variables explicitly, for example
`PrimVar i("i")`.
- **Opaque pointers:** Replace `PrimType::Handle()` with
`PointerType::VoidPointerTy()`. Replace `IsHandle()` tests with explicit
`PointerType` inspection; use `PointerType(element_type, storage_scope)`
when the pointee type is known instead of erasing it to a runtime handle
dtype.
- **TVMScript handles:** Use `arg: T.handle` for a global void-pointer
annotation and `arg = T.handle()` for a standalone value. Use
`T.handle(storage_scope="shared")` for a scoped void pointer. Typed
pointers use forms such as `T.handle("float32")`, `T.handle("float32",
"global")`, or `T.handle("float32", "shared")`. Legacy
`T.handle("void")` input remains parse-compatible, but the printer
canonicalizes it to `T.handle` (or the keyword-only scoped form).
- The separate `tirx.type_annotation` intrinsic used by access-pointer
APIs is unchanged; this migration removes only the duplicate variable
field.
## Validation
- Complete native C++ test executable: 122/122 passed, including
`IRF.CountVar`.
- Relax binding-rewrite suite: 12/12 passed, including transferred-user
bookkeeping.
- Canonical typed/void/scoped TVMScript handle printer and round-trip
checks: 5/5 passed.
## Summary
- Make `PrimExpr` a typed C++ view over `Expr` values whose
`ExprNode::ty` is `PrimType`, instead of using a separate runtime node
class as the proof of primitive-ness.
- Use the shared `ir::Call` node for Relax, TIRX, and primitive-valued
calls, while keeping primitive-only APIs explicit at their semantic
boundaries.
- Keep Python on the general `Expr` surface for primitive-typed values
so `isinstance` behavior does not imply a nominal primitive-expression
subclass.
## Design Rationale
The main advantage of this change is that common expression nodes such
as `Call` can be unified without specializing each one to `PrimType`. A
single `ir::Call` can represent a Relax tensor call, a Relax scalar
call, or a primitive-valued intrinsic call; the result type stored in
`ExprNode::ty` determines whether that particular value can be viewed as
`PrimExpr`.
This keeps the IR node hierarchy focused on expression structure rather
than result-type categories. Nodes that are intrinsically primitive,
such as integer and floating-point literals or TIRX primitive operators,
still have strongly typed C++ APIs and data structures. General nodes
whose result type may vary, such as `Call`, remain general `Expr` nodes
and are narrowed to `PrimExpr` only where primitive-only semantics are
required.
The PR also keeps the compatibility surface practical: C++
primitive-only APIs continue to accept `PrimExpr`, Python exposes a
compatibility predicate for checking the primitive typed category, and
visitors/printers use one natural `Call` path rather than duplicating
Relax and primitive call handling. Missing expression types are
represented explicitly with `Type::Missing()` so constructors can leave
type inference to later analysis without relying on nullable `Type`
values.
This PR lets Relax expressions directly take `PrimExpr` values without
requiring the explicit `PrimValue` wrapper, continuing the Relax IR
unification work by removing Relax-specific leaf/base expression layers.
Summary:
- Remove `LeafExpr` / `LeafExprNode` and use direct expression-node
checks where needed.
- Converge Relax expression typing onto the shared IR `Expr` base.
- Remove the `PrimValue` node wrapper while keeping `relax.prim_value` /
`R.prim_value` as conversion helpers that return existing `PrimExpr`
values unchanged.
- Register direct `PrimExpr` handling through exact concrete node
dispatch, aligned with the `tirx` expression visitor list and excluding
arith iter-map intermediate nodes.
- Inline the private Python primitive conversion helper into public
`relax.prim_value`.
- Handle direct `PrimExpr` values in frontend scalar paths without
assuming a `.value` field on non-immediate expressions.
In the past we have been using `DataType` in PrimExpr.dtype field to
check type information for PrimExpr while still having BaseExpr.ty for
richer type information. DataType is also used both in runtime and
compiler. This PR streamlines the boundary:
- PrimExpr.ty now carries PrimType that replaces original use of
`DataType`
- Runtime use will now favor DLPack DLDataType, removing one layer of
indirection.
- Constants attributes where values are usually runtime values, will use
`DLDataType`
- DataType will be phased out after this PR
We also brings up helper functions in PrimType, but also limits them to
a more concise set so the functions do not grow with the data type codes
in DLPack.
This is a major refactor that changes the IR primitive. It helps to
bring possible future benefits:
- Unified type mechanism through Expr.ty
- Possibility of carry future Type nodes
Migration Guide:
- Use `PrimType` when code reasons about compiler expression types,
tensor element compiler types, or constructs a `PrimExpr`/compiler type.
- Use existing source types such as `expr.ty()`, `ExprOp.expr_ty()`, or
TE tensor element `dtype` where possible instead of rebuilding a type
from dtype text.
- Use raw `DLDataType` for runtime constants, ABI paths, dtype-valued
attrs, and storage/runtime helper logic.
- Prefer direct `PrimType` equality, `MatchesCode(...)`,
`MatchesElementType(...)`, and `WithCode(...)` over local wrappers or
string dtype checks.
Performance:
Using Object type instead of DLDataType would indeed bring some
performance impact to the IR. We have done the following performance
optimizations:
- Make sure most of the outputs reuse one of the PrimType from inputs
- Cache a thread local PrimType based on input so we don't repeatly
realloc
We did benchmarks show that rewrite simplify operation stays within
+-10% overhead of original one. Which merits the refactor given the
benefit the unfication brings
PR #19677 registered every CUDA / Trainium device intrinsic under two Op
names: a flat `tirx.<ns>_<name>` alias plus the canonical
`tirx.<ns>.<name>`. The flat aliases were a migration shim; passes and
codegen that match an intrinsic had to check both spellings (the
dual-name IsOp pattern). The Python builders and TVMScript parser
already canonicalize, so every real Call already carries the canonical
op and the flat aliases were dead weight.
This pr removes the flat device-intrinsic aliases, keeping only the
canonical namespaced ops:
- RegisterDeviceIntrinsic (backend/cuda) and RegisterNKIIntrinsic
(backend/trn) register only the canonical name.
- Drop the flat-only macro registrations for device intrinsics; the
canonical op with all attrs is registered from the alias table. The WMMA
tvm_*_sync / mma_store / mma_fill builtins and the profiling
timer_*_cuda builtins keep their flat names (no namespace / canonical
form, category "builtin").
- Remove the redundant flat tirx.ptx_fetch_register registration.
- C++ consumers that resolved a flat op by name string now use the
canonical name; the ptx_elect_sync / cuda_func_call dual-name matchers
collapse to the canonical check.
- Python: the InjectPTXAsyncCopy round-trip Op.get and the matching test
assertion use the canonical name. call_intrin keeps its flat->canonical
rewrite for back-compat, so user-facing wrappers are unchanged.
- test_op_namespace_cleanup asserts device_intrin op names are canonical
so a flat alias cannot silently reappear.
Generated CUDA is byte-identical: helper names are literals and codegen
dispatches by op name, with the registry resolving the canonical name to
the same helper.
The Jenkins PR title/body linter is comparatively heavy and can report
false positives before the normal CI signal is available.
This removes the check_pr step from the Jenkins prepare flow and drops
the now-unused script-level test coverage.
## Why
ASF INFRA enforces that external GitHub Actions must be pinned to a
commit SHA on the approved allowlist, failing the workflow with "not
allowed in apache/tvm". See the
[policy](https://infra.apache.org/github-actions-policy.html) and the
[approved
allowlist](https://github.com/apache/infrastructure-actions/blob/main/approved_patterns.yml).
## How
- Pin `pre-commit/action` to `2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd`
(v3.0.1)
- Pin `pypa/cibuildwheel` to `294735312765b09d24a2fbec22660ce817587d55`
(v4.1.0)
- Pin `pypa/gh-action-pypi-publish` to
`ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e` (v1.13.0)
- Leave GitHub-owned `actions/*` and the allowlisted
`conda-incubator/setup-miniconda@*` pattern untouched
---------
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
## Summary
Backend loading is easier to maintain when native backend library
discovery, in-tree backend Python hook loading, and out-of-tree entry
point autoload are owned by the backend namespace. This PR consolidates
those paths under `tvm.backend._autoload_backends` while preserving
compatibility routes from the previous top-level helper and
`tvm.base.load_backend_libs`.
- Move backend runtime DSO loading into `tvm.backend._autoload_backends`
- Route `backend.load_all()` through the backend autoload helper
- Keep the previous top-level `_autoload_backends` module as a thin
compatibility import
This PR batches several post-bringup TIRx follow-ups, rebased onto
current `main`.
### Changes
- **op-dispatch**: per-call exec scope via `Tx.<scope>.op`; remove
`ExecScopeStmt`
- **namespaces**: split TIRx op namespaces; remove tile-primitive kind
attrs
- **codegen**: support explicit CUDA launch bounds
- **gemm-async**: support contiguous-axis (K-major) operand slicing
- **backend reorg**: move in-tree GPU backends out of core into
`src/backend/<target>/` and `python/tvm/backend/<target>/`
(codegen/runtime/op), with the corresponding `CMakeLists.txt` /
`cmake/modules` and include-path updates
### Testing
- Builds with `USE_CUDA=ON` / `USE_LLVM=ON`
- The TIRx Python test suite (`tests/python/tirx/`) passes locally
Hongyi Jin