Skip to content

Commit c0406a5

Browse files
authored
[Relax][Frontend][TFLite] Add initial StableHLO builtin operator support (#19536)
## Summary This PR adds initial Relax TFLite frontend support for 29 StableHLO builtin operators from #19519 item I. The covered subset includes pure elementwise ops, BuiltinOptions2 / metadata-based ops, simple shape-manipulation ops, and a take-equivalent subset of `STABLEHLO_GATHER`. StableHLO builtins carry no TFLite-specific quantization or fused-activation metadata, so the implementation uses dedicated converter helpers that bypass the existing TFLite elemwise/QNN code paths. Relates to #19519. ## Changes 1. **Zero-attribute elementwise helpers** - Add `_convert_stablehlo_unary`, `_convert_stablehlo_binary`, and `_convert_stablehlo_ternary` for pure elementwise mapping. - Register 20 ops: unary (`ABS`, `NEGATE`, `COSINE`, `EXPONENTIAL`, `FLOOR`, `LOG`, `LOGISTIC`, `RSQRT`, `TANH`), binary (`ADD`, `SUBTRACT`, `MULTIPLY`, `DIVIDE`, `MAXIMUM`, `MINIMUM`, `POWER`), ternary (`SELECT` → `R.where`), and dtype-dispatched bitwise/logical ops (`AND` / `OR` → logical ops for bool or bitwise ops for integer, `SHIFT_LEFT` → `R.left_shift` for integer). 2. **BuiltinOptions2 infrastructure** - Add `_get_stablehlo_options` helper for parsing `BuiltinOptions2` flatbuffers with enum validation via `getattr(BuiltinOptions2, options_cls.__name__)`. - Register 6 ops: `CONVERT` → `R.astype`, `CLAMP` → `R.minimum(R.maximum(...))`, `CONCATENATE` → `R.concat`, `BROADCAST_IN_DIM` → `R.reshape` + `R.broadcast_to`, `IOTA` → `R.arange` + `R.broadcast_to`, and `COMPARE` → 6 comparison directions (`TOTALORDER` raises `OpNotImplemented`). 3. **Shape-manipulation ops** - `PAD` → `R.nn.pad` in constant mode. The initial PAD path supports non-negative edge padding with zero interior padding and a constant scalar padding value. Interior padding, negative padding, and dynamic padding values raise `OpNotImplemented`. - `DYNAMIC_SLICE` → `R.dynamic_strided_slice`. The initial path supports constant, in-bound start indices only. Runtime start indices and out-of-bounds StableHLO clamping semantics are deferred. 4. **Indexing op** - `GATHER` → `R.take` for the take-equivalent subset only. - Parses the relevant `StablehloGatherOptions` attributes needed to validate this subset: `offset_dims`, `collapsed_slice_dims`, `start_index_map`, `index_vector_dim`, and `slice_sizes`. - Validates the gather axis, collapsed dims, offset dims, slice sizes, and output shape against the expected `R.take` layout. Multi-dimensional and non-take-equivalent gather patterns raise `OpNotImplemented`. 5. **Not included** - `STABLEHLO_RESHAPE`, `STABLEHLO_TRANSPOSE`, and `STABLEHLO_SLICE` are left to another contributor who expressed interest in those ops. - The remaining Issue #19519 StableHLO items are deferred to follow-up PRs: `CBRT`, `REMAINDER`, `SCATTER`, `CONVOLUTION`, `DOT_GENERAL`, `REDUCE`, `REDUCE_WINDOW`, `DYNAMIC_UPDATE_SLICE`, `COMPOSITE`, `CUSTOM_CALL`, `RNG_BIT_GENERATOR`, `SORT`, and `WHILE`. - More general or multi-dimensional `STABLEHLO_GATHER` patterns are also deferred to follow-up work. ## Testing All tests use manually-built minimal TFLite flatbuffers with `tvm.ir.assert_structural_equal`. BuiltinOptions2 ops construct their options via the FlatBuffers schema API, modeled after the existing DILATE test pattern. ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k stablehlo -q ``` ## Result - 29 StableHLO operators registered in the Relax TFLite frontend. - 44 StableHLO test cases covering all registered ops, including structural-equal tests and unsupported/error-path checks: - `COMPARE` with `TOTALORDER` - `PAD` with interior padding, negative padding, and dynamic padding values - `DYNAMIC_SLICE` with runtime starts and out-of-bounds starts - non-take-equivalent or multi-dimensional `GATHER` - All StableHLO TFLite frontend tests pass locally. ## References - Issue #19519 item I: StableHLO operators in TFLite - Related PR #19481: DILATE operator mapping, the first use of BuiltinOptions2 in the TFLite frontend tests
1 parent c309e4e commit c0406a5

2 files changed

Lines changed: 1626 additions & 0 deletions

File tree

0 commit comments

Comments
 (0)