4bcf694cbf211121a600435bca48967eabef360a
6338 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Tianqi Chen
|
4bcf694cbf |
[REFACTOR][IR] Inline ReplaceGlobalVars into AttachGlobalSymbol (#19625)
## Summary `ReplaceGlobalVars` was a public IR-layer API with only one in-tree C++ caller (`relax::AttachGlobalSymbol`). The mechanism used a NodeFunctor vtable populated at static-init time by per-dialect `.cc` files in relax and tirx, which made the IR layer logically depend on its dialects even though the include graph did not show it. Move the dispatch logic into the consumer as file-local mutators and a private helper. Delete the public header, the IR-layer driver, both per-dialect dispatch registrations, the `IRModule.replace_global_vars` python method, and its dedicated test file. The behavior is still covered by `tests/python/relax/test_transform_attach_global_symbol.py` and by the pipelines that include the `AttachGlobalSymbol` pass. |
||
|
Tianqi Chen
|
2f4f4b1de3 |
[REFACTOR][IR][FFI] Bump tvm-ffi (+ SEqHashDef migration) and phase out tvm/ir/repr.h (#19627)
## Summary Two-commit PR: 1. Bump `3rdparty/tvm-ffi` from `3c35034` to `98d0029` and migrate all 21 in-tree `SEqHashDef()` call sites to `SEqHashDefRecursive()` (the conservative variant matching the prior default behavior). Six let-style sites carry `TODO(tqchen)` comments indicating they should flip to `SEqHashDefNonRecursive` after the new tvm-ffi ships on pypi. 2. Phase out `include/tvm/ir/repr.h`. The bumped tvm-ffi now provides ostream `operator<<` for `Any`/`ObjectRef`/`Variant`/`Optional` directly in `tvm/ffi/extra/dataclass.h`, making the in-tree thin wrapper redundant. Rewrite 8 includers, rename `src/ir/repr.cc` → `src/ir/access_path_repr.cc` (preserves `node.AsRepr` + AccessPath/AccessStep `__ffi_repr__` registrations; drops zero-caller `tvm::Dump()`), delete the header. Also fixes a Python-level import regression in `python/tvm/ir/attrs.py` caused by the bump: tvm_ffi 0.1.12.dev changes the field-registration guard from `not hasattr(cls, name)` to `name not in cls.__dict__`, which breaks `DictAttrs` because `DictAttrsNode` registers a reflection field named `"__dict__"` — Python forbids installing a class descriptor with that name via `setattr`. Fix: define `__dict__` as an explicit Python property on `DictAttrs` so the auto-installation is skipped. ## TODO follow-ups After the new tvm-ffi releases on pypi, flip the 6 `SEqHashDefRecursive()` sites that carry `TODO(tqchen)` comments to `SEqHashDefNonRecursive()`. Locations are enumerated in the commit body of commit 1. ## Test plan - [x] Full ninja build clean (638/638). - [x] 118/118 cpptest pass. - [x] `import tvm; tvm.cuda(0).exist` returns True. - [x] `tests/python/all-platform-minimal-test`: 37 passed, 105 skipped. - [x] `tests/python/relax/test_struct_info.py`: 9 passed. - [x] `git grep -nE 'SEqHashDef\(|"tvm/ir/repr\.h"'` is empty. - [x] `pre-commit run --all-files` clean. |
||
|
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. |
||
|
Tianqi Chen
|
f0ac8d62ef |
[REFACTOR][RUNTIME] Phase out tvm::runtime::regex_match (#19620)
## Summary
`tvm::runtime::regex_match` was a thin C++ wrapper that bounced through
a
global `ffi::Function` back into Python's `re.match`. It was introduced
solely to avoid pulling `<regex>` into TVM (libstdc++ dual-ABI conflict
with
pre-cxx11 pytorch wheels). The only C++ caller is the DNNL JSON runtime,
where
every pattern reduces to substring containment — `re.match` anchors at
the
start only, so `.*X.*` is equivalent to `s.find(X) != npos`.
- Remove `src/runtime/regex.{h,cc}` and the Python
`tvm.runtime.regex_match`
global registration.
- Add file-local `contains` / `contains_any` helpers in
`dnnl_json_runtime.cc`
and inline `std::string::find` at the 15 call sites.
- Drop the dead `regex.h` include from
`src/relax/transform/update_param_struct_info.cc`.
No CMakeLists.txt change needed — `src/runtime/*.cc` is picked up by
glob.
`USE_DNNL` is OFF in the ci_gpu container, so DNNL-specific runtime
tests
are not exercised locally. The DNNL translation unit compiles cleanly
with
the inlined helpers, and the full TVM build (636 targets) passes.
|
||
|
Shushi Hong
|
de89da6b18 |
[IR] Rename Call annotations to attrs (#19618)
This PR renames `tirx::CallNode::annotations` to `attrs`, matching the existing Relax `CallNode::attrs` convention. Previously, TIRX Call metadata was stored in a `Map<String, Any>` field named `annotations`. This PR makes it a first-class `Attrs` field instead, so call-level metadata follows the same representation and naming style as Relax calls. |
||
|
YinHanke
|
dcbebe7bfd |
[Relax][Frontend][TFLite] Add UNIDIRECTIONAL_SEQUENCE_RNN converter (#19601)
## Summary This PR adds Relax TFLite frontend support for `UNIDIRECTIONAL_SEQUENCE_RNN` (BuiltinOperator 35), claimed in [#19519](https://github.com/apache/tvm/issues/19519) Group A. The op executes a simple RNN cell over a time sequence. The converter unrolls the time steps at graph-construction time using Relax primitives. Cell equation: ``` h_t = fused_activation(x_t @ W.T + h_{t-1} @ Wr.T + b) ``` ## Changes - **Handler**: `convert_unidirectional_sequence_rnn` registered in `convert_map` (alphabetical, U-region after `UNPACK`) - **Inputs** (5): `input [batch, time, input_size]`, `input_weights [num_units, input_size]`, `recurrent_weights [num_units, num_units]`, `bias [num_units]`, `hidden_state [batch, num_units]` (variable, zero-initialised) - **Output**: `[batch, time, num_units]` (always batch-major) - **time_major=True**: input is transposed to batch-major before unrolling - **Activations**: NONE, RELU, RELU6, TANH, SIGMOID (via `convert_fused_activation_function`) - **Quantized**: raises `OpNotImplemented` (not yet supported) ## Testing Modern TF/Keras (2.x, Keras 3) no longer emits `UNIDIRECTIONAL_SEQUENCE_RNN`; `SimpleRNN` with `unroll=False` lowers to `WHILE`+TensorList ops, and `unroll=True` expands to elementwise ops. Tests therefore follow the same flatbuffer-construction pattern used by the StableHLO op PRs (#19536, #19587). Three tests added to `tests/python/relax/test_frontend_tflite.py`: - `test_unidirectional_sequence_rnn_none_activation` — `tvm.ir.assert_structural_equal` with identity weights / zero bias, NONE activation, time=1 - `test_unidirectional_sequence_rnn_relu_activation` — shape check, random weights, RELU activation, time=3 - `test_unidirectional_sequence_rnn_time_major` — shape check, `time_major=True` input layout ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k unidirectional_sequence_rnn -v ``` All 3 tests pass. pre-commit (ASF header, ruff check, ruff format) all pass. ## References - Issue [#19519](https://github.com/apache/tvm/issues/19519) Group A: Sequence / recurrent model operators Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> |
||
|
HoYi
|
fa66213249 |
[Relax][Frontend][TFLite] Support control-flow multi-subgraph operators (#19616)
## Summary This PR adds Relax TFLite frontend support for the TFLite builtin control-flow / multi-subgraph operator family from #19519 item F: `CALL`, `IF`, `WHILE`, and `CALL_ONCE`. It builds on the multi-subgraph import infrastructure merged in PR #19587. The frontend already accepts TFLite models with extra subgraphs while converting only `Subgraphs(0)` into the Relax `main` function. This PR uses those extra subgraphs as callable or control-flow regions for the TFLite control-flow operators. The supported subset is intentionally pure tensor and guard-first: - `CALL` lowers a referenced TFLite subgraph to a private Relax function and emits a direct call. - `IF` lowers the then/else subgraphs to private Relax functions and emits a private wrapper function containing Relax `If`. - `WHILE` lowers the cond/body subgraphs to private Relax functions and emits a recursive private Relax function for the loop. - `CALL_ONCE` supports the empty-init no-op subset and explicitly rejects non-empty or resource-like init patterns. This PR does not model resource variable side effects. Those cases remain explicitly guarded instead of being imported with incorrect pure functional semantics. ## Design ### Shared Subgraph Lowering The frontend now keeps shared conversion state across the main graph and referenced subgraphs: - `lowered_subgraphs` - `lowered_if_functions` - `lowered_while_functions` - `lowering_stack` - `module_builder` Referenced pure tensor subgraphs are lowered through a recursive `OperatorConverter` using an isolated `ExprTable`, so subgraph tensor bindings cannot overwrite bindings from the main graph. Lowered subgraphs are cached by subgraph index and reused when the same region is referenced more than once. Generated private functions are registered through the shared parent `module_builder`, so nested cases such as `main CALL -> subgraph A -> CALL subgraph B` keep all private functions in the final IRModule. Recursive ordinary `CALL` subgraphs are guarded with `OpNotImplemented`. `WHILE` uses a dedicated recursive wrapper function instead, because recursion is part of the intended Relax representation for the loop itself. ### Boundary Validation The control-flow converters validate subgraph boundaries before lowering: - referenced subgraph indices must be valid - op input/output arity must match the referenced subgraph interface - branch and loop tensor shape/dtype metadata must match the surrounding op - `IF` and `WHILE` conditions must be scalar bool tensors - `WHILE` loop-carried input/output tensors must have matching metadata The shared `_check_subgraph_interface` helper is used by `CALL`, `IF`, and `WHILE` to keep arity and metadata checks consistent across the control-flow operators. `_require_scalar_bool_tensor` accepts both frontend `TensorWrapper` objects and raw TFLite tensors so caller and referenced-subgraph condition checks use the same path. These checks keep the first implementation conservative and make unsupported cases fail with targeted `OpNotImplemented` diagnostics. ### Tuple Outputs TFLite `CALL`, `IF`, and `WHILE` may produce multiple output tensors. The frontend maps those cases to Relax tuple returns: ```text single output -> tensor expression multi output -> Tuple(...) op outputs -> TupleGetItem(...) ``` This keeps the single-output IR simple while covering multi-output calls, multi-output branches, and multi-variable loop state. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `CALL` | `CallOptions.Subgraph()` | private Relax function call | pure tensor subgraphs, single or multiple outputs | | `IF` | `IfOptions.ThenSubgraphIndex()`, `ElseSubgraphIndex()` | private wrapper function containing Relax `If` | scalar bool condition, matching branch I/O metadata | | `WHILE` | `WhileOptions.CondSubgraphIndex()`, `BodySubgraphIndex()` | recursive private Relax function | scalar bool cond output, tensor loop-carried state | | `CALL_ONCE` | `CallOnceOptions.InitSubgraphIndex()` | no-op for empty init subgraph | empty init subgraph only | ## Not Included - Full `CALL_ONCE` resource/variable initialization semantics. - Resource, variant, hashtable, or variable tensor support. - TensorFlow-generated `tf.cond` / `tf.while_loop` smoke tests. - Dynamic-shape loop-state refinements beyond the current static metadata checks. ## Tests The tests manually build minimal TFLite flatbuffers and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. Unsupported-boundary tests use `pytest.raises`. | Test | Coverage | |---|---| | `test_call_subgraph` | basic `CALL` to a pure tensor subgraph | | `test_call_subgraph_multi_output` | `CALL` tuple return and output binding | | `test_call_subgraph_nested_call` | nested `CALL` private function registration | | `test_call_subgraph_invalid_index_unsupported` | invalid `CALL` subgraph index | | `test_call_subgraph_io_mismatch_unsupported` | `CALL` arity mismatch | | `test_call_subgraph_output_metadata_mismatch_unsupported` | `CALL` output metadata guard | | `test_if_subgraphs` | basic `IF` branch selection | | `test_if_subgraphs_multi_output` | `IF` tuple branch returns | | `test_if_subgraphs_non_bool_condition_unsupported` | `IF` condition dtype guard | | `test_if_subgraphs_invalid_index_unsupported` | invalid then/else subgraph index | | `test_if_subgraphs_output_count_mismatch_unsupported` | branch output count guard | | `test_if_subgraphs_input_metadata_mismatch_unsupported` | branch input metadata guard | | `test_if_subgraphs_output_metadata_mismatch_unsupported` | branch output metadata guard | | `test_while_subgraphs` | basic recursive `WHILE` lowering | | `test_while_subgraphs_repeated_cond_body_pair` | shared cond/body loop function cache | | `test_while_subgraphs_two_loop_vars` | multi-variable loop state tuple path | | `test_while_subgraphs_non_bool_condition_unsupported` | `WHILE` cond output dtype guard | | `test_while_subgraphs_invalid_index_unsupported` | invalid cond/body subgraph index | | `test_while_subgraphs_zero_loop_vars_unsupported` | zero-loop-var guard | | `test_while_subgraphs_loop_state_metadata_mismatch_unsupported` | loop state metadata guard | | `test_while_subgraphs_output_count_mismatch_unsupported` | body output count guard | | `test_while_subgraphs_input_metadata_mismatch_unsupported` | cond/body input metadata guard | | `test_while_subgraphs_output_metadata_mismatch_unsupported` | cond/body output metadata guard | | `test_call_once_empty_init_subgraph` | empty `CALL_ONCE` no-op subset | | `test_call_once_non_empty_init_subgraph_unsupported` | non-empty init subgraph guard | | `test_call_once_inputs_outputs_unsupported` | `CALL_ONCE` op I/O guard | | `test_call_once_init_subgraph_io_unsupported` | init subgraph I/O guard | | `test_call_once_invalid_index_unsupported` | invalid init subgraph index | Local validation: ```bash python -m ruff format --check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m pytest \ tests/python/relax/test_frontend_tflite.py \ -k "call_subgraph or if_subgraphs or while_subgraphs or call_once" -q python -m pytest \ tests/python/relax/test_frontend_tflite.py -q ``` Result: ```text ruff format --check: 2 files already formatted ruff check: All checks passed 28 passed, 434 deselected 462 passed ``` ## References - Issue #19519 item F: TFLite control-flow / multi-subgraph operators - PR #19587: StableHLO region-based ops and multi-subgraph model support |
||
|
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
|
||
|
Tianqi Chen
|
e159487b0e |
[REFACTOR][IR] attrs.h follow-up cleanup: drop legacy vtable / rename / phase out AttrFieldInfo (#19615)
## Summary Follow-up to #19607 that continues trimming `attrs.h` and adjacent files. The six commits land independently and each builds clean. - Phase out `OpNode::arguments` and `AttrFieldInfo` — the field stored metadata that no Python tooling, test, or C++ caller (beyond internal sanity checks) read; removing it deletes `AttrFieldInfo` plus ~335 chained `.add_argument(...)` calls. The remaining 12 internal consumers now read `op->num_inputs` and report indexed inputs (`input[i]`). - Drop the (unused) virtual destructor on `BaseAttrsNode` (ffi::Object uses a captured-typed deleter, no virtual dispatch needed) and inline the trivial 3-line `DictAttrs(Map)` constructor into the header. - Rename `BaseAttrsNode` → `AttrsNode`; the `Base` prefix existed only to distinguish from the `AttrsNodeReflAdapter` shim that #19607 removed. The `"ir.Attrs"` FFI registry key is unchanged. - Promote `DictAttrs` to NOTNULLABLE (`TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE` + COW macro). The no-arg `DictAttrs()` constructor already created an empty backing, so every existing call site already produced a defined object; ~15 defensive `attrs.defined()` checks (and a defensive Python `None` fallback in `Function`) are now redundant. - Inline the `WithAttr(DictAttrs, ...)` / `WithAttrs(DictAttrs, ...)` free-function overloads into the TFunc-template wrappers — those overloads had no external callers (no TVM_DLL, no Python binding). - Rename `AttrsWithDefaultValues<T>` → `PassConfigWithDefaults<T>` and move from `attrs.h` to `transform.h`; all 9 consumers are pass-config classes registered via `TVM_REGISTER_PASS_CONFIG_OPTION`. `attrs.h` shrinks from 363 → 262 lines. |
||
|
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>>`.
|
||
|
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
|
||
|
Tianqi Chen
|
3918e14389 |
[REFACTOR][IR] Inline ApplyPassToFunction into relax decompose_ops, delete the util (#19612)
## Summary `ApplyPassToFunction` is a general-purpose wrapper that runs a pass on only the functions in an IRModule whose name matches a regex. Its sole in-tree production callers are `DecomposeOpsForInference` / `DecomposeOpsForTraining` in `src/relax/transform/decompose_ops.cc`, and both callers always supply a literal function name (never a regex pattern). Inlining the logic as a file-local helper simplifies the module-level context and removes an abstraction that exists only to support one use case. - Inline the helper as `ApplyDecomposeToFunction` (exact-name match, not regex) in `src/relax/transform/decompose_ops.cc` - Delete `src/ir/apply_pass_to_function.cc`, its `transform.h` declaration, and the Python wrapper in `python/tvm/ir/transform.py` - Remove two DCE tests (`test_compatibility_with_apply_pass_to_function`, `test_well_formed_output_with_restricted_scope`) that tested the utility's plumbing rather than DCE behavior |
||
|
Tianqi Chen
|
d37b6abd56 |
[REFACTOR][IR] Phase out src/ir/structural_{hash,equal}.cc to tvm-ffi (#19613)
## Summary The tvm-ffi layer now provides fully featured structural-hash and structural-equal APIs (including `GetFirstStructuralMismatch` with `AccessPath` pair output). The two TUs `src/ir/structural_hash.cc` and `src/ir/structural_equal.cc` had become thin adapters with no logic of their own — they forwarded to tvm-ffi and registered the results as `node.Structural*` globals for Python to call. This PR removes the indirection. - **Commit A** (`[REFACTOR][IR]`): relocates the `ffi::ModuleObj` and `ffi::TensorObj` `__data_to_json__`/`__data_from_json__` `TypeAttrDef` registrations from `structural_hash.cc` into `src/ir/module.cc` and `src/runtime/tensor.cc` respectively, both of which already have a `TVM_FFI_STATIC_INIT_BLOCK` for those types. - **Commit B** (`[REFACTOR][PYTHON]`): rewrites the four Python wrappers in `tvm.ir.base` (`structural_equal`, `get_first_structural_mismatch`, `assert_structural_equal`, `structural_hash`) to call `tvm_ffi._ffi_api` directly, bypassing the now-redundant `node.Structural*` globals. `assert_structural_equal` reconstructs the same diagnostic message in Python using `TVMScriptPrinterScript` with `path_to_underline`. - **Commit C** (`[REFACTOR][IR]`): deletes `src/ir/structural_hash.cc` and `src/ir/structural_equal.cc` whose remaining content (the `node.Structural*` FFI global registrations) is now unused. |
||
|
Tianqi Chen
|
b1e1566f82 |
[REFACTOR][IR] Cleanup attrs.h: drop NullValue, AttrsNodeReflAdapter, legacy BaseAttrsNode methods (#19607)
## Overview This PR cleans up `include/tvm/ir/attrs.h` by removing four deprecated abstractions: 1. `NullValue<T>()` sentinel helpers (replaced by `ffi::Optional<T>`) 2. `AttrsNodeReflAdapter<DerivedType>` shim template (Attrs structs now inherit `BaseAttrsNode` directly) 3. `BaseAttrsNode::InitBySeq` / `InitByPackedArgs` legacy initialization methods 4. `DictAttrsNode::InitByPackedArgs` override It also migrates 9 pass-config classes from `Attrs`/`AttrsNodeReflAdapter` to `ffi::Object`, since they are pass configuration objects, not IR attributes. ## Changes **Commit A — Replace NullValue<T>() call sites** (`[REFACTOR][IR] Replace NullValue<T>() call sites with default construction`) - 11 source files: replace `NullValue<T>()` with `T()`, `std::nullopt`, or `DataType::Void()` - `manipulate.h`/`manipulate.cc`: `FlipAttrs::axis` changed from `Integer` to `ffi::Optional<int64_t>` **Commit B — Drop NullValue, AttrsNodeReflAdapter, legacy BaseAttrsNode methods** (`[REFACTOR][IR] Drop NullValue declaration, AttrsNodeReflAdapter, BaseAttrsNode legacy methods`) - `include/tvm/ir/attrs.h`: removes `NullValue<T>`, `InitBySeq`, `InitByPackedArgs`, `AttrsNodeReflAdapter<T>` - `src/ir/attrs.cc`: removes `DictAttrsNode::InitByPackedArgs` definition - `AttrsWithDefaultValues<T>()` broadened to accept any `ffi::ObjectRef` subtype (needed for Commit D) - Removes unused includes: `reflection/accessor.h`, `<functional>`, `<vector>` **Commit C — Subclass BaseAttrsNode directly** (`[REFACTOR][IR] Subclass BaseAttrsNode directly, drop AttrsNodeReflAdapter`) - 17 attrs headers in `include/tvm/relax/attrs/` + `include/tvm/target/virtual_device.h` - All `struct FooAttrs : public AttrsNodeReflAdapter<FooAttrs>` → `struct FooAttrs : public BaseAttrsNode` **Commit D — Migrate pass-config classes to ffi::Object** (`[REFACTOR] Migrate pass-config classes to subclass ffi::Object`) - 9 pass-config classes in `src/s_tir/`, `src/tirx/`, `src/relax/backend/contrib/` - `XConfigNode : public ffi::Object` (was `AttrsNodeReflAdapter<XConfigNode>`) - `XConfig : public ffi::ObjectRef` (was `Attrs`) - Python bindings updated: 7 classes changed from `_ir.Attrs` to `_ffi.Object` ## Design Decisions **`AttrFieldInfo` / `OpNode::arguments` kept**: Pre-flight check revealed `GetArgStructInfo()` in `op_common.h` and `op_common.cc` actively reads `op->arguments` (names, counts). These were not dead metadata — deleting them would break Relax op argument validation. They are kept as-is. **Commit E (trim attrs.h includes) reduced in scope**: Removing `structural_equal.h`, `structural_hash.h`, and `<unordered_map>` from `attrs.h` caused 47 downstream files to fail compilation. Rather than adding explicit includes to 47 files, only clearly-unused includes (`reflection/accessor.h`, `<functional>`, `<vector>`) were removed in Commit B. ## Testing - Build: clean compile with `-DUSE_CUDA=OFF -DUSE_LLVM=ON` - Tests passing: - `tests/python/ir/` (93 passed) - `tests/python/relax/test_analysis.py`, `test_blockbuilder_core.py`, `test_op_manipulate.py`, `test_transform.py` (209 passed) - `tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py`, `test_s_tir_transform_unify_thread_binding.py` (30 passed) - `tests/python/tirx-transform/test_tir_transform_unroll_loop.py`, `test_tir_transform_simplify.py`, `test_tir_transform_remove_no_op.py` (108 passed, 6 xfailed) - Pre-existing failures (unrelated to this PR): `test_s_tir_transform_lower_opaque_block`, `test_s_tir_transform_compact_buffer_region::TestLetBinding::test_compact`, `test_tir_transform_vectorize::test_vectorize_llvm_pure_intrin_fail` |
||
|
HoYi
|
2441461d12 |
[Relax][Frontend][TFLite] Support quantized TFLite import via QDQ decomposition (#19538)
## Summary This PR adds initial quantized TFLite import support to the Relax frontend by preserving tensor quantization metadata and replacing placeholder `_qnn.op.*` frontend calls with an explicit QDQ decomposition: ```text dequantize -> float Relax op -> quantize ``` Before this PR, the Relax TFLite frontend raised `NotImplementedError` as soon as quantization metadata was seen during tensor parsing. This made quantized TFLite models unreachable. This PR keeps `scale`, `zero_point`, and `QuantizedDimension()` in `TensorWrapper.qnn_params`, then uses the existing `R.quantize` / `R.dequantize` operators to lower supported quantized paths. The previous `_qnn.op.*` paths were effectively unreachable for normal quantized TFLite models because `get_tensors()` raised `NotImplementedError` as soon as valid quantization metadata was parsed. After removing that blocker, those paths also needed to be replaced because they depended on undefined `_qnn` helpers and did not handle Relax QDQ, axis remapping, or quantized bias consistently. Closes #19534. ## Design Relax already has `R.quantize` and `R.dequantize` with C++ registration, Python APIs, legalization, and tests. Instead of introducing new fused Relax QNN ops for this first import PR, the frontend now decomposes quantized TFLite operators through QDQ around ordinary Relax float operators. This keeps the change scoped to the Python TFLite frontend and existing Relax QDQ operators, while establishing a working import path first. Fused int8 Relax QNN operators can still be considered later if backend kernel selection requires them. ## Updated Converters | Converter | Replacement | |---|---| | `get_tensors` | Preserve `scale`, `zero_point`, and `QuantizedDimension()` | | `quantize` / `dequantize` helpers | Use `R.quantize` / `R.dequantize` with `axis` | | `convert_quantize` | `float -> Q` and quantized requantize as `DQ -> Q` | | `convert_dequantize` | Use `R.dequantize` | | `convert_relu`, `convert_relu6`, `convert_relu_n1_to_1` | `DQ -> activation -> Q` | | `_convert_elemwise` | Quantized binary ops use `DQ -> op -> fused activation -> Q`; comparisons use `DQ -> compare` | | `convert_reshape` | uint8 different-qparams path uses `DQ -> reshape -> Q` | | `_convert_reduce` | Quantized reduce uses `DQ -> reduce -> Q` | | `convert_conv` | Quantized Conv2D uses `DQ input + DQ weight -> conv2d -> Q` | | `convert_fully_connected` | Quantized FC uses `DQ input + DQ weight -> matmul -> Q` | | `convert_concatenation` | Quantized concat uses `DQ each -> concat -> Q` | | `convert_transpose_conv` | Quantized transpose conv uses `DQ input + DQ weight -> conv2d_transpose -> Q` | | `convert_detection_postprocess` | Inline `_qnn.op.dequantize` calls replaced with `self.dequantize` | All `_qnn.op.*` references are removed, and the stale `# ruff: noqa: F821` suppression is no longer needed. ## Axis Remapping The most correctness-sensitive part of this PR is axis remapping for per-channel weight dequantization after the frontend rewrites TFLite layouts into Relax layouts. | Op | TFLite layout | Relax layout | Axis remap | |---|---|---|---| | Conv2D | `[OC, KH, KW, IC]` | `[KH, KW, IC, OC]` (`HWIO`) | `0 -> 3` | | FullyConnected | `[OC, IC]` | `[IC, OC]` | `0 -> 1` | | TransposeConv | `[OC, KH, KW, IC]` (`OHWI`) | `[IC, OC, KH, KW]` (`IOHW`) | `0 -> 1` | | DepthwiseConv | `[1, KH, KW, C*M]` | `[KH, KW, C, M]` (`HWOI`) | per-channel unsupported | For Conv2D, FC, and TransposeConv, non-zero weight `QuantizedDimension()` values are rejected with `OpAttributeInvalid`, because the supported quantized TFLite weight layout uses output-channel axis 0. Per-channel depthwise convolution is guarded with `OpNotImplemented`. The TFLite depthwise reshape changes the channel-axis semantics in a way that this initial QDQ lowering does not represent directly. ## Bias Handling TFLite INT32/INT64 bias tensors may not store explicit quantization metadata. For quantized Conv2D, FullyConnected, and TransposeConv, the frontend follows the implicit TFLite convention and dequantizes integer bias using: ```text bias_scale = input_scale * weight_scale bias_zero_point = 0 axis = 0 ``` This supports both per-tensor and per-channel weight scales. The per-channel case is covered by a structural regression test that expects vector bias scale. ## Fused Activation Handling Conv2D, FullyConnected, and quantized concat preserve the existing quantized-domain fused activation behavior: ```text float op -> Q -> quantized-domain clip ``` The elemwise QDQ path applies fused activation before the final quantize: ```text DQ -> float binary op -> float fused activation -> Q ``` Both paths are intentional and covered by regression tests: - quantized concat fused `RELU` checks the quantized-domain clip path - quantized add fused `RELU6` checks the float-domain activation-before-Q path This PR also fixes a latent `R.clip` call-site bug in the quantized fused `RELU` helper by using `max=` rather than the unsupported `a_max=` keyword. ## Safety Checks - Quantized elemwise non-comparison outputs must have output qparams. Missing output quantization metadata now raises `OpAttributeInvalid` instead of silently returning a float result. - Per-channel quantization rejects non-zero per-axis zero points, following the TFLite quantization specification. - Per-channel depthwise convolution is explicitly unsupported rather than importing with an incorrect axis interpretation. ## Tests The new tests build minimal TFLite flatbuffers directly and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. Unsupported-boundary tests use `pytest.raises`. The FlatBuffer tests use schema module helpers instead of top-level generated builder functions when needed, so they work with the `tflite` Python package available in CI. | Test | Coverage | |---|---| | `test_tensor_quantization_parameters_are_parsed` | per-tensor and per-axis metadata parsing | | `test_quantize_op_uses_relax_quantize` | TFLite `QUANTIZE` float input | | `test_quantize_op_requantize_uses_dq_q` | TFLite `QUANTIZE` as requantize | | `test_dequantize_op_uses_relax_dequantize` | TFLite `DEQUANTIZE` | | `test_quantized_add_uses_qdq` | quantized ADD with differing input qparams | | `test_quantized_add_fused_relu6_uses_float_clip_before_quantize` | elemwise fused activation before Q | | `test_quantized_add_without_output_qparams_invalid` | invalid missing output qparams guard | | `test_quantized_conv2d_per_tensor_uses_qdq` | Conv2D per-tensor QDQ | | `test_quantized_conv2d_per_channel_weight_uses_remapped_axis` | Conv2D per-channel weight axis `0 -> 3` | | `test_quantized_conv2d_with_int32_bias_dequantizes_bias` | Conv2D INT32 bias scale | | `test_quantized_conv2d_per_channel_weight_with_int32_bias_dequantizes_bias` | Conv2D per-channel vector bias scale | | `test_quantized_concat_uses_qdq` | concat QDQ path | | `test_quantized_concat_fused_relu_uses_quantized_clip` | quantized-domain fused RELU clip | | `test_per_channel_depthwise_conv_unsupported` | per-channel depthwise guard | | `test_uint8_reshape_requantize_uses_dq_reshape_q` | uint8 reshape with different qparams | | `test_transpose_conv_with_int32_bias_dequantizes_bias` | TransposeConv INT32 bias DQ | | `test_quantized_fully_connected_with_int32_bias_dequantizes_bias` | FC INT32 bias DQ | Local validation: ```bash python -m ruff format --check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m pytest tests/python/relax/test_frontend_tflite.py -q ``` Result: ```text 433 passed ``` ## Limitations - This PR prioritizes correct import and explicit Relax IR over fused int8 kernel selection. The generated IR uses QDQ and float Relax operators. - Per-channel depthwise convolution remains unsupported. - The tests are structural IR tests. Numerical comparison against TFLite runtime outputs is left to follow-up work. ## References - Issue #19534: Support quantized TFLite import in Relax frontend - TFLite quantization spec: https://www.tensorflow.org/lite/performance/quantization_spec |
||
|
Tianqi Chen
|
729108cfc4 |
[REFACTOR][RELAX] Fold CalleeCollector into relax DeadCodeElimination (#19603)
## Summary The cross-IR `CalleeCollector` abstraction in `include/tvm/ir/analysis.h` had a single consumer (relax `DeadCodeElimination`) yet forced its per-language visitors to live in separate `analysis/` files registered via a runtime vtable. This PR folds both visitors (relax + tirx) directly into `src/relax/transform/dead_code_elimination.cc` as anonymous-namespace helpers and deletes the now-dead abstraction. The indirection only paid off when multiple unrelated passes shared the visitor; with one consumer, the cross-TU vtable adds compile cost and spreads the implementation across three files. Inlining improves locality without enlarging the consumer's complexity. |
||
|
Shushi Hong
|
cae6cb89b7 |
[IR] Add annotations to Call nodes (#19597)
This PR adds annotation support to `tirx.Call` so downstream codegen
users can attach call-level metadata and preserve it through TIRX
transforms.
What changed:
- Add `CallNode::annotations` and expose it through reflection.
- Add Python `tvm.tirx.Call(..., annotations=...)` support.
- Preserve call annotations in C++ and Python expression mutators.
- Preserve annotations across TIRX/arith passes that rebuild equivalent
calls.
- Print annotated calls as `Tx.Call(..., annotations={...})` and support
script roundtrip.
- Add regression coverage for annotated calls, mutator preservation,
script roundtrip, and simplify preservation.
This pr also cleans some stuff that #19596 didn't clean completely
|
||
|
Shushi Hong
|
59bfb21559 |
[CodeGen][CUDA] Move fast math intrinsic lowering option to PassContext (#19596)
This updates CUDA fast math intrinsic lowering to use a PassContext
option instead of a CUDA Target attribute.
The new option is:
```python
with tvm.transform.PassContext(config={"tirx.enable_fast_math": True}):
...
```
When unset or false, CUDA math intrinsics continue to lower to the
precise CUDA math functions such as expf. When true, tirx.LowerIntrin
prioritizes the cuda.fastmath.* lowering rules, producing fast math
intrinsics such as __expf.
|
||
|
Bl4ckSku11
|
a7463e9b2d |
[RPC][Tracker] Bound msg_size to MAX_TRACKER_MSG_BYTES to prevent unbounded buffer growth (#19586)
Fixes #<issue-number>. Reads of `_msg_size` from the tracker socket are now bounded to `MAX_TRACKER_MSG_BYTES = 1 MiB`, and the 4-byte size header is consumed at read time. Without these checks, a single TCP connection from a peer can grow the tracker process buffer until OOM, and a wire size of 0 starves the parser without ever freeing the bytes. Per the TVM security model the tracker is deployed on trusted networks, so this is filed as a robustness defect, not a security advisory. Apache security team triage (private thread, 2026-05-17) confirmed this is the right channel. ### Test Added regression test in tests/python/contrib/test_rpc_tracker.py that completes the magic handshake, sends an oversized msg_size header (0x7FFFFFFF), and asserts the tracker closes the connection. ### Changes - python/tvm/rpc/tracker.py: bound `_msg_size` to (0, MAX_TRACKER_MSG_BYTES], consume size header on read. - tests/python/contrib/test_rpc_tracker.py: regression test. |
||
|
Tianqi Chen
|
4052880e7c | [BUILD] Modularize device runtime into per-backend DSOs (#19594) | ||
|
hh
|
a1e4cd82fe |
[Relay/ONNX] Add RMSNormalization converter for ONNX opset 23 (#19590)
Add support for the ONNX RMSNormalization operator (opset 23) in the Relax ONNX frontend. This operator is essential for importing LLM models (LLaMA, Gemma, etc.) that use RMS normalization. The implementation: - Maps ONNX RMSNormalization to relax.op.nn.rms_norm - Supports the axis, epsilon, and stash_type attributes - Handles float16 inputs with stash_type=1 (compute in float32) - Includes unit tests comparing against ONNX Runtime |
||
|
HoYi
|
fff3b4bf0d |
[Relax][Frontend][TFLite] Support StableHLO region-based ops and multi-subgraph models (#19587)
## Summary This PR adds Relax TFLite frontend support for 10 additional StableHLO builtin operators from #19519 item I, building on the 29 ops merged in PR #19536. The first 5 ops are direct single-subgraph converters: `CBRT`, `REMAINDER`, `DYNAMIC_UPDATE_SLICE`, `DOT_GENERAL`, and `CONVOLUTION`. The remaining 5 ops are region/subgraph-based: `REDUCE`, `REDUCE_WINDOW`, `SORT`, `SCATTER`, and `COMPOSITE`. To support these, the TFLite frontend is extended to accept multi-subgraph models while still converting only `Subgraphs(0)` into the Relax main function. Region subgraphs are consumed by their parent op converters as needed. Relates to #19519. ## Changes 1. **Single-subgraph ops** - `CBRT` — sign-preserving composite expression: `where(x < 0, -power(-x, 1/3), power(x, 1/3))`. Float dtype only. - `REMAINDER` — truncating remainder via `x - y * trunc(x / y)`, matching StableHLO semantics (sign follows dividend). Float dtype only. - `DYNAMIC_UPDATE_SLICE` — static start indices + static shapes only, lowered to `R.scatter_nd` with a coordinate grid generated via `np.indices`. Runtime starts and out-of-bounds ranges raise `OpNotImplemented`. - `DOT_GENERAL` — canonical 2D matmul subset: no batching dims, `lhs_contracting=[1]`, `rhs_contracting=[0]`, lowered to `R.matmul`. - `CONVOLUTION` — canonical 2D NHWC/HWIO subset with `BatchGroupCount=1`, `FeatureGroupCount=1`, lowered to `R.nn.conv2d`. Non-canonical dimension numbers and grouped/depthwise conv raise `OpNotImplemented`. 2. **Multi-subgraph infrastructure** - Lift `from_tflite()` assertion from `model.SubgraphsLength() == 1` to `model.SubgraphsLength() >= 1`. Only `Subgraphs(0)` is converted into the Relax main function. - Limit `_input_type()` to `Subgraphs(0)` inputs, preventing region parameters from leaking as Relax main function parameters. - Add `_get_stablehlo_simple_body_op` helper for validating and extracting the single operator from a region body subgraph. - Extend test helper `_finish_tflite_model` with `extra_subgraphs` parameter for constructing multi-subgraph TFLite flatbuffers. 3. **Region/subgraph ops** - `REDUCE` — single-op reducer body subgraph. Supports `ADD` → `R.sum`, `MAXIMUM` → `R.max`, `MINIMUM` → `R.min`, `MULTIPLY` → `R.prod`. Init value must match the reducer identity element. - `SORT` — single-op comparator body subgraph. `LT` → ascending sort, `GT` → descending sort via `R.sort`. `IsStable` is not mapped. - `REDUCE_WINDOW` — NHWC 4D 2D-pooling subset with `MAXIMUM` reducer and identity init, lowered to `R.nn.max_pool2d`. BaseDilations must be all 1. - `SCATTER` — single-op update computation body subgraph. Supports `ADD`/`MAXIMUM`/`MINIMUM`/`MULTIPLY` → `R.scatter_nd` with the corresponding reduction mode. Only canonical point-update semantics (no window dims). - `COMPOSITE` — inlines a decomposition subgraph through a recursive `OperatorConverter` with an isolated `ExprTable`, so decomposition tensor bindings cannot overwrite main graph bindings. Only supports composites without `CompositeAttributes`. 4. **Not included** - `STABLEHLO_RESHAPE`, `STABLEHLO_TRANSPOSE`, and `STABLEHLO_SLICE` are left to another contributor. - `WHILE`, `CUSTOM_CALL`, and `RNG_BIT_GENERATOR` are deferred to follow-up PRs. 5. **Bug fix** - Fixed `DYNAMIC_UPDATE_SLICE` scatter_nd indices layout: `np.indices` returns `(rank, *update_shape)` but `scatter_nd` expects `(*update_shape, rank)`. Added `np.moveaxis` to transpose the coordinate axis from first to last position. ## Testing All tests use manually-built minimal TFLite flatbuffers with `tvm.ir.assert_structural_equal`. Region/subgraph tests construct the smallest valid body/comparator/update subgraphs. BuiltinOptions2 ops construct their options via the FlatBuffers schema API. ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k stablehlo -q ``` ## Result - 39 StableHLO operators registered in the Relax TFLite frontend (29 from PR #19536 + 10 from this PR). - 77 StableHLO test cases covering all registered ops, including structural-equal tests and unsupported/error-path checks: - `REMAINDER` truncating semantics - `DYNAMIC_UPDATE_SLICE` with dynamic starts and out-of-bounds starts - `DOT_GENERAL` with non-canonical contracting dimensions - `CONVOLUTION` with non-canonical dimension numbers and `FeatureGroupCount > 1` - `REDUCE` with unsupported reducer and non-identity init value - `SORT` with unsupported comparator and stable sort - `REDUCE_WINDOW` with unsupported reducer and base dilation - `SCATTER` with unsupported reducer and update window dims - `COMPOSITE` with composite attributes and scope isolation - Multi-subgraph model with unused subgraphs - All 77 StableHLO tests pass. ## References - Issue #19519 item I: StableHLO operators in TFLite - PR #19536: First batch of 29 StableHLO ops |
||
|
Javier De Jesus
|
1720d305ae |
[Relax][ONNX] Fix TopK scalar K extraction in from_onnx (#19573)
### Root Cause `TopK._impl_v11` extracted `k` with `int(k.data.numpy())`. ONNX emits `K` as a single-element 1-D tensor constant, so `numpy()` returns a 1-D array and `int()` raises `TypeError: only 0-dimensional arrays can be converted to Python scalars`, failing conversion of any model with a `TopK` node. ### Solution Resolve `k` with `get_constant(inputs[1], params)` and extract the scalar with `.item()`, matching the `Trilu` and `Reshape` converters in the same file. `get_constant` also handles `k` arriving as a parameter when `keep_params_in_input=True`. ### Test Plan `test_topk` in `tests/python/relax/test_frontend_onnx.py` already builds `K` as a single-element 1-D INT64 constant, so it exercises this path. `.item()` returns the scalar for both single-element 1-D and 0-d constants. ### Issue Fixes #19571 |
||
|
ConvolutedDog
|
48f346bb07 |
[RFC][CodeGen][CUDA]: Gate fast math intrinsic lowering behind target option (#19565)
Fix CUDA lowering of standard TIR math intrinsics so they use precise CUDA math functions by default instead of fast-math `__*f` functions. This fixes the default behavior reported in #19546, where operators such as `tirx.exp` could lower to `__expf` even though fast math was not explicitly requested. This change adds a CUDA target attribute, `enable_fast_math`, which defaults to `false`. When the attribute is unset or false, standard math intrinsics lower through the normal CUDA math rule, for example `expf`, `logf`, `sinf`, `cosf`, `powf`, and `rsqrtf` for `float32`. When users explicitly enable the attribute on the target, the lowering pass also checks the `cuda.fastmath.FLowerIntrinsic` rules before the normal CUDA lowering rules. Users can opt in to fast math by constructing a CUDA target with the attribute: ```py tvm.target.Target({"kind": "cuda", "enable_fast_math": True}) target = tvm.target.Target({ "tag": "nvidia/nvidia-a100", "enable_fast_math": True, }) ``` The fast-math lowering path currently covers the CUDA math operators registered with `cuda.fastmath.FLowerIntrinsic`: `tirx.exp`, `tirx.exp10`, `tirx.log`, `tirx.log2`, `tirx.log10`, `tirx.tan`, `tirx.cos`, `tirx.sin`, `tirx.tanh`, and `tirx.pow`. `tirx.rsqrt` is also registered for CUDA lowering so it maps to the CUDA reciprocal-square-root intrinsic instead of being legalized as `1 / sqrt(x)`. Add CUDA codegen tests `tests/python/codegen/test_target_codegen_cuda_fastmath.py` that check the lowered IR, generated CUDA source, and runtime results for the supported math intrinsics across floating point dtypes and both default and fast-math targets. |
||
|
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. |
||
|
Neo Chien
|
bc1a904ec1 |
[Relax][ONNX] Prevent Div divide-by-zero crashes (#19566)
Hi Committers, This PR is trying to fix issues #19541. Any suggestions would be appreciated if you are available. ### Root cause: The ONNX `Div` path in the Relax frontend did not separate two integer-divisor cases: constant zero divisors and dynamic/unknown divisors. As a result, constant integer zero divisors were not rejected during import, and dynamic integer divisors could reach runtime without a guard. When the divisor became zero at runtime, execution could trigger SIGFPE and terminate the process instead of raising a controlled error. ### Solution: This PR applies a minimal, targeted fix in the ONNX frontend `Div` conversion path. It introduces: import-time validation for constant integer divisors containing zero, raising ValueError early. --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |
||
|
Neo Chien
|
bedfcb2b85 |
[Relax][ONNX] Set max_output_boxes_per_class default value to 0 for NonMaxSuppression (#19547)
Hi Committers, This PR is trying to fix issues #19544. Any suggestions would be appreciated if you are available. --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |
||
|
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 |
||
|
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. |
||
|
Sun
|
378c4f3043 |
[BugFix][Relax]: handle ONNX ScatterElements reduction (#19527)
### Summary - Respect the ONNX `reduction` attribute in the Relax ONNX frontend `ScatterElements` converter. - Preserve existing default behavior by mapping missing reduction and ONNX `none` to Relax `update`. - Add focused regression coverage for opset 11 default behavior, opset 16 `add`/`mul`, and opset 18 `none`/`min`/`max`. ### Changes - Added a shared helper to normalize and validate ONNX reduction attributes. - Implemented `ScatterElements` opset 16 and opset 18 converters. - Reused the existing `relax.op.scatter_elements(..., reduction=...)` API. - Reused the same reduction helper in `ScatterND` to keep behavior consistent. ### Test Plan - `python -m py_compile python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.py` - `python -m pytest tests/python/relax/test_frontend_onnx.py::test_gather_elements tests/python/relax/test_frontend_onnx.py::test_scatter tests/python/relax/test_frontend_onnx.py::test_scatter_elements_reduction tests/python/relax/test_frontend_onnx.py::test_scatter_nd -q` ### Issue Fixes #19435 ## Local Verification Notes - WSL conda environment: `/home/thinker/.cache/tvm-conda-onnx` - TVM build directory: `/home/thinker/.cache/tvm-build-onnx` - LLVM runtime check: `tvm.runtime.enabled("llvm") == True` - Relevant ONNX frontend subset: `15 passed, 4 skipped, 2 warnings` - Full `tests/python/relax/test_frontend_onnx.py` was also attempted. It currently has 14 failures in unrelated `Reduce* axes input` and `TopK` tests; running the same selected failures against `origin/main` reproduces them, so they are not introduced by this PR. |
||
|
HoYi
|
c0406a54bc |
[Relax][Frontend][TFLite] Add initial StableHLO builtin operator support (#19536)
## Summary This PR adds initial Relax TFLite frontend support for 29 StableHLO builtin operators from #19519 item I. The covered subset includes pure elementwise ops, BuiltinOptions2 / metadata-based ops, simple shape-manipulation ops, and a take-equivalent subset of `STABLEHLO_GATHER`. StableHLO builtins carry no TFLite-specific quantization or fused-activation metadata, so the implementation uses dedicated converter helpers that bypass the existing TFLite elemwise/QNN code paths. Relates to #19519. ## Changes 1. **Zero-attribute elementwise helpers** - Add `_convert_stablehlo_unary`, `_convert_stablehlo_binary`, and `_convert_stablehlo_ternary` for pure elementwise mapping. - Register 20 ops: unary (`ABS`, `NEGATE`, `COSINE`, `EXPONENTIAL`, `FLOOR`, `LOG`, `LOGISTIC`, `RSQRT`, `TANH`), binary (`ADD`, `SUBTRACT`, `MULTIPLY`, `DIVIDE`, `MAXIMUM`, `MINIMUM`, `POWER`), ternary (`SELECT` → `R.where`), and dtype-dispatched bitwise/logical ops (`AND` / `OR` → logical ops for bool or bitwise ops for integer, `SHIFT_LEFT` → `R.left_shift` for integer). 2. **BuiltinOptions2 infrastructure** - Add `_get_stablehlo_options` helper for parsing `BuiltinOptions2` flatbuffers with enum validation via `getattr(BuiltinOptions2, options_cls.__name__)`. - Register 6 ops: `CONVERT` → `R.astype`, `CLAMP` → `R.minimum(R.maximum(...))`, `CONCATENATE` → `R.concat`, `BROADCAST_IN_DIM` → `R.reshape` + `R.broadcast_to`, `IOTA` → `R.arange` + `R.broadcast_to`, and `COMPARE` → 6 comparison directions (`TOTALORDER` raises `OpNotImplemented`). 3. **Shape-manipulation ops** - `PAD` → `R.nn.pad` in constant mode. The initial PAD path supports non-negative edge padding with zero interior padding and a constant scalar padding value. Interior padding, negative padding, and dynamic padding values raise `OpNotImplemented`. - `DYNAMIC_SLICE` → `R.dynamic_strided_slice`. The initial path supports constant, in-bound start indices only. Runtime start indices and out-of-bounds StableHLO clamping semantics are deferred. 4. **Indexing op** - `GATHER` → `R.take` for the take-equivalent subset only. - Parses the relevant `StablehloGatherOptions` attributes needed to validate this subset: `offset_dims`, `collapsed_slice_dims`, `start_index_map`, `index_vector_dim`, and `slice_sizes`. - Validates the gather axis, collapsed dims, offset dims, slice sizes, and output shape against the expected `R.take` layout. Multi-dimensional and non-take-equivalent gather patterns raise `OpNotImplemented`. 5. **Not included** - `STABLEHLO_RESHAPE`, `STABLEHLO_TRANSPOSE`, and `STABLEHLO_SLICE` are left to another contributor who expressed interest in those ops. - The remaining Issue #19519 StableHLO items are deferred to follow-up PRs: `CBRT`, `REMAINDER`, `SCATTER`, `CONVOLUTION`, `DOT_GENERAL`, `REDUCE`, `REDUCE_WINDOW`, `DYNAMIC_UPDATE_SLICE`, `COMPOSITE`, `CUSTOM_CALL`, `RNG_BIT_GENERATOR`, `SORT`, and `WHILE`. - More general or multi-dimensional `STABLEHLO_GATHER` patterns are also deferred to follow-up work. ## Testing All tests use manually-built minimal TFLite flatbuffers with `tvm.ir.assert_structural_equal`. BuiltinOptions2 ops construct their options via the FlatBuffers schema API, modeled after the existing DILATE test pattern. ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k stablehlo -q ``` ## Result - 29 StableHLO operators registered in the Relax TFLite frontend. - 44 StableHLO test cases covering all registered ops, including structural-equal tests and unsupported/error-path checks: - `COMPARE` with `TOTALORDER` - `PAD` with interior padding, negative padding, and dynamic padding values - `DYNAMIC_SLICE` with runtime starts and out-of-bounds starts - non-take-equivalent or multi-dimensional `GATHER` - All StableHLO TFLite frontend tests pass locally. ## References - Issue #19519 item I: StableHLO operators in TFLite - Related PR #19481: DILATE operator mapping, the first use of BuiltinOptions2 in the TFLite frontend tests |
||
|
Yichen Yan
|
c309e4ea5f |
[TIR] Add cooperative_tensor builtins and metal.cooperative_tensor storage scope (#19423)
part of https://github.com/tile-ai/tilelang/pull/1869 ## Summary Add TIR builtins and storage scope for Metal cooperative_tensor operations (MetalPerformancePrimitives / Metal 4). ## Motivation Apple Metal 4 introduces MetalPerformancePrimitives (MPP) with `matmul2d` using `cooperative_tensor` operands. On M5, this routes to NAX tensor cores; on M1-M4, it falls back to simdgroup matrix instructions. These TIR primitives enable backend codegen to emit MPP calls. ## Changes ### New TIR builtins - `cooperative_tensor_fill(d, index, value, rows, cols)` - `cooperative_tensor_load(d, index, ptr, stride, rows, cols, transpose)` - `cooperative_tensor_store(d, index, ptr, stride, rows, cols, transpose)` - `cooperative_tensor_multiply_accumulate(d, di, a, ai, b, bi, c, ci, M, N, K, trans_a, trans_b)` ### New storage scope - `metal.cooperative_tensor` (`StorageRank::kMetalCooperativeTensor`) ### Files changed - `include/tvm/tirx/builtin.h` — Op declarations - `src/tirx/op/builtin.cc` — Op registrations - `python/tvm/tirx/op.py` — Python wrappers - `python/tvm/script/ir_builder/tirx/ir.py` — Script parser exports - `src/runtime/thread_storage_scope.h` — StorageRank enum + scope parsing These builtins mirror the existing `simdgroup_*` builtins for the older Metal simdgroup matrix API, extended with M/N/K dimension parameters for the matmul2d descriptor. |
||
|
Wei-Cheng Hsu
|
2c76c7955a |
[Relax][Frontend] Add TFLite Frontend Support for CONV_3D_TRANSPOSE (#19530)
This commit adds support for the CONV_3D_TRANSPOSE operator in the Relax TFLite frontend. Key implementations: - Registered CONV_3D_TRANSPOSE to the TFLite op map. - Implemented convert_conv3d_transpose which shares Conv3DOptions with regular Conv3D but handles the distinct tensor input layout [output_shape, weight, data, bias] and the DHWOI kernel layout. - Added calculation for SAME padding that correctly handles transposed convolution semantics, computing padding and output_padding based on dilated kernel and stride sizes. - Added comprehensive unit tests for valid and same padding in test_frontend_tflite.py. Testing: - `python3 -m pytest tests/python/relax/test_frontend_tflite.py -k "test_conv3d_transpose"` Related to: https://github.com/apache/tvm/issues/19519 |
||
|
Neo Chien
|
e370fc7374 |
[Relax][ONNX] Normalize negative indices before the take call for Gather operator (#19525)
Hi Committers, This PR is trying to fix issues https://github.com/apache/tvm/issues/19436. Any suggestions would be appreciated if you are available. ### Root Cause 1. ONNX `Gather` allows negative indices (counting from the end of the target axis). 2. In the Relax ONNX importer, `Gather` was lowered directly to `relax.op.take` without normalizing negative indices first. 3. This created semantic mismatch / incorrect behavior in downstream lowering paths that assume non-negative indices. 4. Test failures were also caused by pytest parametrization issues: - using ONNX `TensorProto` enum values directly as NumPy dtypes, - and tuple-style parametrization triggering fixture interpretation errors. ### Solutions 1. Added conditional negative-index normalization in `Gather._impl_v13`: - apply only for signed index dtypes, - use: `idx < 0 ? idx + axis_extent : idx`, - derive `axis_extent` from shape/runtime expression to support dynamic shapes. 2. Skipped normalization for unsigned index dtypes to avoid redundant graph ops/checks. --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |
||
|
Wei-Cheng Hsu
|
4ab312bf12 |
[Relax][Frontend][TFLite] Add Conv3D support (#19523)
Description
This PR adds support for the CONV_3D operator in the TFLite frontend for
Relax.
Key Changes
- Operator Mapping: Added CONV_3D to the OperatorConverter mapping in
tflite_frontend.py.
- Implementation:
- Implemented convert_conv3d to handle 3D convolution attributes such as
StrideD/H/W, DilationD/H/W, and Padding.
- Correctly handled the TFLite 3D kernel layout, which is expected to be
DHWIO (Depth, Height, Width, Input Channels, Output Channels).
- Integrated support for fused activation functions (ReLU, ReLU6, etc.)
directly following the convolution.
- Unit Tests:
- Added comprehensive tests in
tests/python/relax/test_frontend_tflite.py covering:
- VALID and SAME padding modes.
- Various stride and dilation configurations.
- Verification against expected Relax IR structure.
Testing:
- `python3 -m pytest tests/python/relax/test_frontend_tflite.py -k
"test_conv3d"`
Notes for Reviewers
The implementation follows the existing pattern used for CONV_2D but
extends it to the 5D case (NDHWC layout). I've ensured that the kernel
layout mapping aligns with TVM's R.nn.conv3d requirements.
Related to: https://github.com/apache/tvm/issues/19519
|
||
|
Neo Chien
|
c12debb462 |
[Relax][PyTorch] Fix segfault in from_exported_program when model uses index_put_ with tuple output (#19488)
Hi Committers, This PR is trying to fix issues https://github.com/apache/tvm/issues/18363. Any suggestions would be appreciated if you are available. ### Root Cause - When an ExportedProgram's FX graph output node returns a **nested Python tuple** (e.g., buffer mutation outputs + user-defined tuple returns), `_translate_fx_graph()` passes the raw nested structure directly to the Relax FFI Tuple constructor. - The C++ Array<Expr> initializer cannot handle heterogeneous/nested Python containers, causing a segmentation fault at `expr.cc`. - Additionally, index_put_ (in-place write op) did not update self.env to alias the source tensor to the mutated output, causing subsequent FX nodes that read the same tensor to observe **stale pre-mutation values**. ### Solution - exported_program_translator.py - Added static method `_flatten_output_args()` that recursively walks any Python `tuple/list`, collects only `relax.Expr` leaves, and preserve explicit None outputs as Relax null objects. - Replaced the fragile `assert isinstance(output_args, tuple | relax.Tuple)` guard with a call to `_flatten_output_args()`, producing a clean flat tuple of `relax.Expr` before FFI construction. - base_fx_graph_translator.py - In `_index_put()`, after emitting the `relax.op.index_put(...)` call, added an env alias update: `self.env[source_node] = output` when the target op name starts with `index_put_`, preserving correct in-place mutation semantics for downstream FX nodes. --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |
||
|
Wei-Cheng Hsu
|
e6538517b0 |
[Relax][TFLite] Add gather frontend expected IRModule tests (#19516)
This adds explicit Expected IRModule coverage for TFLite GATHER and GATHER_ND frontend conversion. GATHER_ND uses Relax gather_nd with int64 indices, so the frontend now casts int32 TFLite indices to int64 before emitting the Relax op. This keeps the generated module well-typed and matches the expected Relax IR. Testing: - `python -m pytest tests/python/relax/test_frontend_tflite.py -k "gather"` related to https://github.com/apache/tvm/issues/18971 |
||
|
Soowon Jeong
|
5a7da7a32a |
[BugFix][Relax][Torch] Honor correction in std/var converter (#19512)
## Motivation The PyTorch frontend's `_var` ignored the `correction` kwarg of `aten.var.correction`. `torch.export.run_decompositions()` rewrites both `aten.std.correction` and `aten.std.dim` into `aten.var.correction(..., correction=<value>) → sqrt`, so every `torch.std`/`torch.var` call lands in `_var` — but the correction value was dropped on the floor. The variance was therefore always divided by `n` regardless of what the user requested. Minimal repro (vs PyTorch eager): ``` x = [[1, 2, 3, 4, 5], [2, 2, 2, 2, 2]] torch.std(x, dim=1, unbiased=True) ref: [1.5811, 0.0] # sqrt(2.5) tvm: [1.4142, 0.0] # sqrt(2.0) -- correction silently set to 0 ``` The same omission shows up for explicit `torch.var(x, correction=k)` and any model that relies on the documented Bessel default. ## Fix Route `aten.var.correction` (identified by `OpOverload._overloadname`, not a substring match) to a new `_var_correction` helper. It reads `correction` from `node.kwargs`, treats `None` as 1 to match the overload's `Scalar? correction = None` schema, and scales the existing `relax.op.variance` output by `n / (n - correction)` when `correction != 0`. When `n - correction <= 0`, the multiplier is set to NaN rather than raising — this mirrors PyTorch's documented `max(0, N - correction)` semantics (eager produces NaN with a warning, not an error). Reduction-axis sizes are read from `x.struct_info.shape`. Dynamic sizes raise `NotImplementedError`; static-shape models cover the real-world `torch.export` flow. The legacy fx path through `_var` is intentionally left alone — it has a separate preexisting bug (it reads `args[2]` as `keepdim` even when that slot is `unbiased`), but fixing that here would expand the scope of this PR beyond the `correction` semantics. ## Notes - `_std` is also registered for `"std.correction"` but is unreachable on the default exported-program path because `aten.std.*` always decomposes to `var.correction + sqrt` before dispatch. Sparse-tensor exports that skip `run_decompositions` still hit the old `_std`; that path is out of scope for this fix. - Existing `test_std`/`test_var` encoded the buggy `correction=0` IR for `torch.var(x)` (which defaults to Bessel) and have been updated to expect the correct `R.multiply(var, R.const(15/14))`. New `test_var_correction` covers explicit `correction=2` and `correction=0`. |
||
|
Soowon Jeong
|
61b49bb3f9 |
[BugFix][Relax][Torch] Honor multi-axis dims in torch.flip converter (#19511)
## Motivation
PyTorch's `torch.flip(x, dims=[...])` reverses every listed axis. The
Relax converter `_flip` (`base_fx_graph_translator.py`) instead coerces
the list to a single integer:
```python
if isinstance(dims, list | tuple) and len(dims) > 0:
dims = dims[0]
```
Only the first axis is forwarded to `relax.op.flip`, which is itself
single-axis. The remaining axes are silently dropped.
Minimal repro (vs PyTorch eager) on a `(3, 4)` input with
`dims=[-1, -2]`:
```
ref: [11, 10, 9, 8, 7, 6, 5, 4, ...] # both axes flipped
tvm: [ 3, 2, 1, 0, 7, 6, 5, 4, ...] # only last axis flipped
```
max_abs_diff = 8.0. Both the `torch.export` and legacy fx paths share
this converter, so both are affected.
## Fix
Iterate over `dims` in the converter and emit one `relax.op.flip` per
axis (flips along distinct axes commute, so the order is irrelevant).
A scalar `dims` is wrapped to a single-element list; non-int /
non-sequence arguments still raise `TypeError`.
`relax.op.flip` itself is unchanged: it is used elsewhere as a
single-axis op, and widening its signature would expand the scope of
this fix beyond the PyTorch frontend.
|
||
|
Soowon Jeong
|
82a37dac1d |
[BugFix][Relax][ONNX] Resolve param Vars in Concat to handle mixed Shape/Tensor inputs (#19498)
## Description When `from_onnx(model, keep_params_in_input=True)` is used, every ONNX initializer becomes a `relax.Var` instead of a `relax.Constant`. The `Concat` handler's `is_shape_like()` check only recognizes `relax.ShapeExpr` and 1D-int64 `relax.Constant`, so a 1D-int64 shape value loaded as a Var is no longer recognized. When such a Var is concatenated with a `ShapeExpr` — the standard pattern for dynamic-batch `Reshape` in PyTorch-exported ONNX models — the heterogeneous `Tuple(ShapeExpr, Tensor)` is rejected by `relax.op.concat` with: ``` InternalError: Op(relax.concat) expects the input to be a Tuple of Tensors. However, the given input is R.Tuple(R.Shape([N]), R.Tensor((1,), dtype="int64")) ``` This effectively breaks `keep_params_in_input=True` for any model with dynamic-batch `Reshape` (extremely common in PyTorch ONNX exports). ## Fix Run each `Concat` input through the existing `get_constant` helper before the `is_shape_like` check. This resolves any `Var` that maps to a known param back to its baked `Constant`, restoring the all-shape-like fast path. ## Minimal repro An 8-node ONNX graph (`Shape` → `Slice` → `Concat([dyn_n, [12]])` → `Reshape`) fails with `keep_params_in_input=True` before this PR and passes after. A regression test (`test_concat_with_param_shape_value`) covers this pattern. ## Testing ``` pytest tests/python/relax/test_frontend_onnx.py -k concat ``` 9 passed (1 new + 8 existing). |
||
|
as4230
|
fde09d2052 |
[BugFix][Relax] Fix scatter_elements and scatter_nd CUDA compilation (#19497)
`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.
|
||
|
Bana
|
87bf3022b7 |
[Relax][Frontend][TFLite] Add tests coverage for SPACE_TO_BATCH_ND and BATCH_TO_SPACE_ND (#19499)
**Changes** Add tests in `test_frontend_tflite.py`. Lower S`PACE_TO_BATCH_ND` / `BATCH_TO_SPACE_ND` through TOPI in `tflite_frontend.py`. Use tf.raw_ops.BatchToSpaceND in the test because tf.batch_to_space_nd is not available in this TF build. **Why the TFLite frontend changed** The frontend was calling relax.op.nn.space_to_batch_nd / relax.op.nn.batch_to_space_nd, which aren’t implemented in this checkout. I updated the TFLite frontend to lower these ops via TOPI packed calls so conversion works and the new tests can pass. **Test:** ``` pytest test_frontend_tflite.py -k "test_space_to_batch_nd or test_batch_to_space_nd" ``` related to #18971 |
||
|
HoYi
|
8873a4c8a5 |
[Relax][Frontend][TFLite] Add segment operator mappings (#19491)
## Summary This PR adds Relax TFLite frontend support for the following segment operators from #19412: - `SEGMENT_SUM` - `UNSORTED_SEGMENT_MIN` - `UNSORTED_SEGMENT_PROD` These operators are lowered through `relax.op.scatter_nd` with the corresponding reduction modes. ## Changes ### TFLite Frontend 1. Add TFLite converter mappings for segment operators: - `SEGMENT_SUM` -> `scatter_nd(..., reduction="add")` - `UNSORTED_SEGMENT_MIN` -> `scatter_nd(..., reduction="min")` - `UNSORTED_SEGMENT_PROD` -> `scatter_nd(..., reduction="mul")` 2. Add shared segment lowering logic: - Convert `segment_ids` into scatter indices via `expand_dims`. - Build the output shape from `num_segments` or constant `segment_ids`. - Initialize the scatter base tensor with the correct reduction identity. ### Tests Add TFLite frontend tests for: - `test_segment_sum` - `test_unsorted_segment_min` - `test_unsorted_segment_prod` Each test verifies the imported Relax IR lowers to `R.scatter_nd` with the expected reduction mode and base tensor initialization. ## Testing All targeted tests pass: ```bash python -m pytest \ tests/python/relax/test_frontend_tflite.py::test_scatter_nd \ tests/python/relax/test_frontend_tflite.py::test_segment_sum \ tests/python/relax/test_frontend_tflite.py::test_unsorted_segment_min \ tests/python/relax/test_frontend_tflite.py::test_unsorted_segment_prod \ -q ``` ## References - Issue #19412: TFLite Relax frontend operator support tracking - Related PR #19490: Adds SCATTER_ND support |
||
|
Masahiro Hiramori
|
86794e7d91 |
[Relax][Frontend] Add ParameterList and ParameterDict containers (#19495)
This PR adds first-class `nn.ParameterList` and `nn.ParameterDict` containers to the Relax frontend. These containers provide PyTorch-like list/dict registration for raw `nn.Parameter` objects while preserving Relax frontend semantics: values must be explicit `nn.Parameter` instances, with no automatic tensor-to-parameter conversion. ### Changes - Add public `nn.ParameterList` and `nn.ParameterDict` exports. - Support stable parameter names in traversal: - `params.0`, `params.1` - `params.foo`, `params.bar` - Integrate the new containers with: - `named_parameters()` - `parameters()` - `state_dict()` - `load_state_dict()` - `to(dtype=...)` - `export_tvm()` - `nn.Mutator` - Add focused tests for basic container behavior, type validation, nested traversal, export parameter names, state loading, dtype conversion, and mutator naming. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> |
||
|
ysh329
|
a354b4f59e | [release] Update version to 0.25.dev0 on main branch | ||
|
ysh329
|
af3e4ba814 | [release] Update version to 0.24.0 on main branch | ||
|
as4230
|
bfa07828a6 |
[BugFix][Relax] Add legalize for isnan, isinf, isfinite (#19492)
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.
|
||
|
Bana
|
6157f49205 |
[Relax][Frontend][TFLite] Add BROADCAST_TO, EMBEDDING_LOOKUP, and SELECT_V2 (#19489)
This PR adds support for three new operators in the Relax TFLite frontend:` BROADCAST_TO`, `EMBEDDING_LOOKUP,` and `SELECT_V2.` Passed all newly added unit tests using ``` pytest tests/python/relax/test_frontend_tflite.py -k "test_broadcast_to or test_embedding_lookup or test_select_v2" ``` reference #19412 |
||
|
Bana
|
fbbbae994d |
[Relax][Frontend][TFLite] Add SCATTER_ND operator for Relax TFLite (#19490)
This PR adds support for the `SCATTER_ND` operator in the Relax TFLite frontend. ### Key Changes: - Added handler `convert_scatter_nd` to parse `indices` and `updates`. - Explicitly handles static vs dynamic shape tensor extraction via `to_int_list` and `relax.op.tensor_to_shape`. - Uses `relax.op.zeros` to initialize the base array based on the `updates` precision dtype. - Mapped `SCATTER_ND` to the corresponding `relax.op.scatter_nd()` target. - Registered the translator into `convert_map` and provided the matching unit test in test_frontend_tflite.py. ### Testing: Passed unit tests ``` pytest tests/python/relax/test_frontend_tflite.py::test_scatter_nd ``` Related to #19412 |
||
|
Neo Chien
|
6569cf0ac9 |
[Relax][ONNX] Fix CumSum axis handling: support runtime axis tensor, error on multi-element axis (#19467)
Hi Committers, This PR is trying to fix issues https://github.com/apache/tvm/issues/19437. Any suggestions would be appreciated if you are available. ### Root Cause The original CumSum converter always defaulted to axis=0 when the axis input was a relax.Var (i.e., a runtime tensor), ignoring the actual runtime value. This led to incorrect behavior and did not comply with the ONNX specification. ### Solutions Update CumSum._impl_v14 to: - Check if the axis input is a Constant: require it to have exactly one element, otherwise raise an error. - If the axis input is a relax.Var, raise an error instead of always defaulted to axis=0. --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |