[Relax][Frontend][TFLite] Support STABLEHLO_RNG_BIT_GENERATOR

[Aharrypotter]

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:

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:

ruff check: All checks passed
rng_bit_generator tests: 13 passed
stablehlo tests: 96 passed

References

• Issue [Tracking Issue][TFLite] Remaining builtin operator coverage beyond #19412 #19519 item I: remaining StableHLO operators in TFLite
• tensorflow/lite/kernels/rng_bit_generator.cc, rng_util.cc,
  rng_bit_generator_test.cc

[gemini-code-assist]

Code Review

This pull request adds support for the STABLEHLO_RNG_BIT_GENERATOR operator in the TVM Relax TFLite frontend, including support for UINT32 and UINT64 tensor types. It implements bit-exact TIR kernels for both the Threefry and Philox counter-based PRNG algorithms and adds corresponding unit tests. Feedback on the changes highlights a few issues in the generated TIR kernels, including an invalid import of tirx instead of tir, unsupported s_tir arguments in the @T.prim_func decorators, and an opportunity to avoid generating empty loops when the state length is 2.

[gemini-code-assist]

There is a typo in the import statement. tirx does not exist in tvm.script.parser. It should be imported from tvm.script as tir.

Suggested change

	from tvm.script.parser import tirx as T
	from tvm.script import tir as T

[gemini-code-assist]

The @T.prim_func decorator does not accept an s_tir argument. This will raise a TypeError at runtime. Please remove s_tir=True.

Suggested change

	@T.prim_func(private=True, s_tir=True)
	@T.prim_func(private=True)

[gemini-code-assist]

The @T.prim_func decorator does not accept an s_tir argument. This will raise a TypeError at runtime. Please remove s_tir=True.

Suggested change

	@T.prim_func(private=True, s_tir=True)
	@T.prim_func(private=True)

[gemini-code-assist]

When state_len is 2, state_len - 2 is 0, which generates a loop with 0 iterations (T.serial(0)). To avoid generating redundant/empty loops in the TIR kernel, we can conditionally generate this loop using a Python if statement.

Suggested change

	for tail in T.serial(state_len - 2):
	output_state[tail + 2] = initial_state[tail + 2]
	if state_len > 2:
	for tail in T.serial(state_len - 2):
	output_state[tail + 2] = initial_state[tail + 2]