Files
apache--tvm/python
HoYi c0406a54bc [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
2026-05-11 22:05:20 +08:00
..