Aharrypotter commented on code in PR #19651:
URL: https://github.com/apache/tvm/pull/19651#discussion_r3333785500
##########
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)
+ def kernel(
+ initial_state: T.Buffer((state_len,), "uint64"),
+ output_state: T.Buffer((state_len,), "uint64"),
+ output: T.Buffer(out_shape, out_dtype),
+ ):
+ # A single opaque structured block keeps the imperative kernel as a
+ # well-formed block-structured PrimFunc, as required by the Relax
+ # pipeline (e.g. HasReshapePattern).
+ with T.sblock("rng_bit_generator"):
+ state_key = initial_state[0]
+ state_counter = initial_state[1]
+ key_0 = _u32(state_key & T.uint64(0xFFFFFFFF))
+ key_1 = _u32(state_key >> T.uint64(32))
+ output_state[0] = state_key
+ output_state[1] = state_counter + T.uint64(num_blocks)
+ out_flat = T.decl_buffer((total,), out_dtype, data=output.data)
+ keys = T.decl_buffer((3,), "uint32", scope="local")
+ rotations = T.decl_buffer((8,), "uint32", scope="local")
+ ctr = T.decl_buffer((2,), "uint32", scope="local")
+ keys[0] = key_0
+ keys[1] = key_1
+ keys[2] = key_0 ^ key_1 ^ T.uint32(parity)
+ rotations[0] = T.uint32(13)
+ rotations[1] = T.uint32(15)
+ rotations[2] = T.uint32(26)
+ rotations[3] = T.uint32(6)
+ rotations[4] = T.uint32(17)
+ rotations[5] = T.uint32(29)
+ rotations[6] = T.uint32(16)
+ rotations[7] = T.uint32(24)
+ for block in T.serial(num_blocks):
+ counter = state_counter + _u64(block)
+ ctr[0] = _u32(counter & T.uint64(0xFFFFFFFF)) + key_0
+ ctr[1] = _u32(counter >> T.uint64(32)) + key_1
+ for group in T.serial(5):
+ for step in T.serial(4):
+ rot = rotations[(group * 4 + step) % 8]
+ ctr[0] = ctr[0] + ctr[1]
+ ctr[1] = (ctr[1] << rot) | (ctr[1] >>
(T.uint32(32) - rot))
+ ctr[1] = ctr[1] ^ ctr[0]
+ ctr[0] = ctr[0] + keys[(group + 1) % 3]
+ ctr[1] = ctr[1] + keys[(group + 2) % 3] + _u32(group +
1)
+ for write_index in T.serial(writes_per_block):
+ element = block * writes_per_block + write_index
+ if element < total:
+ out_flat[element] = _store_value(ctr, write_index)
+
+ return kernel
+
+ @T.prim_func(private=True, s_tir=True)
Review Comment:
The same explanation as above.
--
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]
