main
11 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Hongyi Jin
|
05c487d69a |
[FIX][TIRx] Preserve pointer expression types (#20070)
## 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`.
|
||
|
Tianqi Chen
|
80648af29f | [REFACTOR][TIR] Remove buffer type and axis separators (#20019) | ||
|
Tianqi Chen
|
9bfefb7e4b |
[TIRx] Introduce first-class Return statement (#20018)
Return is control flow, but TIRx currently represents it as an Evaluate-wrapped intrinsic call. This prevents return values from participating naturally in statement traversal and requires special-case handling across the pipeline. This change introduces a reflected tirx.Return statement carrying an Expr, wires it through TVMScript, statement visitors and mutators, lowering, storage planning, and C/LLVM code generation, and removes the legacy tirx.ret and T.ret surfaces. |
||
|
Tianqi Chen
|
c717c5b217 | [IR][Relax][TIRx] Unify Var identity (#20004) | ||
|
Tianqi Chen
|
adf8d6a463 |
[TIRx] Phase out duplicate Var type_annotation (#19944)
## 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.
|
||
|
Tianqi Chen
|
275114b327 |
[REFACTOR][IR] Unify PrimExpr with Expr typed view (#19910)
## 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. |
||
|
Tianqi Chen
|
1e1920bcbd |
[REFACTOR][IR] Unify PrimExpr type mechanism to PrimType instead of DataType (#19875)
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 |
||
|
Bohan Hou
|
9db74c7cee |
[TIRx] Update scoped ops and CUDA launch bounds (#19677)
## Summary - replace the block-structured TIRx exec-scope surface with scope-qualified `Tx.<scope>.<op>` namespaces and migrate call sites - split TIRx op namespaces and remove the unused dynamic generic-op fallback - add explicit CUDA launch bounds plumbing through TIRx attrs and split-host-device lowering ## Validation - `git diff --check apache/main..HEAD` - `pre-commit run --from-ref apache/main --to-ref HEAD` |
||
|
Bohan Hou
|
57c638fc7c |
[TIRx] Post-bringup op-dispatch / codegen / TVMScript follow-ups (#19657)
## Summary Follow-up work on top of the TIRx infrastructure bring-up (#19581). It extends the TIRx operator-dispatch, codegen, and TVMScript surfaces with the next batch of low-level programming features for Blackwell-class GPUs, while keeping `s_tir` script support intact. ## Main Changes - **op-dispatch**: warp `ldmatrix`/`stmatrix` copy dispatch; split CUDA copy into register / gmem-smem / `ldgsts` paths; `tcgen05.ld/st` `.16x{64,128,256}b` dispatch with a factory and M=128 layout; element-wise broadcast at the layout level with a copy vec-alignment fix. - **gemm**: CUDA synchronous `mma.sync` tensor-core dispatch; accept a Layout F C operand for M=64 MMAs. - **op**: add the `permute_layout` primitive (replaces `permute_dims`). - **tvmscript**: add the `Tx.jit` decorator, `Tx.constexpr` compile-time params, and `Tx.wg_reg_tile`. - **lower-tirx**: introduce the `Tx.device_entry()` marker (replacing `ScopeKind::kKernel`); canonical thread filters that drop the `Tx.filter` wrapper. - **codegen**: add a typed-pointer byte-offset intrinsic; remove the `entry_cluster_sync` codegen attribute. ## Validation - `pre-commit run` (changed files) — clean - `ninja -C build -j$(nproc)` — builds - `pytest tests/python/tirx/ -n 16` - `1997 passed, 39 skipped, 3 xpassed` - `python -m pytest tests/python/all-platform-minimal-test` - `37 passed, 105 skipped` - `TVM_TEST_TARGETS=llvm pytest tests/python/tirx-analysis tests/python/tirx-base tests/python/tirx-transform -n 16` - `630 passed, 25 skipped, 8 xfailed, 1 xpassed` ## Local CI Notes Several full CI-equivalent jobs are not locally reproducible because this machine is missing parts of the Apache TVM CI environment (e.g., specific `llvm-config` versions, Vulkan, ROCm, ARM/QEMU cross-toolchain, and web/wasm components). The Blackwell/Trainium kernel tests are maintained downstream and are intentionally not part of this PR. |
||
|
Bohan Hou
|
859498dc01 |
[TIRx] Bringup TIRx Infrastructure (#19581)
## Summary This PR adds the initial TIRx support needed for low-level programming of Blackwell-class GPU architectures. As part of the ongoing TIRx refactor, it introduces TVMScript support for directly scripting advanced hardware features without relying on scheduling as the primary programming interface. The change keeps existing `s_tir` script support intact while making direct scripting a first-class path for TIRx programs. ## Main Changes - Add TIRx operator dispatch and layout infrastructure. - Add TVMScript support for new low-level TIRx operations. - Add analysis, transform, and lowering support for TIRx IR nodes. - Add CUDA/Blackwell-oriented codegen and intrinsic coverage. - Add Python and C++ integration points for TIRx scripting and runtime support. ## Validation - `pre-commit run --all-files` - `ninja -C build -j32` - `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16` - `1723 passed, 47 skipped, 32 warnings` - `CUDA_VISIBLE_DEVICES=2 python -m pytest -v tests/python/all-platform-minimal-test` - `37 passed, 105 skipped` - `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis tests/python/tirx-base tests/python/tirx-transform -n 16` - `664 passed, 25 skipped, 9 xfailed, 1 xpassed` ## Local CI Notes Some full CI-equivalent jobs were not locally reproducible because this machine is missing parts of the Apache TVM CI environment, including `llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten, and ARM/QEMU cross-toolchain components. Metal-specific tests were skipped locally because no Metal runtime is available. |
||
|
Tianqi Chen
|
7504e3ed1a |
[REFACTOR][SCRIPT] TVMScript dialect-friendly refactor: per-dialect restructure + dialect registry (#19479)
## Summary Restructure TVMScript to be dialect-agnostic at the script-core layer while letting each extension dialect (TIRX, Relax) own its own per-dialect script subtree. IR is below script in the dependency stack and is NOT a peer dialect — its script handlers stay in the shared core. This PR folds together two coupled refactors that were initially opened as separate PRs (#19478 and the original #19479); they share rename / relocation surface so they ship as one cohesive change. ## What this PR does ### Per-dialect script subtree (originally #19479) - Moves per-dialect printer + builder from `src/script/{printer,ir_builder}/{tirx,relax}/` to `src/{tirx,relax}/script/{printer,builder}/`. - Tightens `src/script/*.cc` CMake glob to the dialect-free core. - Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry (`script.ir_builder.decl_function.<type-key>`); removes cross-dialect includes from the shared core. - Adds `tvm.script.register_dialect` API + `__getattr__` + a `sys.meta_path` finder for Python-side dialect discovery. In-tree dialects (tirx, relax) registered centrally in `python/tvm/__init__.py`. - Drops the obsolete static re-export shims at `python/tvm/script/{parser,ir_builder}/{tirx,relax}/`. ### Dialect-agnostic printer config (originally #19478) - Relocates `include/tvm/ir/script_printer.h` → `include/tvm/script/printer/config.h` next to the rest of the printer's public surface. The header is not IR-specific. - Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for clarity (the macro registers Script as the kRepr callback + per-type vtable dispatch). Aligns with the `TVM_REGISTER_*` family. - Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`, `relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor of a generic `ffi::Map<String, Any> extra_config` keyed by `"<dialect>.<knob>"`. Each call site reads via the templated accessor `config->GetExtraConfig<T>("...", default)`. - Promotes `std::string` config fields to `ffi::String`. After this lands, the script-printer core knows nothing specific about any dialect — new dialects plug in via the registry pattern with zero core edits. Public Python API surface unchanged. |