This pr modernizes test gating. It replaces the heavy
`tvm.testing.Feature` machinery with a thin `tvm.testing.env` module of
`has_*()` capability probes, used via standard pytest.mark + skipif. And
markers move to `pyproject.toml`
TVM's shipped code only uses cuda.bindings — cuda.bindings.nvrtc for the
NVRTC JIT path and cuda.bindings.driver for the NVSHMEM link path, both
in python/tvm/support/nvcc.py; it never uses cuda.core. cuda-python is
now a metapackage that pulls in cuda-bindings + cuda-core (and
cuda-pathfinder), so depending on it drags in cuda-core that TVM does
not need.
Depend directly on cuda-bindings, which provides exactly the nvrtc and
driver submodules TVM imports, and update the user-facing 'pip install
cuda-python' hints to match. A plain cuda-bindings install pulls no
nvidia-* toolkit wheels (those live behind the [all] extra); libnvrtc is
loaded from the system / TVM's CUDA install as before.
IR module cleanup benefits from using a single unique-name primitive
directly at module call sites. This PR renames NameSupply to
UniqueNameSupply and removes redundant wrappers around global variable
naming.
Main changes:
- Rename the public name supply API and header to UniqueNameSupply
- Replace GlobalVarSupply with direct iterator-seeded UniqueNameSupply
usage
- Remove obsolete access-path repr registration now covered by tvm-ffi
Hi Committers,
This PR addresses the `ReduceMax`/ `ReduceMin` part of issue
https://github.com/apache/tvm/issues/19572. Any suggestions would be
appreciated if you are available.
### Root cause:
The ONNX frontend ReduceMax / ReduceMin converters return relax.op.max /
relax.op.min. After legalization these map to topi.max / topi.min, which
fold with a commutative reducer whose combiner is Max(x, y) / Min(x, y).
In codegen, Max(a, b) lowers to select(a > b, a, b) using an **ordered**
float comparison (fcmp ogt), which is false for NaN. As a left-fold (acc
= Max(acc, elem)), NaN propagation becomes **position-dependent** - a
later non-NaN element silently overwrites an earlier NaN.
### Solution:
Adopt the well-defined, **order-independent numpy/IEEE convention**
(matching numpy.max/min and torch.amax/amin): the reduction yields NaN
whenever **any** reduced element is NaN. Minimal, ONNX-frontend-only
change:
- Add a shared helper _reduce_min_max_preserve_nan(reduce_op, data,
axes, keepdims).
- For floating-pint inputs, detect NaN along the reduced axes via
`sum(astype(isnan(data), dtype), axes, keepdims) > 0` and force those
outputs to `NaN` with `where(has_nan, nan, reduce(data))`. The mask
reduces over the **same axes/keepdims**, so it aligns in shape with the
reduced result.
- Keep non-floating(integer) inputs unchanged.
- Route all reduce paths(`_impl_v11`and both reduce branches of
`_impl_v18`) through the helper; the `noop_with_empty_axes` passthrough
is left untouched since it performs no reduction.
### Note on scope (re: #19589 ):
The underlying NaN behavior of Max/Min is the same family of ops
discussed in #19589. Per review comments there, enforcing NaN semantics
at the IR / LLVM-IR level is undesirable(backward-compat with older
LLVM, and portability to CUDA/OpenCL/Vulkan), and a dedicated portable
nanmin/nanmax TIRx intrinsic(like `nearbyint`) would be the preferred
long-term mechanism. This PR deliberately:
- does not touch the IR-level Max/Min lowering, and
- does not rely on the bool reduction of the NaN mask - it uses
`sum(isnan) > 0`, fully sidestepping Max/Min NaN behavior.
---------
Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
## Summary
Backend loading is easier to maintain when native backend library
discovery, in-tree backend Python hook loading, and out-of-tree entry
point autoload are owned by the backend namespace. This PR consolidates
those paths under `tvm.backend._autoload_backends` while preserving
compatibility routes from the previous top-level helper and
`tvm.base.load_backend_libs`.
- Move backend runtime DSO loading into `tvm.backend._autoload_backends`
- Route `backend.load_all()` through the backend autoload helper
- Keep the previous top-level `_autoload_backends` module as a thin
compatibility import
### Root Cause
The ONNX `CumSum` converter in the Relax frontend rejected any node with
`exclusive=1` via a bare
`assert not attr.get("exclusive", False), "Exclusive option not yet
supported."`, so importing a
model that used exclusive cumulative sums failed with `AssertionError:
Exclusive option not yet
supported.`. The underlying op already supports the exclusive form:
`relax.op.cumsum` forwards an
`exclusive` flag to the FFI and `topi/scan.py` implements the exclusive
branch, so the converter
only needed to pass the attribute through.
### Solution
Drop the assert and read the attribute as `exclusive =
attr.get("exclusive", 0) != 0` (matching the
existing `attr.get("reverse", 0) != 0` idiom in the same converter),
then pass it to
`relax.op.cumsum(data, axis, exclusive=exclusive)`. The existing reverse
handling
(`flip -> cumsum -> flip`) composes correctly with exclusive, so
`reverse=1, exclusive=1` lowers to
an exclusive scan over the reversed axis.
### Test Plan
Extended `test_cumsum` in `tests/python/relax/test_frontend_onnx.py` to
parametrize `exclusive` over
`[True, False]`, so `check_correctness` now exercises all four
`(reverse, exclusive)` combinations
against ONNX Runtime:
```
python -m pytest tests/python/relax/test_frontend_onnx.py::test_cumsum -v
```
### Issue
Fixes#19692
### Root cause
In the ONNX `LayerNormalization` spec the bias `B` is optional; when
omitted it should behave as
zeros shaped and typed like the scale `W`. In
`LayerNormalization._impl_v17`, the synthesized zero
bias instead took its shape from `data.struct_info.shape[1]` (an
unrelated data dim) and hardcoded
`dtype="float32"`. For input `[2, 3, 4, 8]` with scale `[8]` and
`axis=-1` this builds a bias of
shape `(3,)` while gamma is `(8,)`, so `relax.op.nn.layer_norm` raises a
size-mismatch
`InternalError`. The float32 hardcode also breaks fp16/bf16 no-bias
models, since gamma, beta, and
data must share a dtype. PyTorch's `nn.LayerNorm(..., bias=False)`
exports exactly this no-bias form.
### Fix
Derive both the shape and dtype of the synthesized zero bias from the
scale, matching the ONNX
semantics for an omitted `B` and the existing torch frontend
(`relax.const(np.zeros(shape), x.struct_info.dtype)`):
```python
if bias is None:
bias = relax.const(_np.zeros(gamma_shape, dtype=scale.struct_info.dtype))
```
`gamma_shape` and the `_np`/`get_const_tuple` imports are already
present. Deriving the dtype from
the scale (rather than the issue's float32-only suggestion) is what also
fixes the fp16/bf16 case.
### Test plan
Added non-square no-bias regression cases to
`test_frontend_onnx.py::test_layer_norm` (the previous
no-bias case was square, which masked the bug): float32
`[2,3,4,8]`/scale `[8]` and float16 with
full `check_correctness`, plus a bf16 importer-only case (ORT's CPU
provider has no bf16
LayerNormalization kernel).
Fixes#19691
## Summary
The in-tree target custom datatype path adds maintenance surface while
current development focuses on core datatypes. This PR phases out the
built-in registry/lowering implementation and keeps the core dtype
behavior intact.
- Remove the target/datatype implementation, BYODT posit build option,
and related Python helpers
- Remove the custom datatype lowering pass from TIRX and S-TIR
finalization pipelines
- Simplify remaining TIRX dtype handling back to built-in/core datatypes
This PR batches several post-bringup TIRx follow-ups, rebased onto
current `main`.
### Changes
- **op-dispatch**: per-call exec scope via `Tx.<scope>.op`; remove
`ExecScopeStmt`
- **namespaces**: split TIRx op namespaces; remove tile-primitive kind
attrs
- **codegen**: support explicit CUDA launch bounds
- **gemm-async**: support contiguous-axis (K-major) operand slicing
- **backend reorg**: move in-tree GPU backends out of core into
`src/backend/<target>/` and `python/tvm/backend/<target>/`
(codegen/runtime/op), with the corresponding `CMakeLists.txt` /
`cmake/modules` and include-path updates
### Testing
- Builds with `USE_CUDA=ON` / `USE_LLVM=ON`
- The TIRx Python test suite (`tests/python/tirx/`) passes locally
### Motivation
`torch.logical_or` and `torch.logical_xor` accept input tensors of any
dtype
(treating any nonzero element as `True`) and always return a `bool`
tensor.
Neither op was handled by the PyTorch frontend. The ExportedProgram
frontend did
not register `logical_or.default` / `logical_xor.default`, and the FX
frontend
did not register `logical_or` / `logical_xor`, so importing a model that
uses
either op failed early with `Unsupported function types`.
This follows up on #19679 (`logical_and`) and addresses the explicit
question
raised in #19743: whether `logical_or` and `logical_xor` need the same
handling.
### Changes
- Add shared `_logical_or` and `_logical_xor` converters in
`BaseFXGraphImporter`
that cast non-bool operands to `bool` before applying
`relax.op.logical_or` /
`relax.op.logical_xor`. Bool operands are passed through unchanged (no
redundant cast).
- Register `logical_or.default` / `logical_xor.default`
(ExportedProgram) and
`logical_or` / `logical_xor` (FX), matching the existing `logical_and`
converter.
- Add standalone `test_logical_or` and `test_logical_xor` to both the FX
and
ExportedProgram test suites, asserting the corrected IR (`astype` to
bool on
each operand, then the logical op, 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.
This fixes the legacy predicated `ptx.cp_async` codegen path used by
`InjectPTXAsyncCopy` for `if_then_else(..., 0)` stores.
The old inline CUDA emission zero-filled the shared-memory destination
when the predicate was false. The TIRx helper-based legacy codegen only
skipped the `cp.async`, leaving the destination slot unchanged. This
restores the previous behavior by emitting an `@!p st.shared.*` zero
store in the generated legacy predicated helper.
The CUDA source snapshot in
`test_s_tir_transform_inject_ptx_async_copy.py` is updated to reflect
the restored false-predicate zero-fill instruction and the current
generated helper-based CUDA source.
The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`,
so the Disco Python wrappers can no longer store instance attributes and
every session construction fails with AttributeError. Declare the
attributes each
wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu`
so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the
disco runtime being present so they skip cleanly on builds (e.g. the pip
wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it.
### Session attribute storage
- `DPackedFunc` / `DModule`: `__slots__ = ("session",)`.
- `Session`: `__slots__ = ("_cache", "_import_python_module")`
This PR fixes two independent test-isolation issues that only surface
when certain test files run together in one pytest session.
1. Fix `_DialectRedirectFinder` duplicate module execution
`_DialectRedirectFinder.find_spec` used to pre-register the redirect
target module under the legacy alias name before returning the alias
spec.
This interacts badly with CPython import logic: when the requested
module name is already in `sys.modules`, CPython may ignore the returned
alias spec and reuse the target module's original spec instead. As a
result, the target source can be executed again under the canonical
module name, creating a duplicate module object.
This caused patches on aliased modules to silently miss the module
object used by existing code. For example,
`unittest.mock.patch("tvm.tirx.script.builder.buffer_store")` patched
the duplicate module, while the tirx parser still held references to the
original one, so `test_scalar_assign_error_not_swallowed` failed with
`DID NOT RAISE`.
This pr removes the pre-registration and let the import machinery
register the alias normally. Since the alias spec is now used,
`_AliasLoader.exec_module` also restores the canonical `__spec__` and
`__loader__` to avoid stale alias metadata on the loaded module.
2. Remove unused `tirx.intrin_test` op registration
`test_s_tir_transform_lower_match_buffer.py` registered a dummy op:
```python
tvm.ir.register_op_attr("tirx.intrin_test", "")
```
This was a leftover from the old TVMScript parser and is no longer
needed. The modern tirx parser eagerly evaluates `intrin_test(...)`
calls into `T.evaluate(0)`, so this op never appears in parsed IR.
The only remaining effect was adding a category-less `tirx.intrin_test`
entry to the global op registry, which could break
`test_registered_tirx_ops_have_exactly_one_category` depending on test
import order.
This pr removes the unused registration.
This PR adds an autoload mechanism for out-of-tree backends, simplifies
TVM's Python library loading, and removes `TVMError` in favor of native
Python errors.
## Autoload out-of-tree backends
Out-of-tree packages can register an autoload callable under the
`tvm.backends` entry-point group (mirroring torch's device-backend
autoload). At `import tvm` startup each entry point is discovered and
its callable invoked once, after the core runtime and the `tvm`
namespace are fully initialized, so an extension can register
ops/targets/funcs or load extra libraries.
```toml
[project.entry-points."tvm.backends"]
tvm_foo = "tvm_foo:_autoload"
```
A failing extension is caught and surfaced via `warnings.warn` so it
cannot break `import tvm`. Autoload can be disabled with
`TVM_DEVICE_BACKEND_AUTOLOAD=0`.
## Simplify library loading
The library-loading path in `base.py` is consolidated around a single
`_LOADED_LIBS` dict (basename to ctypes handle) so downstream and
autoloaded extensions can skip already-loaded libraries; the per-backend
runtime DSO list is folded into `load_backend_libs`. Accumulated cruft
is removed: the Python-3.9 check, the readline shim, the `_FFI_MODE`
ctypes check, the `base.__version__` re-export, and `py_str` (call sites
inline `.decode("utf-8")`).
## Remove TVMError in favor of native Python errors
`TVMError` added a layer atop `RuntimeError` that downstream code had to
import and learn. It is removed; the registered FFI error kinds
(`InternalError`, `RPCError`, `OpError`, `DiagnosticError`,
`ScheduleError`) now subclass `RuntimeError` directly while staying
registered, so the FFI keeps throwing the right kinds. All `TVMError`
imports, `except`/`raise`/`isinstance` uses, and
`pytest.raises(tvm.TVMError)` sites move to the `RuntimeError` builtin.
This pr fixes the failures in
tests/python/nightly/test_nnapi/test_from_exported_to_cuda.py (PyTorch
export -> Relax -> CUDA), which despite the directory name are plain
CUDA exported-program tests.
1. test_index_tensor: aten.index.Tensor with multiple index tensors and
no None entries was lowered through a sequential-take fast path, which
computes an outer product over the index tensors. PyTorch/NumPy advanced
indexing broadcasts the index tensors together and applies them jointly,
so e.g. x[[0, 1], [0, 1]] on a (5, 5, 5, 5) tensor produced shape (2, 2,
5, 5) instead of (2, 5, 5). Route the no-None case directly to
relax.op.index_tensor (topi.adv_index), which implements the correct
broadcast-and-zip semantics. The sequential-take path is kept for the
sliced (None-containing) decomposed-interpolate pattern, whose
orthogonal index shapes make the outer product equivalent.
2. test_copy_ (and any in-place op on a buffer/user input):
functionalization in torch.export prepends mutation outputs
(BUFFER_MUTATION / USER_INPUT_MUTATION) to the graph outputs, and the
importer returned them as part of the Relax function's output tuple.
Callers indexing outputs positionally then read the mutated-buffer value
instead of the model output. Filter the outputs through the graph
signature's output_specs and keep only user-facing outputs (USER_OUTPUT
/ LOSS_OUTPUT). Frontend test expectations that asserted the extra
mutation outputs are updated.
3. test_cross_entropy_module:
- aten.sum over a bool tensor was emitted as relax.op.sum on the bool
input, which keeps dtype bool, so the non-ignored-target count in the
decomposed cross-entropy collapsed to 1.0 and the mean reduction
degenerated to a sum. Match PyTorch type promotion by casting bool and
sub-64-bit integer inputs to int64 (or the explicit dtype argument)
before summing.
- The dlight GPU Fallback rule skipped blocks with zero loops entirely,
so rank-0 kernels (the scalar divide) were marked tirx.is_scheduled
without any thread binding and failed VerifyMemory. Bind such blocks
through the existing add-unit-loop path, and only skip zero-loop blocks
that already launch threads internally (e.g. opaque sort kernels),
detected via thread_extent attrs / thread-bound loops in the block body.
Previously `_generate_triton_kernel` overwrote the user-provided kwargs
with the constexpr dict before calling triton.compiler.compile, so
options such as num_warps passed to T.call_kernel were silently dropped.
Pass constexprs to ASTSource and forward the user kwargs as compile
options.
The pre-3.3 compatibility branches are removed in favor of an explicit
minimum-version check: they were never exercised in CI (which does not
install Triton), and Triton >= 3.3 has shipped with PyTorch since 2.7.
The integration test now matches the actual lowering, where constexpr
parameters (BLOCK_SIZE) appear as runtime kernel arguments in
call_packed, and passes num_warps=8 expecting a thread extent of 256 to
cover the option forwarding.
This pr makes the Hexagon pytest plugin avoid importing Hexagon
build/session machinery at module import time.
This keeps pytest collection lightweight and avoids collection-time
failures from Hexagon-specific runtime/session imports when running
source-tree tests against a TVM wheel or an environment that does not
need to execute Hexagon tests.
Replace TVM's `Diagnostic` / `DiagnosticContext` machinery with the
tvm-ffi
`visit_error_context` mechanism. Validators throw an `ffi::Error` seeded
with the
offending node; leaf pass executors (`ModulePass` / relax `Function` /
`DataflowBlock`) catch and rethrow `EnrichPassErrorWithContext`, which
appends the
failing pass name and a TVMScript-rendered, underlined source location.
`relax.analysis.well_formed` now throws on the first violation; a new
`check_well_formed` returns a bool, and all C++/Python/test callers are
routed
accordingly. `include/tvm/ir/diagnostic.h` and `src/ir/diagnostic.cc`
are deleted.
The enrichment renders with `num_context_lines=10` so a small function
shows in
full with no skipped-lines marker, while a large module stays bounded.
The TVMScript parser diagnostics
(`python/tvm/script/parser/core/diagnostics.py`)
stay self-contained pure-Python with no `DiagnosticContext` dependency,
and
restore multi-line source rendering: a diagnostic whose offending AST
node spans
multiple source lines now renders every spanned line with its gutter
line number
and an underline covering the span. `tvm.error.DiagnosticError` (used by
the
TVMScript parser) is retained.
A rendered end-to-end enriched-error example is posted as a comment
below.
This PR slims `tvm.libinfo` into a thin *info* layer that delegates path
discovery to the `tvm_ffi.libinfo` primitives and never loads libraries.
Loading responsibilities move to `tvm.base`, and the various ad-hoc
path-finding helpers are phased out in favor of the tvm-ffi resolvers.
## Changes
- **libinfo**: add `find_libtvm_runtime()` (resolves `libtvm_runtime`
via
`_find_library_by_basename` + `_resolve_and_validate`) and
`find_tvm_include_path()` (TVM's own `include/`). `find_include_path()`
now
returns `[find_tvm_include_path(), *tvm_ffi.libinfo.include_paths()]`,
folding
in the FFI + dlpack + python-helper include dirs. Remove
`find_lib_path`,
`get_dll_directories`, `use_runtime_lib`, `split_env_var`, and
`load_backend_libs`.
- **base**: receive `load_backend_libs` and the backend DSO list; the
runtime-only switch becomes a strict `TVM_USE_RUNTIME_LIB == "1"` check.
- **rpc**: `with_minrpc` uses `find_libtvm_runtime()` (the `runtime`
kwarg is
retained as an inert back-compat parameter); the rpc server
`load_library`
resolves the literal library name against the current working directory.
- **wasm**: move the `web/dist` asset search into `emcc.find_wasm_lib`,
used by
`emcc.create_tvmjs_wasm` and the tvmjs asset lookup.
- **hexagon**: fix a latent bug where `_get_hexagon_rpc_lib_dir` called
a
non-existent `tvm_ffi.libinfo.find_lib_path`; it now relies solely on
the
`HEXAGON_RPC_LIB_DIR` environment variable.
Validating the apache-tvm wheel in a minimal environment (no torch,
scipy, cloudpickle, or tornado installed) produced 33 pytest collection
errors from module-level imports of optional packages. Add
pytest.importorskip guards so these modules are reported as skipped
instead of erroring during collection.
Indirect import chains guarded:
- tvm.topi.testing imports scipy
- tvm.s_tir.meta_schedule.testing.local_rpc (tvm.rpc.tracker) requires
tornado
- tvm.s_tir.dlight.benchmark imports cloudpickle
Also remove a stray pre-license-header `import pytest` in
test_runtime_builtin_paged_attention_kv_cache_flashinfer.py.
This fixes CoreML Relax partitioning after `FoldDataflowBlockOutput` was
folded into `CanonicalizeBindings`.
`partition_for_coreml` still called
`relax.transform.FoldDataflowBlockOutput()`, but that pass is no longer
exported. This caused CoreML partition tests to fail with:
```text
AttributeError: module 'tvm.relax.transform' has no attribute 'FoldDataflowBlockOutput'
```
The old pass behavior is now covered by `CanonicalizeBindings`, so this
updates CoreML partitioning to call `CanonicalizeBindings()` instead.
The existing CoreML test file now also includes a partition-only
regression test. The end-to-end CoreML tests remain guarded by the
CoreML runtime requirements, while the partition test can run without
coremltools, Xcode, or a CoreML runtime.
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.
### 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.
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>
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>
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>
## 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
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.
`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
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.
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.
### 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.
## 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.
## 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>
### 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.
## 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`
## 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
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.
## 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
## 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
## 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
```
## 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
```
## 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)
Shushi Hong