73 Commits

Author SHA1 Message Date
Tianqi Chen 1a4e037bbb [CI] Bump tvm-ffi with compatible Python wrappers (#20032)
## Summary

- bump tvm-ffi and include the device definition where its `DLDevice`
traits are instantiated
- keep only the required Tensor wrapper layout fix and register
`ir.Type` before reflected `Expr` fields can materialize a fallback
wrapper
- preserve `BaseFunc.with_attr` callers by moving only method-private
results, never the canonical `self` wrapper

## Rationale

The tvm-ffi lifetime update requires a replacement wrapper to fit the
layout already registered for the same type index. `runtime.Tensor`
replaces the core `ffi.Tensor` wrapper, so it must use empty slots. The
ordinary TVM mixins are first-registered with their concrete descendants
and may safely retain normal Python dictionaries; the additional mixin
and explicit-dictionary slot changes are not required.

Object tying also means `BaseFuncCopy(self)` may return `self`. Passing
that wrapper through `_move()` invalidates the caller. The first update
now passes the alias as an lvalue, forcing native copy-on-write to
create a private result. Only later dictionary updates move a result
that is not `self` and has not escaped the method.

## Validation

- built an exact CPython 3.12 wheel from tvm-ffi `21e30c3b1d` and
rebuilt TVM against it
- direct Type/function/detach regressions: 3 passed
- complete IR plus focused Relax coverage: 111 passed
- prior Relax failure set: 157 passed, 9 skipped
- runtime probe for `relax.Function`, `relax.ExternFunc`, and
`tirx.PrimFunc`: original wrappers preserved; single- and
multi-attribute results distinct and valid
- all touched-file pre-commit hooks passed

---------

Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
2026-07-20 14:26:46 +08:00
Tianqi Chen 80648af29f [REFACTOR][TIR] Remove buffer type and axis separators (#20019) 2026-07-17 17:08:13 +08:00
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.
2026-07-16 17:34:50 -04:00
Tianqi Chen c717c5b217 [IR][Relax][TIRx] Unify Var identity (#20004) 2026-07-15 17:12:40 +08:00
Tianqi Chen e479a5dbe7 [RUNTIME][PYTHON] Add explicit Target device conversion (#20005)
## 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
2026-07-15 05:34:21 +08:00
Nanmur a540271e61 [DLight][CUDA] Fix undefined TX in GEMV broadcast epilogue (#19970)
This PR fixes an undefined `TX` reference in the DLight GPU GEMV
inner-reduction schedule.

The broadcast epilogue path splits fused epilogue loops and binds the
inner loop to `threadIdx.x`, but it referenced `TX`, which is not
defined in the surrounding scope or passed into the helper. This PR uses
the existing `TR` tile factor, which is already used for the
`threadIdx.x` direction in this schedule.

A regression test is added to cover the GEMV broadcast epilogue path.

Fixes #19969

Tests:
- `python -m pytest tests/python/s_tir/dlight/test_gpu_gemv.py -q`
- `python -m pytest tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
-q`
- `python -m pytest tests/python/s_tir/dlight/test_gpu_reduction.py -q`

Note: The branch is based on `apache/tvm:main`. Local main-branch test
execution on this Windows machine could not be completed because the
available compiled TVM library is from a v0.25 build and does not match
the latest Python sources; the same tests passed in the matching local
v0.25 environment before rebasing the patch to main.
2026-07-13 14:43:12 -04:00
Tianqi Chen e1e7ac9261 [ARITH] Scope interval constraints to mapped variables (#19963)
## 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.
2026-07-06 15:47:46 -04:00
Tianqi Chen bbfdab79d9 [CI] Repair Python test cleanup regressions (#19955)
## 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.
2026-07-06 16:29:52 +08:00
Tianqi Chen cfb98e938c [CI] Simplify Jenkins pytest execution (#19947)
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.
2026-07-05 09:59:51 -04:00
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.
2026-07-04 21:54:04 -04:00
Tianqi Chen 3452fd4ffa [TEST] Serialize local GPU execution under pytest-xdist (#19942)
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.
2026-07-04 17:49:45 -04:00
Tianqi Chen 99869414de [TIRX] Remove SizeVar in favor of contextual constraints (#19930)
## Rationale

`SizeVar` encodes nonnegativity in runtime subtype identity, which is
fragile under cloning and remapping. Symbolic integer values should use
one `Var` representation, with nonnegative facts recorded in the
analyzer at the use sites that establish them.

## Changes

- Remove `SizeVar` from the C++, Python, TE, TVMScript, FFI, visitor,
and serialization surfaces, and migrate callers to `Var`.
- Preserve the existing Relax constraint ownership model and use
`MarkGlobalNonNegValue` as the canonical path for global nonnegative
facts.
- Preserve `T.handle()` as the normal opaque-handle form. An optional
dtype constructs a typed pointer, with `T.handle("void")` reserved for
an explicit pointer-to-void.
2026-07-03 11:33:14 -04:00
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.
2026-07-01 18:55:33 -04:00
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
2026-06-24 21:31:47 -04:00
Tianqi Chen 1bb5cf6102 [REFACTOR][IR] Unify StructInfo and Type (#19853)
## Summary

- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
2026-06-21 10:12:12 -04:00
Shushi Hong 2a77aaaadd [TIRx] Phase out flat device-intrinsic op aliases (#19838)
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.
2026-06-19 12:25:41 -04:00
Shushi Hong 71466bb737 [Tests] Migrate tvm.testing.parameters() to pytest.mark.parametrize (#19803)
This pr phases out the custom `tvm.testing.parameters()` helper in favor
of native `pytest.mark.parametrize`. `parameters()` itself is left in
place for now and removed in a follow-up, together with updating the
framework self-test
(`tests/python/testing/test_tvm_testing_features.py`) that exercises it.

Migration rules
- A group consumed only by test functions becomes
`pytest.mark.parametrize`.
- Single-name groups are unwrapped from 1-tuples to bare values.
- A group shared by multiple tests uses a module-level named list; a
test that uses only a subset of a group's names is parametrized only on
the names in its signature.
- `pytest.mark.parametrize` is stacked above the existing, unrelated
`tvm.testing.parametrize_targets(...)`, which is kept as-is.

Per-file pytest collection case counts are unchanged, except the two
intentional changes below.

Behavior changes (intentional)
- tests/python/relax/test_training_optimizer_numeric.py: the names `lr`
and `weight_decay` were rebound across three `parameters()` groups, so
`test_sgd` and `test_momentum_sgd` silently used the *adam* group's
`lr`/`weight_decay` (and `test_momentum_sgd` cross-producted with it:
2/6/2 = 10 cases). Native parametrize gives each test its own co-located
group: 2/3/2 = 7 cases. This fixes that latent rebinding bug; the case
count drops 10 -> 7 and `test_momentum_sgd` now exercises its own
`weight_decay` values.
- tests/python/target/test_arm_target.py: its `parameters()` group was
orphaned (no test consumed those names) — removed the dead definition.

Note: for tests that also use `tvm.testing.parametrize_targets`, the
generated test ids reorder the target (e.g. `test_unary[abs-True-llvm]`
-> `test_unary[llvm-abs-True]`); values and case counts are unchanged.
2026-06-16 21:53:40 -04:00
Shushi Hong e4da848e57 [Tests] Modernize test gating (#19777)
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`
2026-06-15 18:50:57 -04:00
Bohan Hou bb6f8aec55 [TIRx] Post-bringup follow-ups: op-dispatch, namespaces, launch bounds, gemm-async, backend reorg (#19757)
This PR batches several post-bringup TIRx follow-ups, rebased onto
current `main`.

### Changes
- **op-dispatch**: per-call exec scope via `Tx.<scope>.op`; remove
`ExecScopeStmt`
- **namespaces**: split TIRx op namespaces; remove tile-primitive kind
attrs
- **codegen**: support explicit CUDA launch bounds
- **gemm-async**: support contiguous-axis (K-major) operand slicing
- **backend reorg**: move in-tree GPU backends out of core into
`src/backend/<target>/` and `python/tvm/backend/<target>/`
(codegen/runtime/op), with the corresponding `CMakeLists.txt` /
`cmake/modules` and include-path updates

### Testing
- Builds with `USE_CUDA=ON` / `USE_LLVM=ON`
- The TIRx Python test suite (`tests/python/tirx/`) passes locally
2026-06-13 21:12:40 -04:00
Shushi Hong 126f1ba7ff [S-TIR][CUDA] Fix legacy predicated cp.async zero fill (#19741)
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.
2026-06-13 07:30:19 -04:00
Shushi Hong 6e6627ec6a [S-TIR][Tests] Mark test_cp_async_in_if_then_else as xfail (#19751)
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.
2026-06-12 19:06:58 -04:00
Shushi Hong 30bf568dd0 [Tests] Check WebGPU volatile allreduce annotation structurally (#19740)
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.
2026-06-12 00:38:19 -04:00
Shushi Hong 768d38bf07 [TIRx] Use canonical PTX async script API in s_tir test (#19739)
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.
2026-06-12 00:38:01 -04:00
Shushi Hong c9a77d6712 [S-TIR][Tests] Fix transform test failures after TIRx bringup (#19735)
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.
2026-06-11 17:34:05 -04:00
Shushi Hong 07e60343bc [Script][Tests] Fix dialect redirect module re-execution and stray category-less tirx.intrin_test op (#19731)
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.
2026-06-11 14:59:34 -04:00
Tianqi Chen 96cba60464 [PYTHON] Autoload backends; simplify library loading; remove TVMError for native errors (#19727)
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.
2026-06-11 13:50:38 -04:00
Shushi Hong 67b0c6cc5f [Tests][S-TIR] Fix stale MetaSchedule sketch expectations and migrate let binds to T.let (#19729)
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.
2026-06-11 07:19:00 -04:00
Shushi Hong def37e30bd [Tests] Skip test modules cleanly when optional deps are missing (#19704)
Validating the apache-tvm wheel in a minimal environment (no torch,
scipy, cloudpickle, or tornado installed) produced 33 pytest collection
errors from module-level imports of optional packages. Add
pytest.importorskip guards so these modules are reported as skipped
instead of erroring during collection.

Indirect import chains guarded:
- tvm.topi.testing imports scipy
- tvm.s_tir.meta_schedule.testing.local_rpc (tvm.rpc.tracker) requires
tornado
- tvm.s_tir.dlight.benchmark imports cloudpickle

Also remove a stray pre-license-header `import pytest` in
test_runtime_builtin_paged_attention_kv_cache_flashinfer.py.
2026-06-10 01:20:50 -04:00
Shushi Hong 219f1d83cf [Refactor][Meta-schedule] Remove meta-schedule as_string mechanism in favor of default representation (#19709)
Python-side meta-schedule classes (`PyCostModel`, `PyFeatureExtractor`,
`PyMeasureCallback`, `PyScheduleRule`, `PyMutator`, `PyPostproc`)
carried an `f_as_string` callback whose only purpose was to produce a
repr-style string (`s_tir.meta_schedule.<SubclassName>(0x...)`) for
`str(...)`.

This mechanism stopped working after #19461 migrated `ReprPrinter` to
the tvm-ffi `__ffi_repr__` mechanism and intentionally removed the
per-type `set_dispatch<Py*Node>` hooks that called back into
`f_as_string`, which broke three `*_as_string` tests:
-
`test_meta_schedule_cost_model.py::test_meta_schedule_cost_model_as_string`
-
`test_meta_schedule_feature_extractor.py::test_meta_schedule_feature_extractor_as_string`
-
`test_meta_schedule_measure_callback.py::test_meta_schedule_measure_callback_as_string`

Rather than restoring the old behavior, this PR removes the mechanism
entirely: the string it produced is just a repr, and tvm-ffi reflection
already provides an auto-generated default repr for every object.
Keeping a dedicated Python → FFI → C++ callback chain alive only to
reproduce that is not worth the complexity.
2026-06-09 21:53:34 -04:00
Shushi Hong fda3220ead [Tests] Fix s_tir tests using removed T.block API in TIRx script (#19706)
Two test files under `tests/python/s_tir/meta_schedule/` use the TIRx
script dialect (`from tvm.script import tirx as T`) but still call the
old TIR block API, which does not exist in `tvm.tirx.script.
This pr fixes this
2026-06-09 18:58:08 -04:00
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`
2026-06-05 21:02:36 -04:00
Tianqi Chen 1240649257 [FFI][REFACTOR] Direct structural APIs to tvm-ffi (#19661)
## 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
2026-06-03 18:57:05 -04:00
Tianqi Chen dea2bf933e [REFACTOR][TIRX] Consolidate split host device stages (#19663)
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.
2026-06-03 15:57:57 -04:00
Tianqi Chen 6b4b866d65 [REFACTOR][ARITH] Phase out arith/scalable_expression; arith no longer proves over scalable vectors (#19638)
## Summary

Phase out `src/arith/scalable_expression.{h,cc}`. The arith layer no
longer attempts to prove anything about scalable vectors — proofs that
depended on `Target::Current()` are removed. Scalable vectors remain a
first-class concept; arith just doesn't reason about their lengths.

## Use-site summary

Only 16 call sites total across 7 symbols (9 live, 7 proof-related).

| Symbol | Live callers (kept) | Proof callers (deleted) | New home |
|---|---|---|---|
| `ExtractVscaleFactor` | 4 × `arith/rewrite_simplify.cc` + 2 ×
`tirx/ir/expr.cc` | — | file-local in each |
| `IsVScaleCall` | 1 × `tirx/op/op.cc` + 1 ×
`tirx/transform/vectorize_loop.cc` | — | inline at use sites |
| `ContainsVscaleCall` | 4 × `arith/rewrite_simplify.cc` + 1 ×
`s_tir/schedule/ir_comparator.cc` | — | inline at use sites |
| `TargetHasVLA` | 2 × `tirx/transform/vectorize_loop.cc` | analyzer.cc
+ const_int_bound.cc | local in vectorize_loop.cc |
| `GetVScaleValues` | 1 × `target/llvm/codegen_aarch64.cc` | analyzer.cc
+ const_int_bound.cc | inlined at codegen_aarch64 |
| `CanProveVscaleExpressionFromKnownValues` | — | analyzer.cc | DELETE |
| `SubstituteVScaleWithKnownValue` | — | internal only | DELETE |

## Changes (6 commits)

1. Move `ExtractVscaleFactor` to file-local anonymous-namespace helpers
in `rewrite_simplify.cc` and `tirx/ir/expr.cc`. Function is small;
per-file duplication is cleaner than a shared header.
2. Inline `IsVScaleCall` / `ContainsVscaleCall` / `TargetHasVLA` at call
sites (1-3 line predicates, anonymous-namespace per consumer `.cc`).
3. Drop the scalable-vector proof scaffolding from `arith/analyzer.cc`
(substitution-proof loop) and `arith/const_int_bound.cc` (vscale
branch). `vscale()` calls fall back to `Everything()` — no special bound
narrowing.
4. Delete `scalable_expression.{h,cc}`. Inline the `GetVScaleValues`
body at `codegen_aarch64.cc` (computes `max_val = vector_width / 8`
floor-rounded to a power of two for the LLVM `vscale_range` attribute).
5. Mark `pytest.mark.xfail` on 19 tests that relied on the deleted
substitution-proof loop.
6. `pre-commit` line-length cleanup.

## Compatibility / intentional regression

This is a hard break for any consumer of the deleted symbols. They were
already in a private header (`src/arith/scalable_expression.h`, not
under `include/`).

19 tests that proved vscale-bearing inequalities on SVE / RVV are
xfailed. The proofs were target-dependent and the new policy is that
arith does not attempt them.
2026-05-28 22:02:42 -04:00
Tianqi Chen 61ae85b9d1 [REFACTOR][PYTHON] Consolidate derived_object into tvm.ir.utils (#19630)
## Summary

`derived_object` was duplicated byte-for-byte across
`python/tvm/runtime/support.py` and
`python/tvm/s_tir/meta_schedule/utils.py`. The function is not a runtime
feature and is used outside meta_schedule (tvm.relax, tvm.tirx), so
neither location was the right home.

Move the single canonical definition into a new
`python/tvm/ir/utils.py`. `tvm.ir` loads before both `tvm.tirx` and
`tvm.s_tir`, so eager top-level imports work from every consumer without
load-order workarounds.

Rewrite all 25 caller imports. Keep the better-typed `cls: type[T] ->
type[T]` signature from the runtime-side copy. After this change
`runtime/support.py` is empty and is removed;
`meta_schedule/__init__.py` drops its now-dead re-export. No alias shims
are left behind — callers update imports directly.
2026-05-27 23:15:43 -04:00
Tianqi Chen ffea531107 [REFACTOR][PYTHON] Lift compiler/CLI/process modules from tvm.contrib to tvm.support (#19624)
## 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.
2026-05-27 15:31:12 -04:00
Tianqi Chen ec3171ab7a [REFACTOR][TIR] Tie AnnotateDeviceRegions/SplitHostDevice/LowerDeviceKernelLaunch together (#19605)
## 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
2026-05-26 22:10:54 -04:00
Tianqi Chen 0388fd0ce0 [REFACTOR][IR] Phase out class Integer and class Bool in Attrs and PassConfig (#19614)
## Summary

Now that the ffi container machinery (Array, Optional, Map, Variant)
accepts bare int64_t and bool, the Integer/Bool ObjectRef wrappers add
no value in attribute fields, pass-config options, function-attr flags,
and OpAttrMap registries — every reader paid an extra .IntValue() /
->value unbox per access for no information gain. This PR is the first
stage of phasing out class Integer and class Bool: migrate the bulk of
those sites at the field-declaration and call-site level. A follow-up
will rewrite the remaining IR-position `Integer(N)` / `Bool(b)`
constructors to `IntImm(...)` / `const_true()` / `const_false()` and
delete the two classes entirely.

- Relax Attrs fields and their container forms (`Array<Integer>` /
`Optional<Array<Integer>>` / `Optional<Integer>` / `Optional<Bool>`)
migrated to bare `int64_t` / `bool` (manipulate.h, nn.h, op.h,
statistical.h, script/builder/frame.h, target/virtual_device.h,
distributed/global_info.h, relax/expr.h).
- OpAttrMap registry (`set_attr<Bool>("FPurity", Bool(true))` ↔
`GetAttrMap<Bool>("FPurity")`) migrated to `bool` across ~38 files.
- PassContext config registrations + `GetConfig<Bool>` /
`GetConfig<Integer>` readers, and function-attr `GetAttr<Bool>` /
`GetAttr<Integer>` readers (~42 files), all migrated; `HasNonzeroAttr`
in `ir/attrs.h` dropped its `.IntValue()` unbox.
- Schedule decision arrays (SampleCategorical candidates, perfect-tile
factors, autobind thread_extents, multi-level-tiling levels) migrated to
`Array<int64_t>` / `Optional<int64_t>` — this is a virtual-signature
change on `ConcreteScheduleNode::SampleCategorical` and related methods,
acceptable per the phase-out intent.
- `Variant<Bool, Array<String>>` for
`LiftTransformParams.shared_transform` migrated to `Variant<bool,
Array<String>>`.
2026-05-26 18:50:45 -04:00
Tianqi Chen 02b130249c [REFACTOR][TIR][ARITH] Phase out ControlFlowGraph, NarrowPredicateExpression, and rename Simplify to StmtSimplify (#19604)
## 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
2026-05-26 15:33:40 -04:00
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.
2026-05-18 16:44:43 -07:00
ConvolutedDog e7a7447929 [Fix][CI]: remove astral-sh/setup-uv from lint workflow (#19554)
This PR fixes https://github.com/apache/tvm/issues/19552.

astral-sh/setup-uv is not on the ASF GitHub Enterprise action allowlist,
causing the Lint workflow to fail with "Startup failure" before any
pre-commit checks run. See
https://github.com/apache/tvm/actions/runs/25743684906 for the failed
reason.

This PR removes the uv setup and sync steps entirely; pre-commit/action
will install and manage pre-commit and all hook dependencies on its own.
This PR also corrected previous lint errors.

After the fix, the CI lint succeeded:
https://github.com/apache/tvm/actions/runs/25775499703/job/75707088129
2026-05-13 12:28:31 +08:00
ConvolutedDog b1918c74fd [Fix][Relax]: ONNX Clip NaN bounds and preserve input NaN (ORT parity) (#19535)
This PR fixes https://github.com/apache/tvm/issues/19533:
- Sanitize floating tensor min/max: replace NaN with +inf/-inf before
topi max/min so bounds match ONNX "unbounded" semantics where NaN bounds
default to no constraint.
- After clamping, preserve NaNs from the input tensor on floating
dtypes.
- Extend check_correctness with equal_nan for float outputs containing
NaN.
- Add parametrized Clip opset-13 tests for NaN min/max tensor bounds.
2026-05-12 19:06:06 +08:00
Soowon Jeong 446bd2dbf0 [BugFix][S-TIR] Wrap bare scalar bodies in DefaultGPUSchedule to avoid root-block crash (#19514)
## 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`.
2026-05-07 00:22:43 +08:00
Neo Chien 7ecf466e33 [S-TIR][Dlight] Add layered fall back strategy to handle missing attr max_shared_memory_per_block (#19453)
Hi Committers,

This PR is trying to fix issues
https://github.com/apache/tvm/issues/19419. Any suggestions would be
appreciated if you are available.

### Root Cause
- auto-detected CUDA might lacks `max_shared_memory_per_block` and it
would cause `KeyError`

### Solutions
- Add layered fall back strategy to handle missing attr
`max_shared_memory_per_block`

---------

Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
2026-04-30 00:36:11 +09:00
wrongtest c01f898ed0 [TIR] Update symbolic index term order in loop fusion (#18406)
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>
2026-04-26 15:44:21 -04:00
Neo Chien 0a0dd3162b [S-TIR][MetaSchedule] Make evolutionary search resilient to trace replay failures (#19438)
Hi Committers,

This PR is trying to fix issues
https://github.com/apache/tvm/issues/17934. Any suggestions would be
appreciated if you are available.

### Root Cause
- During `EvolutionarySearch` candidate generation,
`trace->ApplyToSchedule(...)` could throw `ScheduleError`.
- The exception was propagated through parallel execution and aborted
tuning.
- Error handling was inconsistent between measured and unmeasured paths,
and failure visibility was limited.

### Solutions
- Catch trace replay failures in `ThreadedTraceApply::Apply` and return
`nullopt` instead of crashing.
- Add trace replay failure counting (`trace_fail_counter_`) and accessor
(`TraceFailCount()`).
- Align measured path `PickBestFromDatabase` with unmeasured behavior:
skip invalid candidates and continue.
- Add visible `WARNING` logs when trace replay failures occur (to avoid
silent failures).

---------

Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
2026-04-25 15:01:02 -04:00
Andrey Malyshev 0a79095b1d [S-TIR] Fix cache_read/cache_write region when inner block has T.whe… (#19406)
…re predicate

When the actual buffer access is gated by T.where on a nested (inner)
sblock, the outer block's own predicate is trivially true. Both
cache_write and cache_read were computing cache regions based only on
that outer predicate, producing allocations as large as the full loop
extent instead of the guarded region

  Fix:
- Add CollectNestedBlockPredicates(), a single helper parameterised by
BufferIndexType (kRead / kWrite) that walks the outer block's body,
finds nested sblocks accessing the target buffer, and AND-combines their
predicates after substituting iter-var bindings into the outer scope.
- Add extra_predicate parameter to RelaxBufferRegion() and AND it with
the block's own predicate before region relaxation.
- cache_write: pass the collected nested-write predicate to
RelaxBufferRegion so the cache allocation is tightened.
- cache_read (Case 2 — input buffer): when a non-trivial nested-read
predicate exists, relax the consumer block's declared read region under
that predicate; otherwise fall back to the original scope_block->reads
path (preserves non-int32 dtypes in extents).
2026-04-16 01:07:53 -04:00
Neo Chien 9d13fc04d2 [S-TIR] Fix Segfault when applying Parallel during TIR schedule rewriting (#19403)
Hi Commiters,

This PR is trying to fix issues
https://github.com/apache/tvm/issues/18424. Any suggestions would be
appreciated if you are available.

### Root Cause
Unsafe dynamic-shape dereferences in `AdjustParallelVectorize` The code
assumed IntImm for buffer shape / loop extent and dereferenced directly.
With dynamic shapes, as<IntImmNode>() can be null, which can segfault
before any try/catch handles it.

### Solution
Replaced unsafe `IntImm` assumptions with null checks and
GetLoopIntExtent(...); if contiguous analysis is not possible,
conservatively disables that path instead of dereferencing null.

---------

Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
2026-04-15 20:18:36 -04:00
Neo Chien 828880ebd5 [S-TIR] Fix the ScheduleError from missing guard in CheckInline (#19377)
Hi Commiters,

This PR is trying to fix issues
https://github.com/apache/tvm/issues/18380. Any suggestions would be
appreciated if you are available.

### Root Cause
The ScheduleError comes from a missing guard in CheckInline before
calling `GetScopRoot`

### Solution
Skip Inline if this is the root block

---------

Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
2026-04-11 01:22:35 -04:00
Soowon Jeong 5dda0b3a3a [DLight] Add CPU Reduction schedule rule for softmax-like operators (#19374)
This PR resolves part of #18569

## Description

This PR adds a DLight CPU schedule rule targeting reduction patterns
(softmax, layer norm, RMS norm) that previously had no CPU-specific
schedule.

Without this rule, LLVM auto-vectorization produces suboptimal code for
RVV targets — #18569 reports RVV softmax is **1.34x slower than
scalar**.
The root cause is:
- LLVM fully unrolls the 185-element reduction loop
- Generates harmful fixed-width vectors (`<8 x float>` = 256-bit) on a
  128-bit VLEN target
- No parallelization of the batch axis

### Schedule strategy

1. **Parallelize** leading spatial axes (batch dimension)
2. **Compute-at** all blocks under the spatial loop for data locality
3. **Vectorize** injective blocks (exp, delta, norm) on inner axis
4. **Split + unroll** reduction inner axis to VLEN-sized chunks

### Results (shape=(14, 185), float32, `fast_softmax`)

| Config | ASM Instructions | Vector Insns | LLVM IR Lines |
|--------|:---:|:---:|:---:|
| RV scalar (baseline) | 1,463 | 0 | 1,792 |
| RVV unscheduled (**the bug**) | 3,282 | 960 | 2,345 |
| RVV + this schedule | **1,111** | **105** | **1,338** |

- **66% fewer instructions** vs unscheduled RVV
- **24% fewer instructions** vs scalar baseline
- `fast_softmax` polynomial exp fully vectorizes into RVV
`vfsub`/`vfmul`/`vfmax` instructions with zero scalar `exp()` calls

### Limitations (follow-up work)

- Reduction blocks (max, sum) use split+unroll rather than vectorized
  partial reduction via `rfactor`, because TVM's `rfactor` primitive
  requires the reduction block to be the first child of its enclosing
  loop — incompatible with `compute_at` when multiple blocks share one
  spatial loop. A follow-up PR will register RVV reduction intrinsics
  (`vfredmax`/`vfredusum`) and use `tensorize` to vectorize reductions.

## Testing

```bash
pytest tests/python/s_tir/dlight/test_cpu_reduction.py -v
```

- 14 shape x operator applicability tests
- 2 TIR structure verification tests
- 4 RVV codegen quality tests (code size, exp vectorization, instruction
count)
- Existing `test_cpu_gemv.py` unaffected (10/10 pass)
2026-04-10 14:44:17 -04:00