## 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.
## 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.
## 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.
Bohan Hou