Aharrypotter commented on code in PR #19651:
URL: https://github.com/apache/tvm/pull/19651#discussion_r3333761377
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -7347,6 +7418,161 @@ def get_tensor_shape(self, tensor_wrapper):
)
+# Constants for the Random123 counter-based PRNGs used by
STABLEHLO_RNG_BIT_GENERATOR,
+# matching tensorflow/lite/kernels/rng_util.cc.
+_STABLEHLO_RNG_THREEFRY_PARITY = 0x1BD11BDA
+_STABLEHLO_RNG_PHILOX_MUL_A = 0xD2511F53
+_STABLEHLO_RNG_PHILOX_MUL_B = 0xCD9E8D57
+_STABLEHLO_RNG_PHILOX_WEYL_A = 0x9E3779B9
+_STABLEHLO_RNG_PHILOX_WEYL_B = 0xBB67AE85
+
+
+def _build_stablehlo_rng_bit_generator_primfunc(algorithm, state_len,
out_dtype, out_shape):
+ """Build a bit-exact TIR kernel for STABLEHLO_RNG_BIT_GENERATOR.
+
+ Mirrors the TFLite runtime kernel
(tensorflow/lite/kernels/rng_bit_generator.cc),
+ implementing the Random123 Threefry2x32 (20 rounds) and Philox4x32 (10
rounds)
+ counter-based PRNGs. The kernel reinterprets the uint64 state as uint32
words,
+ advances a 64-bit block counter, and packs the generated words into the
output
+ tensor. The updated state keeps the key unchanged and only advances the
counter,
+ which is the only behaviour the runtime relies on.
+ """
+ from tvm.script.parser import tirx as T
Review Comment:
Thanks for checking this, but this import is intentional for this branch.
The generated RNG kernel uses the TIRx script parser APIs, including
`T.sblock(...)`. In this codebase, the working import for that parser is:
```python
from tvm.script.parser import tirx as T
```
Replacing it with `from tvm.script import tir as T` is not equivalent here
and breaks this branch. The test file already uses the same TIRx import style,
so this is consistent with the local code being exercised by the tests.
For this RNG lowering, the importer needs to generate a block-structured TIR
kernel through TIRx, so I would keep the `tirx` import.
--
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]
