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.
## 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 moves remaining related data structures to s_tir.
- Moves sblock_dependency_info and sblock_scope.
- Moves related analyssis.
- Hides the data_type_rewriter to private functions.
Tianqi Chen