Aharrypotter opened a new pull request, #19651:
URL: https://github.com/apache/tvm/pull/19651
## Summary
This PR adds Relax TFLite frontend support for the TFLite builtin
`STABLEHLO_RNG_BIT_GENERATOR` operator.
Unlike most StableHLO builtins, the TFLite runtime
(`tensorflow/lite/kernels/rng_bit_generator.cc`) implements this op as a
real,
deterministic counter-based PRNG, so the importer must reproduce it
bit-exactly
rather than map it to an existing op:
- one uint64 1-D `initial_state` input, two outputs — uint64 `output_state`
and
the random-bit `output` (int32 / int64 / uint32 / uint64);
- `algorithm` in `{DEFAULT, PHILOX, THREEFRY}`, where `DEFAULT` resolves to
`PHILOX`;
- Random123 Threefry2x32 (20 rounds) and Philox4x32 (10 rounds) with the
fixed
constants from `rng_util.cc`;
- state-length constraints: `THREEFRY` requires `u64[2]`, `PHILOX`/`DEFAULT`
require `u64[2]` or `u64[3]`.
## Design
TVM/Relax has no matching RNG primitive, so the converter generates a TIR
kernel
that mirrors the runtime and emits it through `relax.call_tir` with two
outputs.
The kernel:
- reinterprets the uint64 state as uint32 words and advances a 64-bit block
counter (`final counter = initial_state[1] + num_blocks`);
- runs the selected algorithm per block with all round state materialized
into
local buffers, which keeps the generated IR linear instead of an
exponentially
nested expression tree;
- packs the produced uint32 words back into the output dtype, and writes the
updated state (key unchanged, counter advanced, Philox `u64[3]` tail passed
through) — the only state behaviour the runtime relies on.
The kernel is an `s_tir` PrimFunc wrapped in a single opaque structured
block so
it remains a well-formed block-structured function for the Relax pipeline
(e.g. `HasReshapePattern`). `get_tensor_type_str` and the input
`_decode_type`
map are extended with uint32/uint64 so the uint64 state imports correctly.
Unsupported inputs raise a precise `OpNotImplemented` (non-uint64 / non-1-D
state, mismatched output-state shape, unsupported output dtype, unknown
algorithm, per-algorithm state-length violations).
## Operator Support
| Operator | TFLite options | Relax lowering | Supported subset |
|---|---|---|---|
| `STABLEHLO_RNG_BIT_GENERATOR` |
`StablehloRngBitGeneratorOptions.Algorithm()` from `BuiltinOptions2` |
`call_tir` to a generated bit-exact TIR kernel | THREEFRY (`u64[2]`) and
PHILOX/DEFAULT (`u64[2]`/`u64[3]`); int32/int64/uint32/uint64 output |
## Tests
Tests build minimal RNG flatbuffers, compile, and execute them, comparing the
output and updated state against the verbatim expected vectors from the
TFLite
runtime kernel test (`rng_bit_generator_test.cc`).
| Test | Coverage |
|---|---|
| `test_stablehlo_rng_bit_generator_threefry` | THREEFRY bit-exact, all 4
output dtypes |
| `test_stablehlo_rng_bit_generator_philox` | PHILOX bit-exact, all 4 output
dtypes |
| `test_stablehlo_rng_bit_generator_default_matches_philox` | DEFAULT
resolves to PHILOX |
| `test_stablehlo_rng_bit_generator_deterministic` | run-to-run
bit-identical output |
| `test_stablehlo_rng_bit_generator_unsupported_output_dtype` | output dtype
guard |
| `test_stablehlo_rng_bit_generator_threefry_invalid_state_unsupported` |
THREEFRY `u64[2]` state guard |
| `test_stablehlo_rng_bit_generator_non_uint64_state_unsupported` | uint64
state guard |
Local validation:
```bash
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 rng_bit_generator -q
python -m pytest \
tests/python/relax/test_frontend_tflite.py \
-k stablehlo -q
```
Result:
```text
ruff check: All checks passed
rng_bit_generator tests: 13 passed
stablehlo tests: 96 passed
```
## References
- Issue #19519 item I: remaining StableHLO operators in TFLite
- `tensorflow/lite/kernels/rng_bit_generator.cc`, `rng_util.cc`,
`rng_bit_generator_test.cc`
--
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]
