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