219f1d83cf64e663cbfd6c6c91af91e5ff90d607
6367 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Shushi Hong
|
219f1d83cf |
[Refactor][Meta-schedule] Remove meta-schedule as_string mechanism in favor of default representation (#19709)
Python-side meta-schedule classes (`PyCostModel`, `PyFeatureExtractor`, `PyMeasureCallback`, `PyScheduleRule`, `PyMutator`, `PyPostproc`) carried an `f_as_string` callback whose only purpose was to produce a repr-style string (`s_tir.meta_schedule.<SubclassName>(0x...)`) for `str(...)`. This mechanism stopped working after #19461 migrated `ReprPrinter` to the tvm-ffi `__ffi_repr__` mechanism and intentionally removed the per-type `set_dispatch<Py*Node>` hooks that called back into `f_as_string`, which broke three `*_as_string` tests: - `test_meta_schedule_cost_model.py::test_meta_schedule_cost_model_as_string` - `test_meta_schedule_feature_extractor.py::test_meta_schedule_feature_extractor_as_string` - `test_meta_schedule_measure_callback.py::test_meta_schedule_measure_callback_as_string` Rather than restoring the old behavior, this PR removes the mechanism entirely: the string it produced is just a repr, and tvm-ffi reflection already provides an auto-generated default repr for every object. Keeping a dedicated Python → FFI → C++ callback chain alive only to reproduce that is not worth the complexity. |
||
|
Javier De Jesus
|
0e3250795b |
[Relax][PyTorch] Cast non-bool inputs to bool in logical_and converter (#19679)
### Motivation `torch.logical_and` accepts input tensors of any dtype (treating any nonzero element as `True`) and always returns a `bool` tensor. The PyTorch frontend did not produce that `bool` result. The ExportedProgram frontend lowered `logical_and.default` with `self._binary_op(relax.op.logical_and, operator.and_)`, which kept the operand dtype and emitted `relax.op.logical_and` on non-bool inputs (for example `float32`). `relax.op.logical_and` requires boolean inputs and otherwise fails `LegalizeOps` in the TOPI `logical_and`. The FX frontend did not register `logical_and` at all, so the op was unsupported there. ### Changes - Add a shared `_logical_and` converter in `BaseFXGraphImporter` that casts non-bool operands to `bool` before applying `relax.op.logical_and`. Bool operands are passed through unchanged (no redundant cast). - Point the `logical_and.default` (ExportedProgram) registration at the new converter, and add a `logical_and` (FX) registration that was previously missing, matching the existing `logical_not` converter. - Add a standalone `test_logical_and` to both the FX and ExportedProgram test suites asserting the corrected IR (`astype` to bool on each operand, then `logical_and`, producing a `bool` output). ### Notes The cast to `bool` lowers to an elementwise nonzero test, so it matches PyTorch's "nonzero is True" semantics for float, integer, and NaN inputs. |
||
|
Shushi Hong
|
b172d5ea32 |
[Arith] Make Analyzer a tvm-ffi Object (#19675)
This PR makes `arith::Analyzer` a first-class tvm-ffi object. The implementation splits the previous concrete `Analyzer` class into: - `AnalyzerObj`, the mutable object node that owns analyzer state, sub-analyzers, caches, and bindings - `Analyzer`, a reference-counted `ObjectRef` handle that can be passed across the tvm-ffi boundary This allows Python and C++ to share the same analyzer instance, so bindings, constraints, and cached facts can persist across FFI calls. Public APIs that accept an analyzer now use `const arith::Analyzer&`, while internal helper APIs that only borrow the object continue to use `AnalyzerObj*`. --------- Co-authored-by: Ubospica <ubospica@gmail.com> |
||
|
Neo Chien
|
aa59644072 |
[Relax][ONNX] Preserve NaN in Sign to align with ONNX Runtime (#19674)
Hi Committers, This PR fixes issues https://github.com/apache/tvm/issues/19543. Any suggestions would be appreciated if you are available. ### Root cause: The ONNX frontend `Sign` converter directly returned `relax.op.sign(x)`. After legalization, this maps to `topi.sign`, which is implemented via comparisons (x < 0 ? -1 : x > 0 ? 1 : 0). For `NaN`, both comparisons are false, so TVM produced 0, while ONNX Runtime preserves NaN. This created a frontend semantic mismatch for imported ONNX models. ### Solution: Apply a minimal ONNX-frontend-only fix in `onnx_frontend.py`: - For floating-point inputs, lower `Sign` as `where(isnan(x), x, sign(x))`. - Keep non-floating inputs unchanged (`sign(x)`). --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |
||
|
Bohan Hou
|
9db74c7cee |
[TIRx] Update scoped ops and CUDA launch bounds (#19677)
## Summary - replace the block-structured TIRx exec-scope surface with scope-qualified `Tx.<scope>.<op>` namespaces and migrate call sites - split TIRx op namespaces and remove the unused dynamic generic-op fallback - add explicit CUDA launch bounds plumbing through TIRx attrs and split-host-device lowering ## Validation - `git diff --check apache/main..HEAD` - `pre-commit run --from-ref apache/main --to-ref HEAD` |
||
|
Neo Chien
|
4d9d129c93 |
[Relax][ONNX] Fix Cast operator float->int NaN/Inf handling (#19626)
Hi Committers, This PR is trying to fix issues #19542. Any suggestions would be appreciated if you are available. ### Root cause: FP to INT lowering can be implementation-defined or UB for NaN/Inf and extreme floats, producing backend-dependent results versus ONNX Runtime. ### Solution: Apply a minimal, deterministic frontend sanitization for float to integer Casts: map NaN and ±Inf to 0.0 before astype. This prevents NaN/Inf from reaching backend fptosi/fptoui lowers and yields stable behavior across targets. --------- Co-authored-by: cchung100m <cchung100m@users.noreply.github.com> |
||
|
Tianqi Chen
|
1240649257 |
[FFI][REFACTOR] Direct structural APIs to tvm-ffi (#19661)
## Summary Python callers should reach the canonical tvm-ffi structural helpers directly instead of going through a TVM-side redirect layer. This makes the public tvm.ir bindings exact aliases of the tvm_ffi APIs and exposes get_first_structural_mismatch from tvm.ir. Main changes: - Import structural_equal, get_first_structural_mismatch, and structural_hash directly from tvm_ffi - Remove the pure wrappers from tvm.ir.base while keeping assert_structural_equal's TVM-specific formatting - Update mismatch tests and add identity coverage for the direct bindings |
||
|
Shushi Hong
|
a72d57c616 |
[CI] Derive the version from Git tags via setuptools_scm (#19665)
Replace manual version.py stamping with scikit-build-core's setuptools_scm metadata provider, so local builds no longer call version.py. The Python distribution/runtime version comes from the generated python/tvm/_version.py (libinfo.py reads it with a fallback); the C++ TVM_VERSION is injected from SKBUILD_PROJECT_VERSION_FULL with a #ifndef default in base.h for bare cmake builds. version.py is removed. The publish workflow's wheel build checks out full history (fetch-depth: 0) so setuptools_scm can derive the version, and drops the version.py stamping step. release_process.rst is updated to the tag-driven release flow. |
||
|
Javier De Jesus
|
57912395a8 |
[Relax][PyTorch] Decompose integer pow into repeated multiplication (#19660)
`torch.pow` on an integer tensor returns an integer result, but the PyTorch frontend lowered it to `relax.op.power`, which fails `LegalizeOps` with `power only applies to float` (TOPI `power` / `tvm::pow` requires a floating-point input). This decomposes an integer base raised to a constant non-negative integer exponent into repeated multiplication, so the result stays integral and matches PyTorch. Float bases and non-constant or tensor exponents keep using `relax.op.power` unchanged. The ONNX frontend already uses the same decomposition (`x**3 = x*x*x`). Added structural tests covering both the FX and ExportedProgram import paths. Fixes #19550 |
||
|
Tianqi Chen
|
1382707e8f |
[FFI][IR] Route JSON serialization through tvm-ffi (#19662)
TVM can rely on tvm-ffi's JSON graph serialization helpers directly instead of routing through TVM-side `node.SaveJSON`/`node.LoadJSON` registry entries. This changes `tvm.ir` save/load to call `tvm_ffi.serialization` with `tvm_version` metadata, removes the C++ registry wrapper, and moves the disco debug object path to `ffi::ToJSONGraph`/`ffi::FromJSONGraph` plus JSON parse/stringify. The disco Python wrappers now declare Python attribute storage explicitly for `DRef` and `Session` so `DPackedFunc`/`DModule` and method caches continue to work with the current tvm-ffi object model. The socket address helper also normalizes `localhost` consistently across constructors so the disco socket debug round-trip can bind an IPv4 socket when `localhost` resolves to IPv6 first. Validated locally in an isolated worktree build with `ninja -C build tvm_compiler tvm_runtime_extra`, targeted IR/target tests, `tests/python/disco/test_session.py::test_string_obj`, import smoke, and touched-file pre-commit. |
||
|
Tianqi Chen
|
dea2bf933e |
[REFACTOR][TIRX] Consolidate split host device stages (#19663)
The host/device split flow already runs device-region annotation, host/device function extraction, and device-kernel launch lowering as one consecutive pipeline. Keeping those stages exposed as separate public passes makes the API surface larger than the actual execution model and leaves the stage dependencies spread across multiple files. This change makes `tirx.transform.SplitHostDevice` the single public entry point for that flow, while preserving the existing stage order internally. Changes: - Merge the annotation, splitting, and kernel-launch lowering implementations into `src/tirx/transform/split_host_device.cc` as private sections. - Remove the old public C++ declarations, FFI registrations, and Python wrappers for `AnnotateDeviceRegions` and `LowerDeviceKernelLaunch`. - Replace pipeline call sites that previously invoked the three-stage sequence with one `SplitHostDevice()` call. - Update TIRx and S-TIR tests to exercise the consolidated pass and the reduced public API surface. |
||
|
Shushi Hong
|
a979b2f98c |
[RPC] Import tvm.testing lazily in rpc.testing (#19658)
### Motivation `tvm.testing` imports `pytest` at module load (`tvm/testing/utils.py`). `tvm.rpc.server` imports `tvm.rpc.testing` (to register the `rpc.test.*` helpers), and `tvm.rpc.testing` imported `tvm.testing` at the top level, so a plain `import tvm` / `import tvm.relax` pulls `pytest` in through: ``` tvm.relax -> tvm.runtime.vm -> tvm.rpc -> rpc.server -> rpc.testing -> tvm.testing -> pytest ``` As a result `pytest` is effectively a runtime dependency: a user who installs TVM without `pytest` hits `ModuleNotFoundError: No module named 'pytest'` on import. This is easy to miss because test environments install `pytest`. ### Change `tvm.rpc.testing` only uses `tvm.testing.object_use_count` in a single test helper, so import it lazily at the call site instead of at module top level. This keeps the `rpc.test.*` registration and the helper behavior intact while removing `tvm.testing` (and `pytest`) from the `import tvm` path, so `pytest` can remain a test-only dependency. No functional change; `rpc.testing` is still imported by `rpc.server` and still registers the same global functions. |
||
|
Bohan Hou
|
57c638fc7c |
[TIRx] Post-bringup op-dispatch / codegen / TVMScript follow-ups (#19657)
## Summary Follow-up work on top of the TIRx infrastructure bring-up (#19581). It extends the TIRx operator-dispatch, codegen, and TVMScript surfaces with the next batch of low-level programming features for Blackwell-class GPUs, while keeping `s_tir` script support intact. ## Main Changes - **op-dispatch**: warp `ldmatrix`/`stmatrix` copy dispatch; split CUDA copy into register / gmem-smem / `ldgsts` paths; `tcgen05.ld/st` `.16x{64,128,256}b` dispatch with a factory and M=128 layout; element-wise broadcast at the layout level with a copy vec-alignment fix. - **gemm**: CUDA synchronous `mma.sync` tensor-core dispatch; accept a Layout F C operand for M=64 MMAs. - **op**: add the `permute_layout` primitive (replaces `permute_dims`). - **tvmscript**: add the `Tx.jit` decorator, `Tx.constexpr` compile-time params, and `Tx.wg_reg_tile`. - **lower-tirx**: introduce the `Tx.device_entry()` marker (replacing `ScopeKind::kKernel`); canonical thread filters that drop the `Tx.filter` wrapper. - **codegen**: add a typed-pointer byte-offset intrinsic; remove the `entry_cluster_sync` codegen attribute. ## Validation - `pre-commit run` (changed files) — clean - `ninja -C build -j$(nproc)` — builds - `pytest tests/python/tirx/ -n 16` - `1997 passed, 39 skipped, 3 xpassed` - `python -m pytest tests/python/all-platform-minimal-test` - `37 passed, 105 skipped` - `TVM_TEST_TARGETS=llvm pytest tests/python/tirx-analysis tests/python/tirx-base tests/python/tirx-transform -n 16` - `630 passed, 25 skipped, 8 xfailed, 1 xpassed` ## Local CI Notes Several full CI-equivalent jobs are not locally reproducible because this machine is missing parts of the Apache TVM CI environment (e.g., specific `llvm-config` versions, Vulkan, ROCm, ARM/QEMU cross-toolchain, and web/wasm components). The Blackwell/Trainium kernel tests are maintained downstream and are intentionally not part of this PR. |
||
|
YinHanke
|
4a688ddcbc |
[Relax][Frontend][TFLite] Add EMBEDDING_LOOKUP_SPARSE converter (#19652)
## Summary Add Relax TFLite frontend support for `EMBEDDING_LOOKUP_SPARSE`. This PR adds a converter for `EMBEDDING_LOOKUP_SPARSE` in the Relax TFLite frontend. The implementation supports the `SUM`, `MEAN`, and `SQRTN` combiners and handles higher-rank sparse indices. The sparse aggregation is lowered through `scatter_nd` to match TFLite operator semantics for the supported cases. The PR also adds handcrafted TFLite frontend tests covering: - `SUM` - `MEAN` - `SQRTN` - a 3D indices case ## Testing Ran `tests/python/relax/test_frontend_tflite.py -k 'embedding_lookup_sparse'`. Part of #19519 --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> |
||
|
Javier De Jesus
|
9898909392 |
[Relax][PyTorch] Cast non-bool inputs to bool in logical_not converter (#19645)
### Motivation `torch.logical_not` accepts an input tensor of any dtype (treating any nonzero element as `True`) and always returns a `bool` tensor. The PyTorch frontend previously lowered it with `self._unary_op(relax.op.logical_not)`. `relax.op.logical_not` is a unary arithmetic op that passes its input dtype through, so a non-bool input (for example `float32`) produced a `float32` result instead of the `bool` result PyTorch returns. This is a dtype mismatch against the reference PyTorch semantics for both the FX and ExportedProgram frontends. ### Changes - Add a shared `_logical_not` converter in `BaseFXGraphImporter` that casts non-bool inputs to `bool` before applying `relax.op.logical_not`. Bool inputs are passed through unchanged (no redundant cast). - Point the `logical_not` (FX) and `logical_not.default` (ExportedProgram) registrations at the new converter. - Update the FX test and add a standalone ExportedProgram `test_logical_not` to assert the corrected IR (`astype` to bool, then `logical_not`, producing a `bool` output). ### Notes The cast to `bool` lowers to an elementwise nonzero test, so it matches PyTorch's "nonzero is True" semantics for float, integer, and NaN inputs. |
||
|
HoYi
|
23a0ea8d8b |
[Relax][Frontend][TFLite] Support STABLEHLO_RNG_BIT_GENERATOR (#19651)
## Summary
This PR adds Relax TFLite frontend support for the TFLite builtin
`STABLEHLO_RNG_BIT_GENERATOR` operator.
Unlike most StableHLO builtins, the TFLite runtime
(`tensorflow/lite/kernels/rng_bit_generator.cc`) implements this op as a
real,
deterministic counter-based PRNG, so the importer must reproduce it
bit-exactly
rather than map it to an existing op:
- one uint64 1-D `initial_state` input, two outputs — uint64
`output_state` and
the random-bit `output` (int32 / int64 / uint32 / uint64);
- `algorithm` in `{DEFAULT, PHILOX, THREEFRY}`, where `DEFAULT` resolves
to
`PHILOX`;
- Random123 Threefry2x32 (20 rounds) and Philox4x32 (10 rounds) with the
fixed
constants from `rng_util.cc`;
- state-length constraints: `THREEFRY` requires `u64[2]`,
`PHILOX`/`DEFAULT`
require `u64[2]` or `u64[3]`.
## Design
TVM/Relax has no matching RNG primitive, so the converter generates a
TIR kernel
that mirrors the runtime and emits it through `relax.call_tir` with two
outputs.
The kernel:
- reinterprets the uint64 state as uint32 words and advances a 64-bit
block
counter (`final counter = initial_state[1] + num_blocks`);
- runs the selected algorithm per block with all round state
materialized into
local buffers, which keeps the generated IR linear instead of an
exponentially
nested expression tree;
- packs the produced uint32 words back into the output dtype, and writes
the
updated state (key unchanged, counter advanced, Philox `u64[3]` tail
passed
through) — the only state behaviour the runtime relies on.
The kernel is an `s_tir` PrimFunc wrapped in a single opaque structured
block so
it remains a well-formed block-structured function for the Relax
pipeline
(e.g. `HasReshapePattern`). `get_tensor_type_str` and the input
`_decode_type`
map are extended with uint32/uint64 so the uint64 state imports
correctly.
Unsupported inputs raise a precise `OpNotImplemented` (non-uint64 /
non-1-D
state, mismatched output-state shape, unsupported output dtype, unknown
algorithm, per-algorithm state-length violations).
## Operator Support
| Operator | TFLite options | Relax lowering | Supported subset |
|---|---|---|---|
| `STABLEHLO_RNG_BIT_GENERATOR` |
`StablehloRngBitGeneratorOptions.Algorithm()` from `BuiltinOptions2` |
`call_tir` to a generated bit-exact TIR kernel | THREEFRY (`u64[2]`) and
PHILOX/DEFAULT (`u64[2]`/`u64[3]`); int32/int64/uint32/uint64 output |
## Tests
Tests build minimal RNG flatbuffers, compile, and execute them,
comparing the
output and updated state against the verbatim expected vectors from the
TFLite
runtime kernel test (`rng_bit_generator_test.cc`).
| Test | Coverage |
|---|---|
| `test_stablehlo_rng_bit_generator_threefry` | THREEFRY bit-exact, all
4 output dtypes |
| `test_stablehlo_rng_bit_generator_philox` | PHILOX bit-exact, all 4
output dtypes |
| `test_stablehlo_rng_bit_generator_default_matches_philox` | DEFAULT
resolves to PHILOX |
| `test_stablehlo_rng_bit_generator_deterministic` | run-to-run
bit-identical output |
| `test_stablehlo_rng_bit_generator_unsupported_output_dtype` | output
dtype guard |
| `test_stablehlo_rng_bit_generator_threefry_invalid_state_unsupported`
| THREEFRY `u64[2]` state guard |
| `test_stablehlo_rng_bit_generator_non_uint64_state_unsupported` |
uint64 state guard |
Local validation:
```bash
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 rng_bit_generator -q
python -m pytest \
tests/python/relax/test_frontend_tflite.py \
-k stablehlo -q
```
Result:
```text
ruff check: All checks passed
rng_bit_generator tests: 13 passed
stablehlo tests: 96 passed
```
## References
- Issue #19519 item I: remaining StableHLO operators in TFLite
- `tensorflow/lite/kernels/rng_bit_generator.cc`, `rng_util.cc`,
`rng_bit_generator_test.cc`
|
||
|
YinHanke
|
066bf777b8 |
[Relax][Frontend][TFLite] Add HASHTABLE_LOOKUP converter (#19654)
## Summary Add Relax TFLite frontend support for `HASHTABLE_LOOKUP`. This PR adds a converter for `HASHTABLE_LOOKUP` in the Relax TFLite frontend. The implementation supports non-string value tensors and lowers the lookup through `bucketize`, `take`, and `where` so that missing keys return zero-filled values together with a `uint8` hits mask matching TFLite semantics for the supported cases. The PR also adds handcrafted TFLite frontend tests covering: - 1D float value tensors - 2D float value tensors - the current unsupported string-value case ## Testing Ran `tests/python/relax/test_frontend_tflite.py -k 'hashtable_lookup'`. Part of #19519 |
||
|
Balint Cristian
|
349225ae23 |
[REFACTOR][PYTHON] Revisit lifted support modules from tvm.contrib (#19653)
In continuation of #19624 this catches some unlifted entries. Hope there is no more left, for consistency it now covers comments and perhaps non-active (hotpath) parts. |
||
|
HoYi
|
cf859b927a |
[Relax][Frontend][TFLite] Support STABLEHLO_CUSTOM_CALL (#19649)
## Summary This PR adds conservative Relax TFLite frontend support for the TFLite builtin `STABLEHLO_CUSTOM_CALL` operator. TFLite marks `STABLEHLO_CUSTOM_CALL` as having no runtime kernel. Importing general custom calls as executable Relax operators would therefore give them semantics that TFLite itself does not provide. This PR only supports the metadata-only `Sharding` custom call target, which TensorFlow's StableHLO pipeline treats as an annotation that can be erased. ## Design ### Sharding Annotation Lowering `STABLEHLO_CUSTOM_CALL` now parses `StablehloCustomCallOptions` from `BuiltinOptions2` and reads the `call_target_name`. For `call_target_name == "Sharding"`, the frontend lowers the op to identity: the output tensor is bound to the input expression. This mirrors TensorFlow's handling of Sharding custom calls as metadata annotations. The sharding spec in `backend_config` is intentionally dropped for single-device import. The supported subset is guarded: - exactly one input and one output - input and output shape/dtype metadata must match - `has_side_effect` must be false - `called_computations` must be empty All other custom-call targets raise `OpNotImplemented` with the target name in the diagnostic. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `STABLEHLO_CUSTOM_CALL` | `StablehloCustomCallOptions` from `BuiltinOptions2` | identity for `Sharding`; otherwise unsupported | metadata-only `Sharding` annotations with unchanged tensor metadata | ## Tests The tests manually build minimal StableHLO custom-call TFLite flatbuffers and compare the supported identity path with `tvm.ir.assert_structural_equal`. Unsupported patterns use `pytest.raises`. | Test | Coverage | |---|---| | `test_stablehlo_custom_call_sharding` | `Sharding` annotation lowers to identity | | `test_stablehlo_custom_call_unsupported_target` | unknown external target guard | | `test_stablehlo_custom_call_sharding_side_effect_unsupported` | side-effecting `Sharding` guard | | `test_stablehlo_custom_call_sharding_metadata_mismatch_unsupported` | input/output metadata guard | Local validation: ```bash python -m py_compile \ 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 stablehlo_custom_call -q python -m pytest \ tests/python/relax/test_frontend_tflite.py \ -k stablehlo -q ``` Result: ```text py_compile: passed ruff check: All checks passed stablehlo_custom_call tests: 4 passed stablehlo tests: 81 passed ``` ## References - Issue #19519 item I: remaining StableHLO operators in TFLite - TensorFlow Lite schema marks `STABLEHLO_CUSTOM_CALL` as no runtime support - TensorFlow StableHLO pipeline erases `Sharding` custom calls as metadata annotations |
||
|
HoYi
|
99488d992d |
[Relax][Frontend][TFLite] Support STABLEHLO_WHILE (#19646)
## Summary This PR adds Relax TFLite frontend support for the TFLite builtin `STABLEHLO_WHILE` operator. `STABLEHLO_WHILE` uses StableHLO `BuiltinOptions2` to reference its condition and body region subgraphs. Its loop semantics otherwise match the existing TFLite `WHILE` importer path: loop-carried tensors are passed to the cond/body subgraphs, the cond subgraph returns a scalar bool, and the body subgraph returns the updated loop state. ## Design ### Shared While Lowering The native TFLite `WHILE` converter is refactored through a shared `_convert_while_like` helper. Native `WHILE` and `STABLEHLO_WHILE` now share the same validation and lowering path after their options are parsed: - native `WHILE` reads `WhileOptions` from `BuiltinOptions` - `STABLEHLO_WHILE` reads `StablehloWhileOptions` from `BuiltinOptions2` Both paths lower the referenced cond/body subgraphs to private Relax functions and emit a recursive private Relax function for the loop. ### Boundary Validation `STABLEHLO_WHILE` reuses the same guard-first checks as native `WHILE`: - loop input count must match op output count - cond subgraph input metadata must match loop-carried tensors - cond subgraph must have exactly one output - cond output must be a scalar bool tensor - body subgraph input and output metadata must match loop-carried tensors - referenced cond/body subgraph indices must be valid non-main subgraphs The recursive loop-function cache key now includes the generated function prefix. This prevents native `WHILE` and `STABLEHLO_WHILE` from accidentally sharing a cached loop wrapper if they reference the same cond/body subgraph indices. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `STABLEHLO_WHILE` | `StablehloWhileOptions.CondSubgraphIndex()`, `BodySubgraphIndex()` from `BuiltinOptions2` | recursive private Relax function | tensor loop-carried state, scalar bool cond output, matching cond/body interfaces | ## Tests The tests manually build a minimal StableHLO while TFLite flatbuffer and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. Unsupported patterns use `pytest.raises`. | Test | Coverage | |---|---| | `test_stablehlo_while` | basic `STABLEHLO_WHILE` recursive private function lowering | | `test_stablehlo_while_non_bool_condition_unsupported` | cond output scalar bool guard | | `test_stablehlo_while_invalid_index_unsupported` | invalid cond/body subgraph index guard | | `test_stablehlo_while_output_count_mismatch_unsupported` | body output arity guard | | `test_stablehlo_while_input_metadata_mismatch_unsupported` | cond subgraph input metadata guard | | `test_stablehlo_while_output_metadata_mismatch_unsupported` | body subgraph output metadata guard | Local validation: ```bash python -m py_compile \ 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 stablehlo_while -q python -m pytest \ tests/python/relax/test_frontend_tflite.py \ -k stablehlo -q ``` Result: ```text py_compile: passed ruff check: All checks passed stablehlo_while tests: 6 passed stablehlo tests: 84 passed ``` ## References - Issue #19519 item I: remaining StableHLO operators in TFLite - PR #19587: StableHLO region-based ops and multi-subgraph model support - PR #19616: TFLite control-flow / multi-subgraph support |
||
|
YinHanke
|
e3933804ee |
[Relax][Frontend][TFLite] Support sequence LSTM and RNN operators (#19634)
## Summary Add three TFLite sequence recurrent operators to the Relax frontend, all with coupled input-forget gate (FULL kernel) and float32-only support. - UNIDIRECTIONAL_SEQUENCE_LSTM - BIDIRECTIONAL_SEQUENCE_RNN - BIDIRECTIONAL_SEQUENCE_LSTM From #19519. ## Changes - **UNIDIRECTIONAL_SEQUENCE_LSTM**: same layout as single-step LSTM, unrolls over time and stacks per-step hidden states. Supports time_major, cell_clip, proj_clip, and fused activation. - **BIDIRECTIONAL_SEQUENCE_RNN**: separate fw/bw RNN cells, backward scans in reverse. Supports merge_outputs (concat fw + bw) and split outputs via Tuple. - **BIDIRECTIONAL_SEQUENCE_LSTM**: 48-input operator with fw/bw LSTM cells sharing the same input tensor. States at indices 35-38. - All converters propagate final states to exp_tab for multi-step correctness. - Peephole, projection, layer norm, and aux input are not supported (raise OpNotImplemented). ## Testing - `test_unidirectional_sequence_lstm_none_activation` — output shape [batch, time, num_units] - `test_bidirectional_sequence_rnn_none_activation` — merge_outputs=True, shape [batch, time, 2*num_units] - `test_bidirectional_sequence_lstm_none_activation` — merge_outputs=True, shape [batch, time, 2*num_units] ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k "sequence_lstm or sequence_rnn" -v ``` |
||
|
HoYi
|
b971a75de4 |
[Relax][Frontend][TFLite] Add TFLite Resource Variable and Static Hashtable Import Support (#19639)
## Summary This PR adds incremental Relax TFLite frontend support for the resource variable initialization subset: - `VAR_HANDLE` - `ASSIGN_VARIABLE` - `READ_VARIABLE` It builds on the TFLite control-flow / multi-subgraph support from #19616, especially `CALL_ONCE`. TFLite commonly represents initialization through a `CALL_ONCE` init subgraph, then uses resource handles from the main subgraph to read initialized variables. This PR supports that constrained initialization pattern without introducing general mutable runtime state into Relax. The PR also adds explicit frontend guards for the TFLite builtin hashtable operators: - `HASHTABLE` - `HASHTABLE_IMPORT` - `HASHTABLE_FIND` - `HASHTABLE_SIZE` These operators are intentionally left unsupported for now. TFLite builtin hashtable kernels are not generic tensor maps: their runtime implementations cover the `int64 -> string` and `string -> int64` table variants, and correct import requires proper `TensorType.STRING` support. Rejecting the operators is safer than lowering a synthetic numeric table semantics that TFLite does not actually implement. ## Design ### Shared Initialization State The frontend now keeps resource initialization data in shared conversion state: - `conversion_state["resource_values"]` - `conversion_state["in_call_once_init"]` This state is shared by the main graph converter and the `CALL_ONCE` init subgraph converter. Each converter instance still keeps its own local `self.resource_handles` map, keyed by TFLite tensor name. Resource variables use `container + shared_name` from `VarHandleOptions` when present, falling back to the handle tensor name. This keeps tensor-name bindings scoped to each subgraph while allowing init subgraphs and the main graph to agree on the same logical resource. ### CALL_ONCE Init Subgraphs `CALL_ONCE` now accepts a non-empty init subgraph when all operators are in the supported initialization subset: - `VAR_HANDLE` - `ASSIGN_VARIABLE` The init subgraph still must have no inputs and no outputs. The converter first checks every operator against the allowlist, then converts the init subgraph with a fresh `ExprTable` and shared conversion state. The init subconverter deliberately shares the parent `BlockBuilder`. This is safe for the current subset because all supported init operators update importer state and return `None`; they do not emit Relax bindings. A comment documents that this should be revisited if future `CALL_ONCE` init operators emit Relax expressions. ### Resource Variables `VAR_HANDLE` is declarative. It registers the output resource tensor in the current converter's local `resource_handles` map and returns `None`. `ASSIGN_VARIABLE` is accepted only while converting a supported `CALL_ONCE` init subgraph. It resolves the resource handle through the init converter's local handle map and stores the assigned tensor expression in shared `conversion_state["resource_values"]`. `READ_VARIABLE` resolves the main graph resource handle and returns the initialized expression from shared state. If the resource has not been initialized by a supported `CALL_ONCE` path, the frontend raises `OpNotImplemented`. This supports the common static-initialization inference pattern while avoiding incorrect lowering for runtime mutation. ### Hashtable Operators `HASHTABLE` registers the table handle and validates the dtype pair against TFLite kernel constraints (`int64/string` or `string/int64`). `HASHTABLE_IMPORT` in a supported `CALL_ONCE` init subgraph captures static metadata (table size, key/value dtypes) but does not store actual string data, because Relax does not yet support `TensorType.STRING`. `HASHTABLE_SIZE` returns a scalar Relax constant for statically imported tables. `HASHTABLE_FIND` is rejected with `OpNotImplemented` because Relax cannot represent TFLite string tensors or the runtime lookup semantics. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `VAR_HANDLE` | `VarHandleOptions` | handle registration only | main graph and supported `CALL_ONCE` init subgraphs | | `ASSIGN_VARIABLE` | `AssignVariableOptions` | store initialized Relax expression in shared importer state | supported `CALL_ONCE` init subgraphs only | | `READ_VARIABLE` | `ReadVariableOptions` | return initialized Relax expression | resource must have supported static initialization | | `HASHTABLE` | `HashtableOptions` | handle registration + dtype validation | validates `int64/string` or `string/int64` pair, rejects other combinations | | `HASHTABLE_IMPORT` | `HashtableImportOptions` | store static metadata (size, key/value dtype) | `CALL_ONCE` init subgraphs only, constant key/value shape validation | | `HASHTABLE_FIND` | `HashtableFindOptions` | unsupported guard | requires future `TensorType.STRING` support in Relax | | `HASHTABLE_SIZE` | `HashtableSizeOptions` | scalar Relax constant | returns `[size]` int64 for statically imported tables | ## Safety Checks - `ASSIGN_VARIABLE` outside `CALL_ONCE` initialization raises `OpNotImplemented`. - `READ_VARIABLE` without supported initialization raises `OpNotImplemented`. - `CALL_ONCE` init subgraphs with inputs or outputs remain unsupported. - `CALL_ONCE` init subgraphs containing operators outside the resource-variable initialization allowlist remain unsupported. - TFLite builtin hashtable operators raise `OpNotImplemented` until the frontend can model their real int64/string table semantics. ## Not Included - Runtime `ASSIGN_VARIABLE` mutation in the main graph. - Runtime resource-state threading through Relax function parameters and returns. - Cross-subgraph resource handle aliasing beyond the static `container/shared_name` matching pattern. - Multiple runtime writes with ordering semantics. - TFLite builtin hashtable lowering. - `TensorType.STRING` import support. ## Tests The tests manually build minimal TFLite flatbuffers and compare imported Relax IR with `tvm.ir.assert_structural_equal`. Unsupported patterns use `pytest.raises`. | Test | Coverage | |---|---| | `test_resource_variable_call_once_init_read` | `CALL_ONCE` init subgraph with `VAR_HANDLE + ASSIGN_VARIABLE`, then main graph `READ_VARIABLE` | | `test_assign_variable_main_subgraph_unsupported` | runtime/main graph `ASSIGN_VARIABLE` guard | | `test_read_variable_uninitialized_unsupported` | `READ_VARIABLE` without supported initialization guard | | `test_hashtable_call_once_import_find_unsupported` | hashtable init/find path remains unsupported | | `test_hashtable_call_once_import_size_unsupported` | hashtable init/size path remains unsupported | | `test_hashtable_import_main_subgraph_unsupported` | main graph `HASHTABLE_IMPORT` remains unsupported | | `test_hashtable_size_uninitialized_unsupported` | uninitialized `HASHTABLE_SIZE` remains unsupported | Local validation: ```bash python -m py_compile \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py 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 "resource_variable or read_variable_uninitialized or hashtable" -q python -m pytest \ tests/python/relax/test_frontend_tflite.py -q ``` Result: ```text py_compile: passed ruff format --check: files already formatted ruff check: All checks passed targeted resource/hashtable tests: 6 passed full test_frontend_tflite.py: 472 passed ``` |
||
|
YinHanke
|
576e60e974 |
[Relax][Frontend][TFLite] Add LSTM and SVDF converter (#19633)
## Summary Add LSTM (coupled input-forget) and SVDF single-step converters to the TFLite frontend. Both are float32-only; quantized variants are not supported yet. From #19519. ## Changes - **LSTM**: FULL kernel type, coupled input-forget gate only. Peephole, projection, and layer norm are not supported - **SVDF**: Standard SVDF with feature projection + time filtering + bias + fused activation - Both converters validate unsupported modes (quantized, non-coupled LSTM) with clear error messages ## Testing - `test_lstm_none_activation` — verifies LSTM converter produces correct IR shapes (batch, input_size) → (batch, num_units) with 3 params (input, h_state, c_state) - `test_svdf_none_activation` — verifies SVDF converter produces correct IR shapes (batch, input_size) → (batch, num_filters) with 2 params (input, state) ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k "lstm or svdf" -v ``` ## References - TFLite LSTM spec: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/lstm.cc - TFLite SVDF spec: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/svdf.cc |
||
|
Tianqi Chen
|
878903f574 |
[REFACTOR][IR] Delete class Bool and class Integer boxed-type wrappers (#19636)
## Background `class Integer : public IntImm` and `class Bool : public IntImm` were thin wrappers sharing `IntImmNode` with no separate node class and no FFI registration. They existed to provide implicit int→Integer constructors and a `.IntValue()` / `operator bool()` accessor, but the same functionality is available directly through `IntImm`. ## What this PR does Migrates all call sites away from `Integer` / `Bool` and then deletes the class definitions. The changes are split into four commits, each independently buildable: **Commit 1 – [REFACTOR][TIR]** Replace IR-position `Integer(N)` / `Bool(b)` constructors with `IntImm(DataType::Int(32), N)` / `IntImm(DataType::Bool(), b)` across ~62 source files (arith, relax analysis, s_tir schedule state, transform passes, codegen). **Commit 2 – [REFACTOR][SCHEDULE]** Migrate `Schedule` and `MetaSchedule` trace-boxing code: `Integer(N)` attrs in `TracedSchedule` → `IntImm(DataType::Int(32), N)`; `ffi::Array<Integer>` schedule-rule parameters → `int64_t`; `Bool(b)` attrs → `IntImm(DataType::Bool(), b)`. **Commit 3 – [REFACTOR][TOPI]** Migrate topi container signatures (`ffi::Array<Integer>` → `ffi::Array<int64_t>`) and update all internal usages (`.IntValue()` → plain int64_t, `.defined()` → removed, `->value` → direct indexing). Also handles stray `Integer` / `Bool` variables in clml codegen, make_packed_api, infer_layout_utils, and relax distributed code. **Commit 4 – [REFACTOR][IR]** Delete `class Bool`, `class Integer`, `TypeTraits<Bool>`, and `TypeTraits<Integer>` from `include/tvm/ir/expr.h`. ## Canonical replacements | Old | New | |-----|-----| | `Integer(N)` | `IntImm(DataType::Int(32), N)` | | `Bool(b)` | `IntImm(DataType::Bool(), b)` | | `x.IntValue()` | `x->value` | | `x` as bool | `x->value != 0` | | `ffi::Array<Integer>` | `ffi::Array<int64_t>` | ## Testing - All 118 C++ unit tests pass (`./cpptest`) - `tests/python/s_tir/` — 1251 passed (14 pre-existing failures unrelated to this change, all in TIR transform tests with annotation-mismatch errors) - `tests/python/relax/` — passes (excluding pre-existing torch/torchvision import failures in frontend tests) |
||
|
YinHanke
|
e89570fa83 |
[Relax][Frontend][TFLite] Add RNN converter (#19632)
## Summary Add Relax TFLite frontend support for `RNN` (BuiltinOperator 23), claimed in [#19519](https://github.com/apache/tvm/issues/19519) Group A. Single-step RNN cell: ``` h = fused_activation(x @ W.T + h @ Wr.T + b) ``` ## Changes - **Handler**: `convert_rnn` registered in `convert_map` (alphabetical, after `RANGE`) - **Inputs** (5): `input [batch, 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, num_units]` - **Activations**: all fused activations via `convert_fused_activation_function` - **Quantized**: raises `OpNotImplemented` ## Testing Two tests added to `tests/python/relax/test_frontend_tflite.py`: - `test_rnn_none_activation` — `tvm.ir.assert_structural_equal` with identity weights, NONE activation - `test_rnn_relu_activation` — shape check, random weights, RELU activation ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k rnn -v ``` ## References - Issue [#19519](https://github.com/apache/tvm/issues/19519) Group A: Sequence / recurrent model operators |
||
|
Sun
|
0d70112300 |
[Relax][Frontend][TFLite] Add REDUCE_WINDOW support (#19637)
## Summary Add Relax TFLite frontend support for the builtin `REDUCE_WINDOW` operator. This covers the ordinary TFLite op only, not `STABLEHLO_REDUCE_WINDOW`. The converter parses `ReduceWindowOptions` from `BuiltinOptions2`, validates the static window attributes, and lowers supported reduce functions through `topi.sliding_window` plus Relax reductions. Supported modes: - `ADD` - `MUL` - `MINIMUM` - `MAXIMUM` - `ALL` - `ANY` Empty output shapes are handled directly with `relax.op.zeros`. Quantized `REDUCE_WINDOW`, dynamic window attributes, and unsupported reduce functions remain rejected with explicit errors. ## Testing - `python -m py_compile 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 reduce_window -q -p no:tvm.testing.plugin` - `python -m pytest tests/python/relax/test_frontend_tflite.py -k "reduce_window or reduction_ops" -q -p no:tvm.testing.plugin` - `conda run -n test python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.py` ## Related Related to #19519. |
||
|
Tianqi Chen
|
d26ea6ff51 |
[REFACTOR][SCRIPT] tvmscript streamline: lift printer.h, restore one-way dep, migrate dialect config to extra_config (#19631)
## Background The `tvm::ir` layer previously had a reverse dependency on `tvm::script`, injected via the `TVM_OBJECT_ENABLE_SCRIPT_PRINTER()` macro that added a `Script()` member method to IR node types (IRModule, PrimExpr, Buffer, PrimFunc, Stmt). This violated the intended one-way dependency: `script` should depend on `ir`, never the other way around. Additionally, `PrinterConfigNode` accumulated dialect-specific fields (`tir_prefix`, `tir_import_module`, `tirx_prefix`, `relax_prefix`) that created leakage between the generic printer infrastructure and dialect internals. ## Changes This PR restores the clean dependency direction and encapsulates dialect config properly, in 5 commits: 1. **Lift TVMScript entry point into `script/printer/printer.h`**: New header `include/tvm/script/printer/printer.h` introduces: - `tvm::Script()` free function replacing `TVMScriptPrinter::Script()` static method - `TVMScriptPrinter` class with vtable (`NodeFunctor<std::string(...)>`) - `TVM_REGISTER_SCRIPT_AS_REPR` macro for registering per-type repr callbacks 2. **Drop `TVM_OBJECT_ENABLE_SCRIPT_PRINTER` macro**: Remove the macro from all IR headers (`ir/expr.h`, `ir/module.h`, `tirx/buffer.h`, `tirx/function.h`, `tirx/stmt.h`), eliminating the reverse `ir` → `script` dependency. All call sites of `.Script()` member methods updated to use `tvm::Script()`. 3. **Move dialect-specific `PrinterConfig` fields to `extra_config`**: Remove `tir_prefix`, `tir_import_module`, `tirx_prefix`, `relax_prefix` from `PrinterConfigNode`. Dialect internals now read their config via `GetExtraConfig<T>(key, fallback)` with dotted keys (e.g., `"tirx.prefix"`). `buffer_dtype` is kept as a top-level field alongside `int_dtype`/`float_dtype` since it is a shared scalar-literal default, not a dialect-specific knob. 4. **Python: drop dialect kwargs, expose `extra_config`**: Update `PrinterConfig`, `Scriptable.script()`, `Scriptable.show()`, `Scriptable._relax_script()`, and `BasePyModule.script()` to use `extra_config: dict | None = None` instead of individual dialect kwargs. The tirx auto-switch logic is preserved. 5. **Fix transitive include breakage**: Explicitly add direct includes for `config.h` and `node_functor.h` where headers previously relied on transitive paths through `expr.h`/`module.h`. ## Testing - C++ unit tests: 118/118 pass - TVMScript printer tests: 771 passed, 1 skipped, 1 xfailed - TIR namespace tests (`tests/python/tirx/test_printer_tir_namespaces.py`): 13/13 pass - Relax AST printer tests: 24/24 pass - Minimal platform tests: 37/37 pass - Pre-commit (ASF headers, ruff, clang-format): all clean |
||
|
Tianqi Chen
|
61ae85b9d1 |
[REFACTOR][PYTHON] Consolidate derived_object into tvm.ir.utils (#19630)
## Summary `derived_object` was duplicated byte-for-byte across `python/tvm/runtime/support.py` and `python/tvm/s_tir/meta_schedule/utils.py`. The function is not a runtime feature and is used outside meta_schedule (tvm.relax, tvm.tirx), so neither location was the right home. Move the single canonical definition into a new `python/tvm/ir/utils.py`. `tvm.ir` loads before both `tvm.tirx` and `tvm.s_tir`, so eager top-level imports work from every consumer without load-order workarounds. Rewrite all 25 caller imports. Keep the better-typed `cls: type[T] -> type[T]` signature from the runtime-side copy. After this change `runtime/support.py` is empty and is removed; `meta_schedule/__init__.py` drops its now-dead re-export. No alias shims are left behind — callers update imports directly. |
||
|
Tianqi Chen
|
23db01f1fa |
[REFACTOR][RUNTIME] Structural reorganization: locality moves for thread_map, texture, minrpc, disco, contrib (#19628)
## Background The TVM runtime has been growing organically. Several headers and directories live at the top level of `src/runtime/` despite only being consumed by a single backend subsystem. This PR applies the **locality principle**: code that has exactly one consumer moves to live next to that consumer. ## Changes ### Move 1: `thread_map.h` → `src/runtime/vulkan/` `ThreadMap` is only used by Vulkan device API headers. Moving it under `src/runtime/vulkan/` reflects this exclusive ownership. ### Move 2: `texture.h` → `src/runtime/opencl/` Texture storage utilities are OpenCL/Adreno-specific. Moving the header under `src/runtime/opencl/` makes ownership clear. ### Move 3: `minrpc/` → `src/runtime/rpc/minrpc/` The minrpc mini-RPC implementation belongs logically under the existing `src/runtime/rpc/` subtree. All consumers already live under rpc/ or reference it as a child of rpc/. ### Move 4: Introduce `src/runtime/extra/` boundary `disco/` and `contrib/` are the sole source directories for `libtvm_runtime_extra`. Grouping them under `src/runtime/extra/` makes the `libtvm_runtime_extra` build boundary visible in the filesystem, matching the modular runtime split introduced in #19444. - `src/runtime/disco/` → `src/runtime/extra/disco/` - `src/runtime/contrib/` → `src/runtime/extra/contrib/` - Public `include/tvm/runtime/disco/` is unchanged. ### Drive-by fixes - `apps/android_rpc/…/tvm_runtime.h`: Drop stale `minrpc_logger.cc` include (file no longer exists) and fix stale `tvm-ffi/src/ffi/extra/testing.cc` path to `tvm-ffi/src/ffi/testing/testing.cc`. ## Test Plan - [x] Full build (`ninja -j$(nproc)`) — succeeds - [x] `./cpptest` — 118 tests passed - [x] Python smoke: `tvm.__version__` + `tvm.cuda(0).exist` — pass - [x] `tests/python/all-platform-minimal-test` — 37 passed, 105 skipped - [x] `tests/python/runtime/test_runtime_rpc.py` — 2 passed, 21 skipped - [x] `tests/python/runtime/test_rpc_base.py` — 2 passed - [x] `pre-commit run --all-files` — all hooks pass |
||
|
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 |