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