Aharrypotter commented on code in PR #19651:
URL: https://github.com/apache/tvm/pull/19651#discussion_r3333774644
##########
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
+
+ total = 1
+ for dim in out_shape:
+ total *= int(dim)
+ is_64bit = out_dtype in ("int64", "uint64")
+ block_words = 2 if algorithm == "threefry" else 4
+ out_word_count = total * (2 if is_64bit else 1)
+ num_blocks = (out_word_count + block_words - 1) // block_words
+ writes_per_block = block_words // (2 if is_64bit else 1)
+ parity = _STABLEHLO_RNG_THREEFRY_PARITY
+ mul_a, mul_b = _STABLEHLO_RNG_PHILOX_MUL_A, _STABLEHLO_RNG_PHILOX_MUL_B
+ weyl_a, weyl_b = _STABLEHLO_RNG_PHILOX_WEYL_A, _STABLEHLO_RNG_PHILOX_WEYL_B
+
+ def _u32(value):
+ return T.Cast("uint32", value)
+
+ def _u64(value):
+ return T.Cast("uint64", value)
+
+ def _store_value(words, write_index):
+ # Pack the generated uint32 words into one output element,
reinterpreting
+ # the bit pattern into the (possibly signed) output dtype.
+ if is_64bit:
+ low = _u64(words[2 * write_index])
+ high = _u64(words[2 * write_index + 1])
+ return T.reinterpret(out_dtype, low | (high << T.uint64(32)))
+ return T.reinterpret(out_dtype, words[write_index])
+
+ if algorithm == "threefry":
+
+ @T.prim_func(private=True, s_tir=True)
Review Comment:
`s_tir=True` is also intentional here.
The PrimFunc body is wrapped in:
```python
with T.sblock("rng_bit_generator"):
...
```
That block-structured form is needed by the Relax pipeline. In particular,
downstream checks such as `HasReshapePattern` expect the PrimFunc body to be an
`SBlockRealizeNode`. Without `s_tir=True`, the parser rejects this style of
structured block usage, so simply removing the flag is not a valid fix.
There is existing in-tree precedent for this form as well, for example
kernels using:
```python
@T.prim_func(private=True, s_tir=True)
```
in `python/tvm/relax/frontend/nn/op.py`.
So the combination of `tirx`, `T.sblock(...)`, and `@T.prim_func(...,
s_tir=True)` is deliberate: it keeps the custom imperative RNG kernel
block-structured and acceptable to the Relax pipeline. I would not apply the
suggested change.
--
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]
