Aharrypotter opened a new pull request, #19639:
URL: https://github.com/apache/tvm/pull/19639
## Summary
This PR adds incremental Relax TFLite frontend support for the TFLite
resource
variable and static hashtable operator family:
- `VAR_HANDLE`
- `ASSIGN_VARIABLE`
- `READ_VARIABLE`
- `HASHTABLE`
- `HASHTABLE_IMPORT`
- `HASHTABLE_FIND`
- `HASHTABLE_SIZE`
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 variables or query initialized tables. This PR supports that constrained
initialization pattern without introducing general mutable runtime state into
Relax.
The implementation is split into two reviewable commits:
- Resource variable subset: `VAR_HANDLE`, `ASSIGN_VARIABLE`, `READ_VARIABLE`
- Hashtable subset: `HASHTABLE`, `HASHTABLE_IMPORT`, `HASHTABLE_FIND`,
`HASHTABLE_SIZE`
The supported subset is intentionally conservative:
- Resource variables can be assigned only inside a supported `CALL_ONCE` init
subgraph.
- Resource variables can be read from the main graph only after supported
initialization.
- Hashtables can be imported only inside a supported `CALL_ONCE` init
subgraph
using constant key/value buffers.
- Hashtable lookup and size are lowered as pure Relax expressions after
static
initialization is captured.
- Runtime mutation, dynamic hashtable import, and full resource-state
threading
remain unsupported.
## Design
### Shared Initialization State
The frontend now keeps resource initialization data in shared conversion
state:
- `conversion_state["resource_values"]`
- `conversion_state["hashtable_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 handle
maps:
- `self.resource_handles`
- `self.hashtable_handles`
The local maps are keyed by TFLite tensor name and map the tensor handle to a
stable resource key. Resource variables use `container + shared_name` from
`VarHandleOptions` when present, falling back to the handle tensor name.
Hashtables use `HashtableOptions.table_id` when present, also 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`
- `HASHTABLE`
- `HASHTABLE_IMPORT`
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.
### Static Hashtables
`HASHTABLE` is declarative and registers the output table handle in the
current
converter's local `hashtable_handles` map.
`HASHTABLE_IMPORT` is accepted only while converting a supported `CALL_ONCE`
init subgraph. The imported keys and values must come from constant TFLite
buffers. The frontend stores:
- keys expression
- values expression
- table size
- per-key value shape
in shared `conversion_state["hashtable_values"]`.
`HASHTABLE_SIZE` lowers to a scalar Relax constant for statically initialized
tables.
`HASHTABLE_FIND` lowers to pure Relax tensor operations:
```text
expand_dims(query_keys)
equal(query_keys, table_keys)
astype(bool_matches, int32)
argmax(matches, axis=-1)
take(table_values, matched_indices, axis=0)
max(matches, axis=-1) > 0
broadcast default value
where(has_match, selected_value, default_value)
```
The `where` mask makes the all-mismatch case correct even though `argmax`
over
all-zero matches selects index 0. For vector or higher-rank per-key values,
the
match mask is expanded and broadcast to the output shape.
## 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 only | main graph
and supported `CALL_ONCE` init subgraphs |
| `HASHTABLE_IMPORT` | `HashtableImportOptions` | store static table
keys/values in shared importer state | supported `CALL_ONCE` init subgraphs
only, constant keys/values |
| `HASHTABLE_FIND` | `HashtableFindOptions` | pure Relax lookup expression |
one-dimensional query keys, statically imported table |
| `HASHTABLE_SIZE` | `HashtableSizeOptions` | scalar Relax constant |
statically imported table |
## Safety Checks
- `ASSIGN_VARIABLE` outside `CALL_ONCE` initialization raises
`OpNotImplemented`.
- `HASHTABLE_IMPORT` outside `CALL_ONCE` initialization raises
`OpNotImplemented`.
- `READ_VARIABLE` without supported initialization raises
`OpNotImplemented`.
- `HASHTABLE_FIND` and `HASHTABLE_SIZE` require a table initialized by a
supported `CALL_ONCE` subgraph.
- `CALL_ONCE` init subgraphs with inputs or outputs remain unsupported.
- `CALL_ONCE` init subgraphs containing operators outside the supported
initialization allowlist remain unsupported.
- `HASHTABLE_IMPORT` requires one-dimensional keys and a values tensor whose
leading dimension matches the number of keys.
- `HASHTABLE_FIND` validates that vector or higher-rank table values match
the
suffix of the declared output shape.
## Not Included
- Runtime `ASSIGN_VARIABLE` mutation in the main graph.
- Runtime or dynamic `HASHTABLE_IMPORT`.
- General resource-state threading through Relax function parameters and
returns.
- Cross-subgraph resource handle aliasing beyond the static
`container/shared_name` and `table_id` matching pattern.
- Multiple runtime writes with ordering semantics.
- Full TFLite resource semantics for mutable variables, tables, and aliases.
- Partial or dynamic initialization subgraphs that emit Relax bindings.
## Tests
The tests manually build minimal TFLite flatbuffers and compare imported
Relax
IR with `tvm.ir.assert_structural_equal`. Unsupported patterns use
`pytest.raises`. The all-mismatch hashtable lookup also runs through the
Relax
VM to verify the `where` default-value path.
| 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` | static hashtable import and
scalar-value lookup lowering |
| `test_hashtable_call_once_import_find_all_mismatch` | VM execution for all
query keys absent, validating default values |
| `test_hashtable_call_once_import_find_vector_values` | static hashtable
lookup with vector per-key values |
| `test_hashtable_call_once_import_size` | `HASHTABLE_SIZE` lowering to
scalar const |
| `test_hashtable_import_main_subgraph_unsupported` | runtime/main graph
`HASHTABLE_IMPORT` guard |
| `test_hashtable_size_uninitialized_unsupported` | `HASHTABLE_SIZE` without
supported initialization 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 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 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: 7 passed
full test_frontend_tflite.py: 474 passed
```
## Limitations
- This PR intentionally models only static initialization captured during
import. It does not implement a runtime resource store.
- Hashtable keys and values must be constant at import time.
- `HASHTABLE_FIND` currently assumes query keys are one-dimensional. Per-key
values may be scalar, vector, or higher-rank tensors as long as the
imported
values tensor has shape `[num_keys, ...]`.
## References
- Issue #19519 item G: TFLite resource / variable / hashtable operators
- PR #19616: TFLite control-flow / multi-subgraph support
- Related operator family: `CALL_ONCE` initialization subgraphs
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
