## 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.
This PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
## Summary
This PR introduces `AllocBufferNode`/`AllocBuffer` as a single TIR
statement that both allocates memory and declares a buffer into scope.
This replaces the previous pattern of `Allocate(var, dtype, shape, cond,
DeclBuffer(buf, body))` with the simpler `AllocBuffer(buf, body)`.
### Main changes
- **New IR node** `AllocBufferNode` with fields `{buffer, annotations,
body}` — same semantics as `DeclBuffer` but also allocates memory
- **TVMScript**: `T.alloc_buffer(shape, dtype, scope)` now emits
`AllocBuffer` directly (statement-level allocation).
`T.sblock_alloc_buffer(...)` for SBlock-level buffer allocation (full
parameter set)
- **All codegen backends** (C, CUDA, Metal, OpenCL, WebGPU, LLVM, NVPTX,
AMDGPU, SPIR-V) updated to handle `AllocBufferNode`
- **All TIR transforms** (storage_rewrite, flatten_buffer,
vectorize_loop, lower_warp_memory, etc.) updated
- **All S-TIR transforms** (compact_buffer_region, merge_shared_memory,
inject_double_buffer, etc.) updated
- **Removed `AllocateNode`** entirely — `AllocBuffer` is now the sole
allocation primitive
- **Removed `AllocDescriptor`** from merge_shared_memory_allocations —
uses `Buffer` objects directly
- **Added `AllocBuffer::ConstantAllocationSize()`** inline helper method
### Design rationale
The old `Allocate + DeclBuffer` pair was a historical artifact:
`AllocateNode` stored raw fields (`buffer_var`, `dtype`, `extents`,
`condition`) separate from the `Buffer` object, requiring pattern
matching (`IsAllocateDeclBufferPattern`) to reconstruct the buffer
association. `AllocBuffer` unifies this into a single node with a proper
`Buffer` reference, simplifying codegen backends and transform passes.
225 files changed, ~3500 insertions/deletions (net near-zero, mostly
mechanical migration).
## Test plan
- [x] All TIR base tests pass
- [x] All TIR transform tests pass
- [x] TVMScript roundtrip tests pass
- [x] S-TIR transform tests pass
- [x] Codegen tests pass
- [x] All-platform minimal tests pass
- [x] C++ functor tests pass
- [x] Pre-commit clean (clang-format, ruff, etc.)
This PR phases out alloc const node in the TIR.
This node was oroginally introduced to include embedded weights into the
allocation. However, the presence of the particular IR couples the data
allocation and the weight placement, which is not as desirable especialy
when weights get large. A better approach is to have extra annotation on
the allocation and store weights separately either outside module or as
part of module/function attribute.
As a result, we phases out this node which can help us to simplify code
logic in the codebase.
This PR renames tir.Block to SBlock. This clearly indicate the
scheduable property of the block and is a prereq for followup stir
passes refactor.
Main changes:
- Data structure change from Block to SBlock
- Syntax change from T.block to T.sblock
Prior to this commit, the automatic `T.reads()` and `T.writes()`
annotations were only generated for buffers appearing as function
arguments, as `T.alloc_buffer` in a `T.block`, or as `T.match_buffer`
in a `T.block`. However, inferred `T.reads()` for a buffer defined by
the `"tir.BindParams"` pass would be erroneously missing. These
annotations may be required for correct scheduling (see discussion in
[PR#16660](https://github.com/apache/tvm/pull/16660)).
This commit updates the TVMScript parsing to infer `T.reads()` and
`T.writes()` annotations for buffers defined with `DeclBuffer` nodes.
The current unittest folder is too large and contains too many files and
too many components. This PR refactors the unittest folder by moving the
files to the corresponding folders.
Bohan Hou