## 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.
## 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
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.
## 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
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.
Tianqi Chen