The GitHub Actions job "CI" on tvm.git/main has succeeded. Run started by GitHub user tlopex (triggered by tlopex).
Head commit for run: 99488d992de65ac9e6299548c673c3ca95ef98c2 / HoYi <[email protected]> [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 Report URL: https://github.com/apache/tvm/actions/runs/26704590617 With regards, GitHub Actions via GitBox --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
