Commit c0406a5
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 tests1 parent c309e4e commit c0406a5
2 files changed
Lines changed: 1626 additions & 0 deletions
File tree
- python/tvm/relax/frontend/tflite
- tests/python/relax
0 commit comments