Aharrypotter opened a new pull request, #19646:
URL: https://github.com/apache/tvm/pull/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
--
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]
