## Summary
Compiler Targets can carry device-type semantics that runtime
device-name parsing does not preserve.
- add `tvm.device_from_target` for canonical Target-to-Device
translation
- use explicit runtime constructors where the device kind is fixed
- update target-derived utilities, tests, and documentation to use the
explicit boundary
## Rationale
Analyzer constraint scopes continue to provide loop-positive facts to
constant-bound and rewrite proofs. During IntSet relaxation, however, a
scoped domain constraint is a refinement only for a variable explicitly
present in the relaxation map. Applying it to an unmapped variable
reinterprets a free parameter as a relaxation domain and can let a
loop-local symbol survive recursive interval evaluation.
## Changes
- Apply scoped IntSet constraints only to variables already present in
the relaxation map; unmapped variables remain free parameters.
- Remove the finite-bound restoration fallback and its stronger
parametric-bound contract.
- Add a direct compact-buffer regression that prevents a loop-local
variable from escaping into a function-scope allocation extent.
## Summary
- Keep the Python test launcher close to plain `pytest -n auto`, move
nightly tests under `tests/nightly/python`, remove obsolete launchers
and collection bookkeeping, and partition CPU/GPU jobs with explicit
`gpu` marker expressions.
- Repair exact-pointer regressions at their owning boundaries: packed
raw-string ABI values, CUDA/Metal matrix intrinsic pointers, internal TE
extern offsets, MetaSchedule scalar annotations, localized
auto-tensorization scope matching, and typed DLTensor fixture fields.
- Preserve typed workspace calls in TIR and cast pointer-returning
external calls in CodeGenC, covered by a plain-TIRx 1024-byte global
workspace that is compiled as C++.
- Finish phasing out value-bearing Relax `R.Prim` annotations by
requiring an explicit dtype, removing obsolete value-based contracts,
and expressing the DISCO rank-dependent slices as explicit scalar
`call_tir` inputs.
- Gate the distributed callback on the optional DISCO runtime, NCCL, and
at least two GPUs so capability-limited jobs skip instead of failing.
- Remove the non-demonstrating pointer probe, use direct TVMScript
comparison for packed strings, and remove the four designated legacy
testing modules.
The seven repaired CPU categories cover packed raw strings (7 failures),
CUDA/Metal matrix access-pointer types (7), internal TE extern offsets
(1), a typed DLTensor fixture (1), MetaSchedule scalar annotations (1),
CodeGenC workspace return casts (12), and localized auto-tensorization
storage-scope matching (19).
## Validation
- Base: `ded6ad8dd212869c881efb5590f8a33fc972728e`
- Head: `a7277e86dbcfe0638c8c252d36760859c4ab4297`
- All 35 locally available original failing node IDs pass across the
focused runs.
- The full focused TE, TIR builtin-lowering, and CodeGenC files pass: 61
tests.
- The complete touched Relax/TVMScript set plus
PlanAndUpdateBufferAllocationLocation passes with 784 passed, 20
skipped, and 1 expected failure.
- The DISCO callback collects and skips when its runtime or two-GPU
environment is unavailable.
- Six direct mapping tests, twelve tensor-core sketches, and the dp4a
sketch pass unchanged.
- The compiler rebuild, branch-wide pre-commit hooks, and full-range
whitespace checks pass.
- The 13 broad CBLAS/TFLite nodes remain dependency-gated; their owning
TE and generated-C regressions compile.
No merge is included in this change.
This PR simplifies Jenkins pytest execution around standard pytest-xdist
behavior.
- Runs each already-filtered CPU/GPU suite once with `-n auto`; the
broad suite keeps load-group scheduling because its order-sensitive
cases require it.
- Removes external sharding, wrapper/profile code, JUnit XML generation
and publication, the skipped-test XML consumer, obsolete suite naming,
and orphaned helpers.
- Retains one inert `task_clear_pytest.sh` entry point only because PR
jobs evaluate their Jenkinsfile from the trusted base branch before
checking out the PR; it performs no cleanup or reporting and can be
removed after this pipeline lands.
- Corrects stale broad-suite paths and explicit target guards, and
migrates a scalar stride test to the current `T.handle` pointer
semantics while preserving its negative lowering check.
- Prevents nested MetaSchedule/XGBoost unit tests from multiplying CPU
fanout without serializing the full suite.
- Builds only the `tvm_runtime` target for the secondary GPU
configuration and removes its unconsumed `gpu2` artifact upload.
The result reduces parallelism to one layer managed by pytest-xdist
while preserving GPU filtering and native failure visibility.
## 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.
Add tvm.testing.run_with_gpu_lock backed by the existing
tvm_ffi.utils.FileLock. Migrate live local GPU tests to acquire the
machine-local lock around device execution, synchronization, host
transfer, and checks while leaving target construction and compilation
outside the critical section.
Replace the custom xdist scheduler with standard xdist_group placement
for the order-dependent test family. RPC tests retain dynamic port
allocation and per-test process isolation rather than gaining a broad
category lock.
## 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.
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.
This pr modernizes test gating. It replaces the heavy
`tvm.testing.Feature` machinery with a thin `tvm.testing.env` module of
`has_*()` capability probes, used via standard pytest.mark + skipif. And
markers move to `pyproject.toml`
This fixes the legacy predicated `ptx.cp_async` codegen path used by
`InjectPTXAsyncCopy` for `if_then_else(..., 0)` stores.
The old inline CUDA emission zero-filled the shared-memory destination
when the predicate was false. The TIRx helper-based legacy codegen only
skipped the `cp.async`, leaving the destination slot unchanged. This
restores the previous behavior by emitting an `@!p st.shared.*` zero
store in the generated legacy predicated helper.
The CUDA source snapshot in
`test_s_tir_transform_inject_ptx_async_copy.py` is updated to reflect
the restored false-predicate zero-fill instruction and the current
generated helper-based CUDA source.
This pr marks it xfail with a TODO so the s_tir/transform CI enrollment
(#19737) is not blocked; the mark should be removed once the CSE
determinism fix land.
This pr updates the WebGPU multi-warp allreduce test to check the
generated `tirx.volatile` allocation annotation structurally instead of
matching the exact TVMScript printer output.
The test is intended to verify that `LowerThreadAllreduce` marks the
generated shared allocation as volatile. It previously checked for the
exact string:
```python
"tirx.volatile": T.bool(True)
```
However, the current printer emits the same annotation as:
```python
annotations={"tirx.volatile": True}
```
The transform behavior is unchanged; only the printer spelling differs.
This patch walks the generated TIRX body and checks for an `AllocBuffer`
with `tirx.volatile=True`, which matches the actual semantic requirement
of the test without depending on bool literal formatting.
Update `test_multiplication_nodes_are_inlined` to use the current TIRX
PTX async script namespace:
- `T.ptx.cp_async.commit_group()`
- `T.ptx.cp_async.wait_group(0)`
The test still used the older top-level names `T.ptx_commit_group()` and
`T.ptx_wait_group(0)`, which are not exposed by the current
`tvm.tirx.script` namespace. This caused parsing to fail before
`InjectPTXAsyncCopy` could be tested.
This keeps the test aligned with the rest of the TIRX PTX async tests
and with the TVMScript printer output, without adding extra legacy
aliases to the public script namespace.
This PR fixes 11 test failures in `tests/python/s_tir/transform/`
introduced as side effects of the TIRx bringup (#19581 / 859498dc01), in
three independent commits.
### 1. LowerOpaqueBlock: update expected IR for buffer metadata
annotations
`LowerOpaqueBlock` now emits `buffer_allocated_addr` and
`buffer_data_alignment` annotations on lowered allocations (intentional
in #19581: the annotations are consumed downstream by `codegen_cuda.cc`
/ `codegen_trn.cc`; the alignment value 64 comes from
`kAllocAlignment`). The tests' expected IR predates this, so
`assert_structural_equal` failed on the missing annotations.
Fix: update the expected IR in
`test_s_tir_transform_lower_opaque_block.py` to carry the annotations
(`T.decl_buffer(...)` → `T.alloc_buffer(..., annotations={...})`). Fixes
6 tests.
### 2. DefaultGPUSchedule: parse scalar-block test in s_tir mode
#19581 added a well-formedness rule rejecting `SBlockRealize` in
`tirx=True` mode, which is correct — sblocks are s_tir-mode constructs.
The hand-written `Before`/`Expected` modules in
`test_scalar_block_no_loops` were the only ones in the file still using
plain `T.prim_func`, so they failed at parse time before the pass under
test even ran.
Fix: parse both modules with `T.prim_func(s_tir=True)`, consistent with
every other test in the file. Fixes 1 test.
### 3. InjectPermutedLayout: match legacy PTX intrinsics by canonical
name
#19581 registers device intrinsics under two Op identities: a flat
builtin name (returned by `builtin::xxx()` in C++) and a canonical
dotted name (e.g. `tirx.ptx.ldmatrix_legacy`, produced when TVMScript /
tensor intrinsics are parsed). `InjectPermutedLayout` only compared with
`same_as(builtin::...)`, so it silently skipped rewriting the swizzled
shared-memory offsets of parsed legacy-form calls, leaving the expected
swizzle index expressions unmatched.
Fix: match `ptx_ldmatrix_legacy` / `mma_store_legacy` by both the
builtin Op and the canonical name via an `IsOp` helper, following the
existing pattern in `lower_warp_memory.cc` and `codegen_cuda.cc`. Only
the legacy intrinsic forms fold shared-memory access into
`tvm_access_ptr` + offset; non-legacy forms address shared memory
through `BufferLoad` and are already handled by the BufferLoad visitor,
so the unreachable `InternalError` throw is replaced by a pass-through.
(`mma_store_legacy` has no dotted alias, hence the asymmetric name
strings.) Fixes 4 tests.
This PR fixes two independent test-isolation issues that only surface
when certain test files run together in one pytest session.
1. Fix `_DialectRedirectFinder` duplicate module execution
`_DialectRedirectFinder.find_spec` used to pre-register the redirect
target module under the legacy alias name before returning the alias
spec.
This interacts badly with CPython import logic: when the requested
module name is already in `sys.modules`, CPython may ignore the returned
alias spec and reuse the target module's original spec instead. As a
result, the target source can be executed again under the canonical
module name, creating a duplicate module object.
This caused patches on aliased modules to silently miss the module
object used by existing code. For example,
`unittest.mock.patch("tvm.tirx.script.builder.buffer_store")` patched
the duplicate module, while the tirx parser still held references to the
original one, so `test_scalar_assign_error_not_swallowed` failed with
`DID NOT RAISE`.
This pr removes the pre-registration and let the import machinery
register the alias normally. Since the alias spec is now used,
`_AliasLoader.exec_module` also restores the canonical `__spec__` and
`__loader__` to avoid stale alias metadata on the loaded module.
2. Remove unused `tirx.intrin_test` op registration
`test_s_tir_transform_lower_match_buffer.py` registered a dummy op:
```python
tvm.ir.register_op_attr("tirx.intrin_test", "")
```
This was a leftover from the old TVMScript parser and is no longer
needed. The modern tirx parser eagerly evaluates `intrin_test(...)`
calls into `T.evaluate(0)`, so this op never appears in parsed IR.
The only remaining effect was adding a category-less `tirx.intrin_test`
entry to the global op registry, which could break
`test_registered_tirx_ops_have_exactly_one_category` depending on test
import order.
This pr removes the unused registration.
This PR adds an autoload mechanism for out-of-tree backends, simplifies
TVM's Python library loading, and removes `TVMError` in favor of native
Python errors.
## Autoload out-of-tree backends
Out-of-tree packages can register an autoload callable under the
`tvm.backends` entry-point group (mirroring torch's device-backend
autoload). At `import tvm` startup each entry point is discovered and
its callable invoked once, after the core runtime and the `tvm`
namespace are fully initialized, so an extension can register
ops/targets/funcs or load extra libraries.
```toml
[project.entry-points."tvm.backends"]
tvm_foo = "tvm_foo:_autoload"
```
A failing extension is caught and surfaced via `warnings.warn` so it
cannot break `import tvm`. Autoload can be disabled with
`TVM_DEVICE_BACKEND_AUTOLOAD=0`.
## Simplify library loading
The library-loading path in `base.py` is consolidated around a single
`_LOADED_LIBS` dict (basename to ctypes handle) so downstream and
autoloaded extensions can skip already-loaded libraries; the per-backend
runtime DSO list is folded into `load_backend_libs`. Accumulated cruft
is removed: the Python-3.9 check, the readline shim, the `_FFI_MODE`
ctypes check, the `base.__version__` re-export, and `py_str` (call sites
inline `.decode("utf-8")`).
## Remove TVMError in favor of native Python errors
`TVMError` added a layer atop `RuntimeError` that downstream code had to
import and learn. It is removed; the registered FFI error kinds
(`InternalError`, `RPCError`, `OpError`, `DiagnosticError`,
`ScheduleError`) now subclass `RuntimeError` directly while staying
registered, so the FFI keeps throwing the right kinds. All `TVMError`
imports, `except`/`raise`/`isinstance` uses, and
`pytest.raises(tvm.TVMError)` sites move to the `RuntimeError` builtin.
Fix the s_tir tests broken or left stale by two upstream changes.
* test_meta_schedule_space_cuda.py (cap, dil, gmm, t2d, nrm, sfm, cbr,
tbg) and test_meta_schedule_space_cuda_async.py (c2d): #18927 expanded
DefaultCUDA unroll_max_steps from {0, 16, 64, 512, 1024} to {0, 16, 32,
64, 128, 256, 512, 1024} without updating the recorded SampleCategorical
decisions. Remap the indices (2->3, 3->6, 4->7) so each test keeps
sampling the same unroll value; every sketch was re-verified by
replaying the trace and structurally comparing against the expected
module.
* T.let migration: since #19581 the TIRx parser treats `v: T.int32 =
expr` as a mutable local-scalar buffer instead of an immutable bind,
which is now spelled `v: T.let[T.int32] = expr` (a Bind node, the same
form te.create_prim_func emits). Tests whose intent is a bind are
migrated to the new spelling: reduction combiner temporaries
(add_rfactor, lower_cross_thread_reduction) and let-dependent passes
(compact_buffer_region, hoist_expression, remove_undef).
* Also convert reduction temporaries in still-green tests
(cross_thread_reduction rule, compute_inline, schedule utilities,
parallel_vectorize_unroll postproc, dlight general reduction, relax
cuda_graph) so the hand-written workloads match the canonical Bind form
instead of feeding rules a mutable-scalar body.
## Summary
Python callers should reach the canonical tvm-ffi structural helpers
directly instead of going through a TVM-side redirect layer. This makes
the public tvm.ir bindings exact aliases of the tvm_ffi APIs and exposes
get_first_structural_mismatch from tvm.ir.
Main changes:
- Import structural_equal, get_first_structural_mismatch, and
structural_hash directly from tvm_ffi
- Remove the pure wrappers from tvm.ir.base while keeping
assert_structural_equal's TVM-specific formatting
- Update mismatch tests and add identity coverage for the direct
bindings
The host/device split flow already runs device-region annotation,
host/device function extraction, and device-kernel launch lowering as
one consecutive pipeline. Keeping those stages exposed as separate
public passes makes the API surface larger than the actual execution
model and leaves the stage dependencies spread across multiple files.
This change makes `tirx.transform.SplitHostDevice` the single public
entry point for that flow, while preserving the existing stage order
internally.
Changes:
- Merge the annotation, splitting, and kernel-launch lowering
implementations into `src/tirx/transform/split_host_device.cc` as
private sections.
- Remove the old public C++ declarations, FFI registrations, and Python
wrappers for `AnnotateDeviceRegions` and `LowerDeviceKernelLaunch`.
- Replace pipeline call sites that previously invoked the three-stage
sequence with one `SplitHostDevice()` call.
- Update TIRx and S-TIR tests to exercise the consolidated pass and the
reduced public API surface.
## Summary
Lifts 10 host-toolchain / CLI / process / utility modules from
`python/tvm/contrib/` to a new `python/tvm/support/` package, and
deletes two dead contrib shims.
`tvm.support` is the home for Python helpers that integrate TVM with
external CLIs and host-side tools — compilers, archivers, subprocess
pools, and build-info queries. These are load-bearing internal pieces
that TVM's compile/link/run paths depend on. `tvm.contrib` is reserved
for optional vendor SDK integrations and experimental features. The
distinction is documented in the `tvm.support` package docstring.
Moved (one commit each):
- `tvm.contrib.cc` → `tvm.support.cc`
- `tvm.contrib.nvcc` → `tvm.support.nvcc`
- `tvm.contrib.rocm` → `tvm.support.rocm`
- `tvm.contrib.ndk` → `tvm.support.ndk`
- `tvm.contrib.xcode` → `tvm.support.xcode`
- `tvm.contrib.clang` → `tvm.support.clang`
- `tvm.contrib.emcc` → `tvm.support.emcc`
- `tvm.contrib.popen_pool` → `tvm.support.popen_pool`
- `tvm.contrib.utils` → `tvm.support.utils`
- `tvm.contrib.tar` → `tvm.support.tar`
Deleted:
- `tvm.contrib.spirv` — single `optimize()` wrapping `spirv-opt`; zero
importers.
- `tvm.contrib.rpc` — self-deprecation shim with "removed in 0.5"
banner; honoring it.
Package conversion:
- `python/tvm/support.py` → `python/tvm/support/__init__.py` with
inclusion-rule docstring.
- `libinfo()` extracted into `python/tvm/support/libinfo.py`.
- `FrontendTestModule` dropped (audit confirmed zero callers outside its
own definition).
## Compatibility
Hard break — no `tvm.contrib.<mod>` re-export shims. All callers updated
in this PR.
C++-side FFI registry keys (`tvm.contrib.nvcc.*`, etc.) are unchanged —
only the Python module path moves. Renaming the FFI keys is a separate
follow-up.
## Summary
These three passes are logically a single host/device split step;
having intermediaries between them obscures the model and blocks
folding them into one pass. This PR moves each intermediary to the
position its actual ordering constraint allows, so that
`AnnotateDeviceRegions`, `SplitHostDevice`, and
`LowerDeviceKernelLaunch` run consecutively in every pipeline.
## Rationale
- `MergeSharedMemoryAllocations` moves **before**
`AnnotateDeviceRegions`
(the only legal position: `LowerDeviceKernelLaunch` requires at most
one dyn-shmem allocation per kernel, so Merge cannot move past Lower).
- `MakePackedAPI` moves **after** `LowerDeviceKernelLaunch` (Lower's
`kCallingConv = kDeviceKernelLaunch` flag causes `MakePackedAPI` to
correctly skip device kernels; the host body's lowered
`tvm_call_packed` is transparent to `MakePackedAPI`'s subroutine
rewriter).
- `FP8StorageLegalize` / `BF16StorageLegalize` move **after**
`MakePackedAPI` (their `buffer_map.size()==0` ICHECK requires
`MakePackedAPI` to have cleared the map).
Prereq for Phase 2: collapsing the three consecutive passes into a
single `tirx.transform.SplitHostDevice` with three commented regions.
## Test plan
- [x] tests/python/tirx-transform/ target-pass unit tests (25 pass)
- [x]
tests/python/s_tir/transform/test_merge_dynamic_shared_memory_allocations.py
(5 pass)
- [x] tests/python/tirx-transform/test_tir_transform_fp8_legalize.py /
test_tir_transform_bf16_legalize.py (13 pass)
- [x] tests/python/codegen/test_target_codegen_c_host.py /
test_target_codegen_device.py (6 pass including
test_subroutine_call — verifies Risk #2)
- [x] pre-commit run --all-files clean
- [ ] CI: lint / Windows / MacOS
## Summary
This PR cleans up technical debt in the TIR simplification machinery via
two commits:
**Commit 1: Phase out ControlFlowGraph and NarrowPredicateExpression**
- Remove `ControlFlowGraph` (~2360 lines) from `src/tirx/analysis/` —
used only in
non-default config paths that are no longer maintained
- Remove `NarrowPredicateExpression` from `src/arith/` — sole non-test
caller was `ControlFlowGraph`
- Remove gated config fields `propagate_knowns_to_prove_conditional` and
`propagate_knowns_to_simplify_expressions` from `SimplifyConfig`
- Remove `use_dataflow_analysis` from `RemoveNoOpConfig`
- Delete the associated test files and test cases that tested the
now-removed paths
- ~3800 lines deleted
**Commit 2: Rename Simplify → StmtSimplify**
- Rename `src/tirx/transform/simplify.{h,cc}` → `stmt_simplify.{h,cc}`
- Rename C++ identifiers: `Simplify` → `StmtSimplify`, `SimplifyConfig`
→ `StmtSimplifyConfig`
- Rename FFI keys: `"tirx.Simplify"` → `"tirx.StmtSimplify"`,
`"tirx.transform.Simplify"` → `"tirx.transform.StmtSimplify"`
- Update Python wrappers and all call sites (~40 files)
- Clarifies that this pass operates on statements (distinct from
expression-level `arith::Analyzer::Simplify()`)
## Test plan
- [x] `tests/python/tirx-transform/test_tir_transform_simplify.py` — 52
tests pass
- [x] `tests/python/tirx-transform/test_tir_transform_remove_no_op.py` —
18 pass, 5 xfail
- [x] `tests/python/arith/` — full arith test suite passes
- [x] `tests/python/tirx-transform/` — full suite: 315 passed, 8
xfailed, 1 xpassed (pre-existing vectorize failure unrelated to this
change)
- [x] `pre-commit run --all-files` — all hooks pass
## 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.
## Problem
Closes#17873.
`DefaultGPUSchedule` crashes when a PrimFunc body is a bare
`SBlockRealize` (a fully-scalar op with no enclosing loops and no iter
vars):
```
ValueError: Check failed: (sref->parent != nullptr) is false:
Cannot add loops on top of the root block
```
Minimal repro (TVMScript decorators are omitted in this snippet to
satisfy the PR-body lint; the regression test uses the regular
`T.prim_func` form):
```
ir_module:
prim_func main(a: Buffer((), "float32"),
b: Buffer((), "float32"),
c: Buffer((), "float32")):
func_attr({"target": target("nvidia/geforce-rtx-3080")})
with sblock("scalar_add"):
c[()] = a[()] + b[()]
s_tir.transform.DefaultGPUSchedule()(M) # crashes
```
## Root Cause
The realized `scalar_add` block is itself the prim_func body's root
sref — it has no parent stmt to mutate. `ThreadBind`
(`src/s_tir/transform/default_gpu_schedule.cc`) reaches the
`loops.empty()` branch and calls `sch->AddUnitLoop(block)`, which fails
the `sref->parent != nullptr` check in `s_tir::AddUnitLoop`
(`src/s_tir/schedule/primitive/loop_transformation.cc:1166`).
The schedule infrastructure additionally requires the prim_func body
to be an `SBlockRealize` whose block is the function's root
(`GetRootPrimFunc` in `src/s_tir/schedule/analysis/analysis.cc:53`),
so the body cannot simply be wrapped in a top-level `For`.
## Fix
Before constructing the schedule, rewrite GPU-bound PrimFuncs whose
body is a bare-leaf `SBlockRealize` so the realized block is no longer
the root. The wrap conditions are intentionally narrow:
1. `func->body` is `SBlockRealize`,
2. the realized block has empty `iter_vars`, and
3. the block's body is not `For` or `SBlockRealize` (i.e. it is a leaf
computation, not the well-formed implicit root that wraps a loop
nest produced by the rest of the pipeline).
When all three hold, the body becomes:
```
SBlockRealize(
block=SBlock("root", body=
For(u, 0, 1, kSerial,
SBlockRealize(iter_values=[u],
block=<original block, iter_vars=[IterVar(0..1, vu, kDataPar)]>))))
```
The synthesised 1-extent data-parallel iter keeps
`iter_values.size() == iter_vars.size()` for downstream checks, and the
new For loop gives `ThreadBind` a real loop to bind to `blockIdx.x` /
`threadIdx.x`. Already-scheduled functions and host-only PrimFuncs are
skipped via the existing `IsScheduledOnGPU` / `kIsScheduled` gating.
## Testing
```
pytest tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
```
10 passed (9 existing + 1 new `test_scalar_block_no_loops`). End-to-end
compile + execute on RTX 3080 (sm_86): the scalar repro returns the
expected `2.0 + 3.0 = 5.0`.
This change just keep stride terms order the same with fused loop order
in `fuse` primitive. In symbolic circumstances, previous form suffer
from simplification issues and would make the expression tree much
complex in following lowering steps.
Take [M, N] tiling as an example, the previous binding form after
```python
i, j = sch.get_loops(block_b)
i0, i1 = sch.split(i, factors=[None, 64])
j0, j1 = sch.split(j, factors=[None, 16])
sch.reorder(i0, j0, i1, j1)
sch.fuse(i0, j0)
```
would be like (i_0_j_0_fused in `[0, ceildiv(M, 64) * ceildiv(N, 16)]`
```
vi = T.axis.spatial(M, i_0_j_0_fused % ((N + 15) // 16 * ((M + 63) // 64)) // ((N + 15) // 16) * 64 + i_1)
```
instead of more simple version
```
vi = T.axis.spatial(M, i_0_j_0_fused // ((N + 15) // 16) * 64 + i_1)
```
This is because unfortunately we do not know `ceildiv(N, 16) *
ceildiv(M, 64) == ceildiv(M, 64) * ceildiv(N, 16)` in rule based
simplifications. And then certain analysis (for example, region
estimation) may fail to give concise estimations, due to complex dynamic
expression trees.
Co-authored-by: baoxinqi <bao.xinqi@intellif.com>
## Summary
This adds gating logic on top of #17699 to support optional subgroup
shuffle
primitives based on a compile-time flag.
## Problem
The PR #17699 always generates subgroup shuffle ops when targeting
WebGPU.
However, not all WebGPU devices support subgroups. We need a way to:
- Default to shared memory reductions (universally compatible)
- Optionally enable subgroup shuffles for devices that support them
## Solution
Implement gating via TVM target parameter:
- Default `thread_warp_size=1` disables warp reductions (uses shared
memory + barriers)
- Add target parser `UpdateWebGPUAttrs()` that sets
`thread_warp_size=32` when `supports_subgroups=true`
- Add `--enable-subgroups` CLI flag in mlc-llm to surface the option to
users
The gating happens at the reduction path selection level
(`IsWarpReduction()` in
`lower_thread_allreduce.cc`), ensuring subgroup ops are never generated
unless explicitly enabled.
## Testing
Tested with Llama-3.2-1B-q4f16_1. Baseline (no flag) uses shared memory
reductions;
with flag, generates subgroupShuffle* ops.
Both the generated WGSLs here:
https://gist.github.com/ksgr5566/301664a5dda3e46f44092be4d09b2d4f
Benchmarking:
https://gist.github.com/ksgr5566/c9bd5bc5aadba999ec2f2c38eb0c49b3
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
- Remove `body` field from `AllocBufferNode` and `DeclBufferNode`,
making them flat statements consistent with `Bind`
- Buffer scope extends to end of enclosing scope via flat `SeqStmt`
semantics
- 60 files changed across core IR, codegen backends, transforms, script
IR builder, and tests
## Test plan
- All existing test suites pass (tir-transform, tir-base, tvmscript,
s_tir, codegen, C++)
## Summary
Rename `LetStmtNode`/`LetStmt` to `BindNode`/`Bind` and remove the
`body` field.
The variable defined by `Bind(var, value)` is now visible in all
subsequent
statements within the same enclosing scope, rather than being scoped to
a nested body.
This flattens deeply nested let-chains into sequential
`SeqStmt([Bind(...), Bind(...), ...])`,
making the IR easier to read, transform, and analyze.
## Key Changes
- **New `BindNode`**: `{var, value}` — no body field. Variable scope is
the enclosing
statement's body (For, IfThenElse, AllocBuffer, etc.)
- **ScopeStack pattern**: Passes that need scope-aware cleanup
(ConvertSSA, CSE,
tir_visitor_with_path) use `ScopeStack` instead of manual save/restore
or RAII wrappers
- **All passes migrated**: 89 files updated across codegen backends, TIR
transforms,
S-TIR transforms, analyses, TVMScript printer/parser/ir_builder
## 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 removes unused `from tvm import te` imports from 25 test files
across the codebase, continuing the ongoing TE → TVMScript migration
cleanup.
Changes:
- Remove unused `from tvm import te` from 24 test files in
tir-transform/, s_tir/transform/, codegen/, arith/,
all-platform-minimal-test/, and testing/
- Replace `te.var("x")` with `tvm.tir.Var("x", "int32")` in
`test_s_tir_transform_decorate_device_scope.py` (the only file where te
was actually used)
- Clean up stale `tvm.tir.ir_builder` comment references in
`test_tir_transform_convert_ssa.py`
This PR migrates the s_tir related transform passes into s_tir namespace
instead. This set of changes can minimize the overall tir namespace to
make it more focused.
Tianqi Chen