Aharrypotter commented on code in PR #19651:
URL: https://github.com/apache/tvm/pull/19651#discussion_r3333912058
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -6926,6 +6928,240 @@ def
test_stablehlo_options_missing_payload_unsupported():
_load_model_from_buffer(buf)
+def _build_stablehlo_rng_model(algorithm, state_len, out_shape,
out_tensor_type, const_state=None):
+ """Build a STABLEHLO_RNG_BIT_GENERATOR model.
+
+ When ``const_state`` is provided, the uint64 initial state is embedded as a
+ constant tensor (no graph input); otherwise it is a graph input.
+ """
+ builder = flatbuffers.Builder(1024)
+
+ _tfl_stablehlo_rng_opts.StablehloRngBitGeneratorOptionsStart(builder)
+
_tfl_stablehlo_rng_opts.StablehloRngBitGeneratorOptionsAddAlgorithm(builder,
algorithm)
+ rng_opts =
_tfl_stablehlo_rng_opts.StablehloRngBitGeneratorOptionsEnd(builder)
+
+ rng_builtin =
_get_stablehlo_builtin_operator("STABLEHLO_RNG_BIT_GENERATOR")
+ rng_code = _build_operator_code(builder, rng_builtin)
+
+ main_tensors = [
+ _build_tensor(builder, 0, [state_len],
tensor_type=_tfl_tensor_type.UINT64),
+ _build_tensor(builder, 1, [state_len],
tensor_type=_tfl_tensor_type.UINT64),
+ _build_tensor(builder, 2, list(out_shape),
tensor_type=out_tensor_type),
+ ]
+ rng_op = _build_operator(
+ builder,
+ 0,
+ [0],
+ [1, 2],
+
builtin_options2_type=_tfl_builtin_options2.StablehloRngBitGeneratorOptions,
+ builtin_options2=rng_opts,
+ )
+ main_subgraph = _build_subgraph(
+ builder,
+ tensors=main_tensors,
+ operators=[rng_op],
+ inputs=[] if const_state is not None else [0],
+ outputs=[1, 2],
+ )
+
+ state_data = None
+ if const_state is not None:
+ state_data = np.array(const_state, dtype="uint64").tobytes()
+ buffers = [
+ _build_buffer(builder, data=state_data),
+ _build_buffer(builder),
+ _build_buffer(builder),
+ ]
+ return _finish_tflite_model(
+ builder,
+ subgraph=main_subgraph,
+ operator_codes=[rng_code],
+ buffers=buffers,
+ )
+
+
+def _run_stablehlo_rng_model(algorithm, state_len, out_shape, out_tensor_type,
init_state):
+ """Import, compile, and execute an RNG model, returning (output_state,
output)."""
+ buf = _build_stablehlo_rng_model(algorithm, state_len, out_shape,
out_tensor_type)
+ mod = _load_model_from_buffer(buf)
+ ex = tvm.compile(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ result = vm["main"](tvm.runtime.tensor(np.array(init_state,
dtype="uint64")))
+ return result[0].numpy(), result[1].numpy()
+
+
+# Expected vectors are taken verbatim from the TFLite runtime kernel test
+# (tensorflow/lite/kernels/rng_bit_generator_test.cc), guaranteeing bit-exact
parity.
+_RNG_THREEFRY_EXPECTED = {
+ "int32": [43444564, -2144348869, -315321645, -549236733, 1672743891,
-54463903],
+ "uint32": [43444564, 2150618427, 3979645651, 3745730563, 1672743891,
4240503393],
+ "int64": [
+ -9209908263526143660,
+ -2358953802017238317,
+ -233920680524772397,
+ 2658481902456610144,
+ -2022031683723149139,
+ -2324041912354448873,
+ ],
+ "uint64": [
+ 9236835810183407956,
+ 16087790271692313299,
+ 18212823393184779219,
+ 2658481902456610144,
+ 16424712389986402477,
+ 16122702161355102743,
+ ],
+}
+_RNG_THREEFRY_STATE = {"int32": [1, 5], "uint32": [1, 5], "int64": [1, 8],
"uint64": [1, 8]}
+_RNG_PHILOX_EXPECTED = {
+ "int32": [-263854262, 1366700262, 495645701, -1243243882, 89414891,
1917262711],
+ "uint32": [4031113034, 1366700262, 495645701, 3051723414, 89414891,
1917262711],
+ "int64": [
+ 5869932932755744586,
+ -5339691813646437371,
+ 8234580641674714347,
+ 2641225993340350124,
+ 1962472297844690804,
+ -3580856229565614135,
+ ],
+ "uint64": [
+ 5869932932755744586,
+ 13107052260063114245,
+ 8234580641674714347,
+ 2641225993340350124,
+ 1962472297844690804,
+ 14865887844143937481,
+ ],
+}
+_RNG_PHILOX_STATE = {
+ "int32": [1, 4, 3],
+ "uint32": [1, 4, 3],
+ "int64": [1, 5, 3],
+ "uint64": [1, 5, 3],
+}
+
+
[email protected](
+ "out_dtype,out_tensor_type",
+ [
+ ("int32", _tfl_tensor_type.INT32),
+ ("uint32", _tfl_tensor_type.UINT32),
+ ("int64", _tfl_tensor_type.INT64),
+ ("uint64", _tfl_tensor_type.UINT64),
+ ],
+)
+def test_stablehlo_rng_bit_generator_threefry(out_dtype, out_tensor_type):
+ """TFLite STABLEHLO_RNG_BIT_GENERATOR THREEFRY matches the runtime kernel
bit-exactly."""
+ state, output = _run_stablehlo_rng_model(
+ _tfl_rng_algorithm.THREEFRY, 2, [2, 3], out_tensor_type, [1, 2]
+ )
+ assert output.flatten().tolist() == _RNG_THREEFRY_EXPECTED[out_dtype]
+ assert state.tolist() == _RNG_THREEFRY_STATE[out_dtype]
+
+
[email protected](
+ "out_dtype,out_tensor_type",
+ [
+ ("int32", _tfl_tensor_type.INT32),
+ ("uint32", _tfl_tensor_type.UINT32),
+ ("int64", _tfl_tensor_type.INT64),
+ ("uint64", _tfl_tensor_type.UINT64),
+ ],
+)
+def test_stablehlo_rng_bit_generator_philox(out_dtype, out_tensor_type):
+ """TFLite STABLEHLO_RNG_BIT_GENERATOR PHILOX matches the runtime kernel
bit-exactly."""
+ state, output = _run_stablehlo_rng_model(
+ _tfl_rng_algorithm.PHILOX, 3, [2, 3], out_tensor_type, [1, 2, 3]
+ )
+ assert output.flatten().tolist() == _RNG_PHILOX_EXPECTED[out_dtype]
+ assert state.tolist() == _RNG_PHILOX_STATE[out_dtype]
+
+
+def test_stablehlo_rng_bit_generator_default_matches_philox():
+ """TFLite STABLEHLO_RNG_BIT_GENERATOR DEFAULT resolves to the PHILOX
algorithm."""
+ state, output = _run_stablehlo_rng_model(
+ _tfl_rng_algorithm.DEFAULT, 3, [2, 3], _tfl_tensor_type.INT32, [1, 2,
3]
+ )
+ assert output.flatten().tolist() == _RNG_PHILOX_EXPECTED["int32"]
+ assert state.tolist() == _RNG_PHILOX_STATE["int32"]
+
+
+def test_stablehlo_rng_bit_generator_deterministic():
+ """Re-running the imported RNG kernel yields identical bit-exact output."""
+ buf = _build_stablehlo_rng_model(_tfl_rng_algorithm.PHILOX, 3, [3, 3],
_tfl_tensor_type.INT32)
+ mod = _load_model_from_buffer(buf)
+ ex = tvm.compile(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ init = tvm.runtime.tensor(np.array([7, 8, 9], dtype="uint64"))
+ first = vm["main"](init)
+ second = vm["main"](init)
+ np.testing.assert_equal(first[1].numpy(), second[1].numpy())
+ np.testing.assert_equal(first[0].numpy(), second[0].numpy())
+
+
+def test_stablehlo_rng_bit_generator_constant_state():
+ """A constant uint64 initial state imports and stays bit-exact (no graph
input)."""
+ buf = _build_stablehlo_rng_model(
+ _tfl_rng_algorithm.THREEFRY, 2, [2, 3], _tfl_tensor_type.INT32,
const_state=[1, 2]
+ )
+ mod = _load_model_from_buffer(buf)
+ assert len(mod["main"].params) == 0
+ ex = tvm.compile(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ result = vm["main"]()
+ assert result[1].numpy().flatten().tolist() ==
_RNG_THREEFRY_EXPECTED["int32"]
+ assert result[0].numpy().tolist() == _RNG_THREEFRY_STATE["int32"]
Review Comment:
I think the `compile + VirtualMachine` calls in these RNG tests are
intentional, though they cover slightly different things.
`STABLEHLO_RNG_BIT_GENERATOR` is different from the other StableHLO TFLite
ops in this file: the importer generates a custom TIR implementation of the
TFLite runtime RNG algorithm, rather than lowering to an existing Relax op. A
structural-equality test would only prove that the generated module has the
expected shape; it would not catch mistakes in the Threefry/Philox rounds,
counter advancement, word packing, or signed/unsigned reinterpretation.
The nightly E2E path in `_verify_model` also does not cover this case. These
tests build a raw TFLite flatbuffer directly, so there is no TensorFlow
callable/reference path for nightly CI to run and compare. The expected vectors
below are copied from the TFLite runtime kernel tests, so compiling and
executing the imported module here is what actually checks bit-exact parity
with TFLite.
--
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]
