- `LegalizeOps` skips grouped `conv1d/2d/3d` when channel size is
symbolic
- The only blocker is `topi.nn.conv`'s divisibility `assert`s, which
fail on symbolic `PrimExpr`; the grouped compute already handles
symbolic dims
## How
- Skip the divisibility check when the channel size is not a constant
int
- Remove the symbolic-channel guards from the conv legalize functions
## Why
Fixes#19925. `relax.shape_to_tensor` had no legalization, so it always
lowered to the host packed func `relax.run.shape_to_tensor`, producing a
CPU tensor regardless of the target device.
## How
- Register a legalization that emits the shape values as a `call_tir` TE
kernel, passing symbolic dims via `tir_vars`.
- Fall back to the packed func when the shape values are unknown
(`ShapeStructInfo` without values).
This PR fixes Relax bucketize lowering to pass the correct integer
output dtype to TOPI searchsorted.
Previously, both LegalizeOps and DispatchSortScan passed the input
tensor dtype as the output dtype. For float input tensors, this caused
TOPI searchsorted to receive a float output dtype, which later failed
during binary-search lowering because bucket indices must be integer
values.
This patch derives the output dtype from bucketize's out_int32
attribute:
- int32 when out_int32=True
- int64 otherwise
A numerical ExportedProgram frontend test is added to cover:
- right=False
- right=True
- out_int32=False
- out_int32=True
- float input values on bucket boundaries
Test:
python -m pytest
tests/python/relax/test_frontend_from_exported_program.py -k "bucketize"
-q
Remove the Relax-specific Id indirection and use Var/DataflowVar object
identity directly.
Type-changing rewrites now remap definitions, uses, and binding lookups
coherently while preserving reflection, serialization, and the
DataflowVar subtype. Existing tests are adjusted for the API change; no
new test files or test cases are added.
Validation: the compiler and C++ tests build successfully; focused C++
coverage passes 4/4; the affected existing Python matrices pass 549
tests with 2 expected xfails; source censuses, diff checks, and
applicable hooks pass.
## 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.
## 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.
This PR lets Relax expressions directly take `PrimExpr` values without
requiring the explicit `PrimValue` wrapper, continuing the Relax IR
unification work by removing Relax-specific leaf/base expression layers.
Summary:
- Remove `LeafExpr` / `LeafExprNode` and use direct expression-node
checks where needed.
- Converge Relax expression typing onto the shared IR `Expr` base.
- Remove the `PrimValue` node wrapper while keeping `relax.prim_value` /
`R.prim_value` as conversion helpers that return existing `PrimExpr`
values unchanged.
- Register direct `PrimExpr` handling through exact concrete node
dispatch, aligned with the `tirx` expression visitor list and excluding
arith iter-map intermediate nodes.
- Inline the private Python primitive conversion helper into public
`relax.prim_value`.
- Handle direct `PrimExpr` values in frontend scalar paths without
assuming a `.value` field on non-immediate expressions.
This PR introduces Relax AnyType as the primary top/base type spelling,
replacing the previous ObjectType naming for the type that represents
any Relax value.
Changes:
- Add AnyType/AnyTypeNode with relax.AnyType registration and keep
ObjectType/R.Object compatibility aliases.
- Update Relax type analysis, type visitors, opaque function defaults,
and script printer/parser handling to use AnyType/R.Any.
- Migrate affected Python/C++ call sites, docs, and focused tests to the
new spelling.
Validation:
- cmake --build build --parallel 16
- Focused Relax/TVMScript pytest: 704 passed, 1 xfailed
- pre_commit run --files <changed files>
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
## Summary
This PR adds focused Relax TFLite frontend coverage for the remaining
non-quantized builtin operators tracked by #18971:
- `SQUEEZE`
- `REVERSE_SEQUENCE`
- `UNPACK`
- `ZEROS_LIKE`
The tests manually build minimal TFLite flatbuffers and compare the
imported
Relax IR with `tvm.ir.assert_structural_equal`. This keeps the coverage
on the
frontend importer itself, without depending on TensorFlow converter
rewrites or
constant folding.
The PR also adds first-class Relax support for `reverse_sequence`.
TFLite
`REVERSE_SEQUENCE` was previously routed through:
```text
R.call_dps_packed("topi.reverse_sequence", ...)
```
That is not executable as a runtime packed call because
`topi.reverse_sequence`
is a TE compute and expects TE tensors during lowering. The frontend now
emits
`R.reverse_sequence`, and `LegalizeOps` lowers it through TOPI to TIR:
```text
TFLite REVERSE_SEQUENCE
-> R.reverse_sequence
-> LegalizeOps
-> topi.reverse_sequence
-> R.call_tir
```
## Design
### TFLite Operator Tests
The new TFLite tests use hand-built flatbuffers for small importer
fixtures:
- `SQUEEZE` checks axis handling and direct Relax `squeeze` lowering.
- `REVERSE_SEQUENCE` checks import to `R.reverse_sequence`, rejects the
old
`R.call_dps_packed("topi.reverse_sequence", ...)` path, compiles the
module,
and runs it with the VM.
- `UNPACK` checks multi-output lowering through Relax tuple output
handling.
- `ZEROS_LIKE` checks direct Relax zero-like tensor creation.
### Relax reverse_sequence Operator
The PR adds a public Relax operator:
```python
relax.op.reverse_sequence(data, seq_lengths, seq_axis=1, batch_axis=0)
```
The operator uses `ReverseSequenceAttrs` with `seq_axis` and
`batch_axis`.
Type inference preserves the input tensor's shape, dtype, and vdevice,
and
validates the statically known constraints:
- `data` must be a tensor.
- `seq_lengths` must be a 1-D tensor.
- `seq_lengths` dtype must be `int32` or `int64`.
- `seq_axis` and `batch_axis` must be in `[-ndim, ndim)` when the input
rank is
known.
- `seq_lengths.shape[0]` must match the batch-axis extent when both
shapes are
statically available.
The op is exported through Python as `relax.op.reverse_sequence` and
through
the script builder as `R.reverse_sequence`.
### Legalization
`relax.reverse_sequence` is registered in `LegalizeOps` and lowered with
`bb.call_te`:
```python
bb.call_te(
topi.reverse_sequence,
data,
seq_lengths,
seq_axis,
batch_axis,
primfunc_name_hint="reverse_sequence",
)
```
This produces `R.call_tir` in the legalized Relax module, keeping
runtime
execution on the normal TOPI/TIR path.
### TOPI Packed Registration
The Python TOPI wrapper already accepts `batch_axis`:
```python
topi.reverse_sequence(a, seq_lengths, seq_axis=1, batch_axis=0)
```
The C++ packed registration only forwarded the first three arguments, so
Python
calls that provided `batch_axis` would drop it before reaching the TOPI
compute.
The registration now forwards the fourth argument and keeps the old
three-argument call form compatible by defaulting `batch_axis=0`.
## Operator Support
| Operator | TFLite options | Relax lowering | Supported subset |
|---|---|---|---|
| `SQUEEZE` | `SqueezeOptions.SqueezeDims()` | `R.squeeze` | static
squeeze axes from TFLite options |
| `REVERSE_SEQUENCE` | `ReverseSequenceOptions.SeqDim()`, `BatchDim()` |
`R.reverse_sequence` legalized to TOPI/TIR | tensor input, 1-D
int32/int64 `seq_lengths`, valid `seq_axis` and `batch_axis` |
| `UNPACK` | `UnpackOptions.Axis()`, `Num()` | Relax tuple output |
static axis and output count from TFLite options |
| `ZEROS_LIKE` | none | `R.zeros_like` | tensor input |
## Not Included
- Quantized TFLite `REVERSE_SEQUENCE` support.
- A runtime DPS packed implementation for `topi.reverse_sequence`.
- Changes to TOPI compute semantics.
- ONNX `ReverseSequence` importer support.
## Tests
The tests cover both the TFLite frontend fixtures and the new Relax op:
| Test | Coverage |
|---|---|
| `test_squeeze` | imports TFLite `SQUEEZE` to Relax `squeeze` |
| `test_reverse_sequence` | imports TFLite `REVERSE_SEQUENCE` to
`R.reverse_sequence`, avoids the old TOPI DPS packed call, compiles, and
runs through VM |
| `test_unpack` | imports TFLite `UNPACK` as multi-output Relax tuple
handling |
| `test_zeros_like` | imports TFLite `ZEROS_LIKE` to Relax `zeros_like`
|
| `test_op_correctness` | `relax.op.reverse_sequence(...).op` resolves
to `relax.reverse_sequence` |
| `test_reverse_sequence_infer_ty` | static shape, unknown dtype,
unknown ndim, symbolic shape, and vdevice propagation |
| `test_reverse_sequence_infer_ty_wrong_inputs` | non-tensor
`seq_lengths`, wrong rank, wrong dtype, invalid axes, and static batch
mismatch |
| `test_reverse_sequence` in `test_transform_legalize_ops_manipulate.py`
| `LegalizeOps` emits `R.call_tir` and exercises `seq_axis=0,
batch_axis=1` |
Local validation:
```bash
python -m pytest tests/python/relax/test_op_manipulate.py \
-k reverse_sequence -q
python -m pytest tests/python/relax/test_transform_legalize_ops_manipulate.py \
-k reverse_sequence -q
python -m pytest --noconftest tests/python/relax/test_frontend_tflite.py \
-k "reverse_sequence or squeeze or unpack or zeros_like" -q
```
Result:
```text
cmake build: passed
py_compile: passed
ruff format --check: 9 files already formatted
ruff check: All checks passed
clang-format --dry-run --Werror: passed
pre-commit run --files: passed
test_op_manipulate.py -k reverse_sequence: 3 passed
test_transform_legalize_ops_manipulate.py -k reverse_sequence: 1 passed
test_frontend_tflite.py -k "reverse_sequence or squeeze or unpack or zeros_like": 4 passed
```
## References
- Issue #18971: TFLite non-quantized operator unit-test coverage
- TFLite `REVERSE_SEQUENCE` builtin semantics
## Why
relax.nn.conv{1,2,3}d_transpose with dilation > 1 silently bailed in
legalize and then crashed in VM codegen with an opaque error.
## How
- Lower dilation > 1 by zero-filling (dilating) the kernel, then reusing
the existing TOPI transposed-conv compute (1D/2D/3D).
- Unsupported non-NCHW layouts and out_layout != data_layout keep their
existing passthrough (left for downstream/BYOC codegen such as CLML),
unchanged.
- Add a 2D-dilation structural test.
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
## Related Issue
Closed#19695
## Why
Building any module containing relax.nn.dropout crashed in VM codegen
because the op had no real legalization, and the ONNX frontend could not
import it
## How
- Legalize nn.dropout to pass the input through with an all-ones mask,
matching its (output, mask) tuple result.
- Add and register a Dropout converter in the ONNX frontend.
- Add legalize structural and ONNX onnxruntime-parity tests.
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Summary:
- Remove the Relax-specific PrimType node/API and use canonical
ir.PrimType for dtype-only scalar types.
- Update parser, printer, analysis, op inference/legalization, and tests
to avoid value-bearing PrimType semantics.
- Preserve scalar values where needed by reading PrimValue expressions
directly instead of storing values in the type.
## Related Issue
closes#19690
## Why
ONNX AffineGrid carries an align_corners attribute, but the Relax op
ignored it and always produced the align_corners=1 grid.
## How
- Add align_corners field to AffineGridAttrs (mirrors GridSampleAttrs).
- Thread the flag through the op, legalize pass, and TOPI compute.
- Pass the ONNX attribute through in the frontend instead of dropping
it.
## 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
## 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
## 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>>`.
## 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.
`topi.scatter_elements` and `topi.scatter_nd` emit bare `T.parallel`
loops in their te.extern IRBuilder bodies which trips `VerifyMemory` on
CUDA targets:
RuntimeError: Memory verification failed
...
Did you forget to bind?
CPU (LLVM) is unaffected.
This fix makes the IRBuilder body in both `topi/scatter_elements.py` and
`topi/scatter.py` target-aware. When `Target.current()` is a GPU target
it emits thread bindings instead of `T.parallel`.
Fixes#19451.
The Relax ops `relax.isnan`, `relax.isinf`, and `relax.isfinite` are
registered in the op registry and emit valid IR, but lack FLegalize
entries so LegalizeOps() leaves them unlowered and relax.build() crashes
at codegen with:
InternalError: CodeGenVM cannot handle this intrinsic now:
Op(relax.isnan)
Same crash on `isinf` and `isfinite`, on both LLVM and CUDA targets.
The TOPI implementations already exist (python/tvm/topi/math.py). The
fix is three register_legalize calls following the
_call_topi_without_attr pattern used by the other unary ops in the same
file.
Fixes#19452.
This PR adds the ATAN2 operator to the Relax TFLite frontend.
Introduces relax.op.atan2 as a new binary elementwise primitive (TOPI
broadcast op, Relax registration, legalization to topi.atan2, script
parser support) and registers ATAN2 in the TFLite convert_map. It reuses
the TIR primitive tvm::atan2 so this PR is the higher-layer plumbing.
Validation:
python -m pytest tests/python/relax/test_op_binary.py
python -m pytest tests/python/relax/test_frontend_tflite.py -k binary
Addresses the ATAN2 item under #19412.
## Summary
Restructure TVMScript to be dialect-agnostic at the script-core layer
while letting each extension dialect (TIRX, Relax) own its own
per-dialect script subtree. IR is below script in the dependency
stack and is NOT a peer dialect — its script handlers stay in the
shared core.
This PR folds together two coupled refactors that were initially
opened as separate PRs (#19478 and the original #19479); they
share rename / relocation surface so they ship as one cohesive
change.
## What this PR does
### Per-dialect script subtree (originally #19479)
- Moves per-dialect printer + builder from
`src/script/{printer,ir_builder}/{tirx,relax}/` to
`src/{tirx,relax}/script/{printer,builder}/`.
- Tightens `src/script/*.cc` CMake glob to the dialect-free core.
- Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry
(`script.ir_builder.decl_function.<type-key>`); removes
cross-dialect includes from the shared core.
- Adds `tvm.script.register_dialect` API + `__getattr__` + a
`sys.meta_path` finder for Python-side dialect discovery.
In-tree dialects (tirx, relax) registered centrally in
`python/tvm/__init__.py`.
- Drops the obsolete static re-export shims at
`python/tvm/script/{parser,ir_builder}/{tirx,relax}/`.
### Dialect-agnostic printer config (originally #19478)
- Relocates `include/tvm/ir/script_printer.h` →
`include/tvm/script/printer/config.h` next to the rest of the
printer's public surface. The header is not IR-specific.
- Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for
clarity (the macro registers Script as the kRepr callback +
per-type vtable dispatch). Aligns with the `TVM_REGISTER_*`
family.
- Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`,
`relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor
of a generic `ffi::Map<String, Any> extra_config` keyed by
`"<dialect>.<knob>"`. Each call site reads via the templated
accessor `config->GetExtraConfig<T>("...", default)`.
- Promotes `std::string` config fields to `ffi::String`.
After this lands, the script-printer core knows nothing specific
about any dialect — new dialects plug in via the registry pattern
with zero core edits. Public Python API surface unchanged.
This PR fixes partial TFLite frontend support for dynamic `FILL` and
`SPLIT_V`.
Key changes:
- Allow `FILL` to accept runtime `dims` tensors by converting them with
`relax.op.tensor_to_shape` before calling `relax.op.full`.
- Allow `SPLIT_V` to accept runtime `size_splits` tensors by decomposing
the
op into `cumsum` and `dynamic_strided_slice`.
- Use the TFLite output tuple arity as the source of truth for dynamic
`SPLIT_V`, instead of relying on `size_splits` static shape information.
- Add TFLite frontend regression tests covering dynamic `FILL` import
and
dynamic `SPLIT_V` import/compile behavior.
This addresses the `FILL` and `SPLIT_V` items under the "Fix partial
implementations" section of #19412.
Validation:
```bash
python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.py
```
```bash
python -m pytest tests/python/relax/test_frontend_tflite.py \
-k "split_v_dynamic or fill_dynamic_dims" -v
```
Result:
- All checks passed
## Summary
This PR completes the TFLite `NON_MAX_SUPPRESSION_V5` implementation in
Relax by adding support for `soft_nms_sigma != 0`.
It extends `relax.vision.non_max_suppression` with soft-NMS attributes,
updates the TFLite frontend to consume the soft-NMS outputs correctly,
and aligns the TOPI implementation with LiteRT's reference behavior.
Relates to #19412.
## Changes
1. **Relax / TOPI soft-NMS support**
- Extend `NonMaximumSuppressionAttrs` with `soft_nms_sigma` and
`score_threshold`.
- Thread the new attributes through Relax op registration, Python
wrapper, and legalization.
- Add soft-NMS handling to TOPI classic NMS so
`relax.vision.non_max_suppression` can represent the
`NON_MAX_SUPPRESSION_V5` behavior.
2. **TFLite frontend support for `NON_MAX_SUPPRESSION_V5`**
- Remove the previous `soft_nms_sigma != 0` unsupported-path guard in
the TFLite frontend.
- Forward `soft_nms_sigma` and `score_threshold` into
`relax.vision.non_max_suppression`.
- Handle the soft-NMS return path explicitly so the frontend reads
decayed scores from the processed NMS output instead of re-reading the
original score tensor.
3. **Soft-NMS correctness fixes**
- Fix the soft-NMS path so boxes whose scores fall below the threshold
after decay are invalidated consistently.
- Keep returned indices and decayed scores aligned in both the TOPI TIR
implementation and the NumPy reference implementation.
- Update the soft-NMS candidate selection logic to re-pick the current
best candidate after each decay step, matching LiteRT's
reference behavior.
- Align the Gaussian decay formula with LiteRT.
4. **Test coverage**
- Add Relax tests for soft-NMS struct-info inference and legalization.
- Add Relax E2E tests covering reordered outputs after score decay and
other soft-NMS follow-up cases.
- Add TFLite frontend tests for `NON_MAX_SUPPRESSION_V5` with
`soft_nms_sigma != 0`.
- Add IR checks to verify that `soft_nms_sigma` and `score_threshold`
are forwarded correctly.
## Testing
```bash
python -m pytest -n 1 tests/python/relax/test_op_vision.py -k "all_class_non_max_suppression or get_valid_counts or nms" -v
python -m pytest tests/python/relax/test_frontend_tflite.py -k "nms_v5" -v
```
## Result:
- Relax vision tests passed locally
- TFLite `NON_MAX_SUPPRESSION_V5` coverage added for both hard-NMS and
soft-NMS paths
## Summary
- Implemented the TFLite `DETECTION_POSTPROCESS` operator conversion to
Relax IR.
- Wires up the previously unimplemented operator to support object
detection post-processing workflows in Relax.
- Relates to #18928
## Changes
- **Operator Registration**: Implemented `convert_detection_postprocess`
in `python/tvm/relax/frontend/tflite/tflite_frontend.py`.
- **Core Logic**:
- Integrated `multibox_transform_loc` for coordinate decoding and
variance scaling.
- Supported `use_regular_nms` attribute to switch between all-class NMS
and class-agnostic NMS paths.
- Leveraged `all_class_non_max_suppression` for efficient box filtering.
- **Output Alignment**: Used `topk`, `gather_nd`, and `where` operators
to ensure the output tensors (boxes, classes, scores, num_detections)
match the TFLite specification in terms of shape and layout.
- **Attribute Validation**: Added strict validation for required custom
options such as `num_classes`, `max_detections`, and scaling factors.
## Validation
Verified with linting and pre-commit hooks:
```bash
# Lint check
python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py
# Pre-commit checks
python -m pre_commit run --files python/tvm/relax/frontend/tflite/tflite_frontend.py
```
Result:
- **Passed**: All static checks and style guidelines are met.
## Summary
Add Relax `roi_pool` support and wire it through the ONNX frontend for
`MaxRoiPool`.
## Changes
- add `relax.vision.roi_pool`, including attrs, Python wrapper, struct
info inference, and legalization
- add TOPI `roi_pool` compute for NCHW layout
- support ONNX `MaxRoiPool` in the Relax ONNX frontend
- handle empty / out-of-bound pooled bins according to ONNX/reference
semantics, returning `0` instead of propagating invalid reductions
- add regression tests for Relax op inference, legalization, and ONNX
frontend import
- add out-of-bound ROI coverage to make sure fully invalid pooled bins
still match ONNX Runtime
## Validation
- `pytest tests/python/relax/test_op_vision.py -k roi_pool`
- `pytest tests/python/relax/test_frontend_onnx.py -k 'max_roi_pool'`
This PR completes the `MaxRoiPool` portion of the Relax ONNX frontend
operator work tracked in #18945.
Introduce relax.nn.conv3d_transpose (attrs, C++ inference/layout, Python
API) and lower it to TOPI group_conv3d_transpose_ncdhw when using
NCDHW/IODHW with dilation 1, matching the conv2d_transpose legalization
policy.
Wire the Relax ONNX frontend to emit conv3d_transpose for 5D inputs.
Extend tests for ONNX, struct info, LegalizeOps, and TVMScript
round-trip; fix ConvTranspose test output spatial size to include
output_padding.https://github.com/apache/tvm/issues/18945
## Summary
Add `relax.vision.get_valid_counts` and classic
`relax.vision.non_max_suppression` for object-detection post-processing
pipelines.
`get_valid_counts` performs score-based bounding box filtering and
compacts valid boxes to the front of each batch. Classic
`non_max_suppression` performs flexible IoU-based suppression on
filtered boxes, complementing existing `all_class_non_max_suppression`
for custom post-processing workflows.
This PR implements the Relax-level registration, legalization, TOPI
compute, and test coverage for both operators.
## Changes
**Relax op registration and legalization:**
- C++ op functions, FFI registration, and struct info inference for both
operators (`vision.h`, `vision.cc`)
- Python wrappers with Relax docstrings (`vision.py`)
- Legalization to `topi.vision.get_valid_counts` and
`topi.vision.non_max_suppression`
- Additional struct-info validation for `score_index`, `id_index`, and
`coord_start` when `elem_length` is statically known
**TOPI and testing:**
- Full TOPI implementation for `get_valid_counts`
- Reimplementation of classic `non_max_suppression` in TOPI
- NumPy reference implementations in `tvm.topi.testing` for both
operators
- Op-level tests for struct info inference, legalization, invalid
attribute ranges, and e2e numerical correctness
- Stronger legalization tests that verify both `relax.call_tir`
introduction and removal of the original Relax vision op
## Limitations
- Attribute range validation for `score_index`, `id_index`, and
`coord_start` is only enforced when the input `elem_length` is
statically known during struct-info inference.
- Classic `non_max_suppression` follows the existing Relax / TOPI API
shape and is intended for single-class or class-aware custom
post-processing flows, distinct from `all_class_non_max_suppression`.
## Validation
```bash
pytest tests/python/relax/test_op_vision.py -k "get_valid_counts" -v
pytest tests/python/relax/test_op_vision.py -k "test_nms_" -v
```
All related tests passed.
Introduce relax.vision.multibox_transform_loc with
MultiboxTransformLocAttrs: decode center-size offsets against ltrb
priors, softmax on class logits, and optional clip, threshold masking,
and background score zeroing. Register the C++ op with FInferStructInfo
checks for shapes and dtypes (including batch and 4*N consistency).
Legalize to topi.vision.multibox_transform_loc.
Add tests for struct inference, invalid inputs, Legalize+e2e on LLVM,
attribute branches, and TVMScript roundtrip. Add a standalone numpy
reference under topi/testing (not exported from tvm.topi.testing to
avoid pulling scipy).
Update TFLite frontend NotImplementedError text for
DETECTION_POSTPROCESS and NON_MAX_SUPPRESSION_V5 to note multibox is
available and link tracking issue #18928.
## Summary
Add `relax.image.affine_grid` operator for Spatial Transformer Networks,
along with PyTorch and ONNX frontend integration.
TOPI compute (`topi.image.affine_grid`) already exists. This PR
completes the Relax-level registration and frontend support, following
the existing `resize2d` / `grid_sample` pattern.
## Changes
**Relax op registration:**
- C++ op function, FFI registration, and struct info inference
(`resize.h`, `resize.cc`)
- Python wrapper with flexible size parameter handling (`image.py`)
- Legalization to `topi.image.affine_grid` with `PrimExpr` → `int`
conversion
- Op-level tests (struct info inference + e2e numerical correctness) and
legalization test
**PyTorch frontend:**
- Converter for `aten.affine_grid_generator.default`
- Layout conversion from TVM `[N,2,H,W]` to PyTorch `[N,H,W,2]` via
`permute_dims`
- Single-kernel path is 5.6x faster than the decomposed path (30+ ops)
- Structural IR test and numerical correctness test
**ONNX frontend:**
- `AffineGrid` converter with `_impl_v20` (opset 20, when the op was
first introduced)
- Support for constant size tensor `[N,C,H,W]`
- Layout conversion from TVM `[N,2,H,W]` to ONNX `[N,H,W,2]`
- End-to-end correctness test against ONNX Runtime
## Limitations
- Only `align_corners=True` is supported (matches current TOPI
implementation)
- Only 2D affine grid is supported
## Validation
```bash
pytest tests/python/relax/test_op_image.py -k "affine_grid" -v # 8 passed
pytest tests/python/relax/test_transform_legalize_ops_image.py -k "affine_grid" -v # 1 passed
pytest tests/python/relax/test_frontend_from_exported_program.py -k "affine_grid" -v # 2 passed
pytest tests/python/relax/test_frontend_onnx.py -k "affine_grid" -v # 1 passed
```
All 12 tests passed.
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
## Summary
Add Relax `roi_align` support and wire it through the ONNX and PyTorch
frontends.
## Changes
- add `relax.vision.roi_align`, including attrs, Python wrapper, struct
info inference, and legalization
- add TOPI `roi_align` compute and keep both legacy and aligned ROIAlign
semantics
- support ONNX `RoiAlign`, including `coordinate_transformation_mode`
handling for `output_half_pixel` and `half_pixel`
- support PyTorch `torchvision.ops.roi_align` in the exported-program
frontend, including the `aligned` flag
- add regression tests for Relax op inference, legalization, TVMScript
parsing, ONNX frontend import, and PyTorch frontend import
- add aligned ROIAlign test coverage to make sure sub-pixel RoIs no
longer use the legacy `min=1.0` clamp
## Validation
- `pytest tests/python/relax/test_op_vision.py -k roi_align`
- `pytest tests/python/relax/test_tvmscript_parser_op_vision.py -k
roi_align`
- `pytest tests/python/relax/test_frontend_onnx.py -k roi_align`
- `pytest tests/python/relax/test_frontend_from_exported_program.py -k
roi_align`
This PR completes the Relax/ONNX/Torch roi_align work tracked in #18928.
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 `FewShotTuning` pass from Relax transform (C++ implementation,
Python bindings, and test file)
- The pass is unused in the current codebase and can be safely removed
## Files Changed
- `include/tvm/relax/transform.h` — Remove declaration
- `python/tvm/relax/transform/__init__.py` — Remove from imports
- `python/tvm/relax/transform/transform.py` — Remove Python function
- `src/relax/transform/few_shot_tuning.cc` — Delete (C++ implementation)
- `tests/python/relax/test_transform_few_shot_tuning.py` — Delete (test
file)
This PR enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
This PR moves remaining related data structures to s_tir.
- Moves sblock_dependency_info and sblock_scope.
- Moves related analyssis.
- Hides the data_type_rewriter to private functions.
Guan-Ming Chiu