## 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.
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.
## 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>
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
Tianqi Chen