gemini-code-assist[bot] commented on code in PR #19651:
URL: https://github.com/apache/tvm/pull/19651#discussion_r3331739181
##########
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)
+ def kernel(
+ initial_state: T.Buffer((state_len,), "uint64"),
+ output_state: T.Buffer((state_len,), "uint64"),
+ output: T.Buffer(out_shape, out_dtype),
+ ):
+ 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)
+ for tail in T.serial(state_len - 2):
+ output_state[tail + 2] = initial_state[tail + 2]
Review Comment:
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.
```suggestion
if state_len > 2:
for tail in T.serial(state_len - 2):
output_state[tail + 2] = initial_state[tail + 2]
```
##########
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:
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`.
```suggestion
from tvm.script import tir as T
```
##########
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 `@T.prim_func` decorator does not accept an `s_tir` argument. This will
raise a `TypeError` at runtime. Please remove `s_tir=True`.
```suggestion
@T.prim_func(private=True)
```
##########
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:
The `@T.prim_func` decorator does not accept an `s_tir` argument. This will
raise a `TypeError` at runtime. Please remove `s_tir=True`.
```suggestion
@T.prim_func(private=True)
```
--
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]
