[GitHub] [tvm] zhuzilin commented on a change in pull request #8085: [Relay][PRNG] Support generating data of any shape in threefry_generate

[GitBox <gi...@apache.org>]

zhuzilin commented on a change in pull request #8085:
URL: https://github.com/apache/tvm/pull/8085#discussion_r636583911



##########
File path: python/tvm/topi/random/kernel.py
##########
@@ -306,7 +309,10 @@ def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr):
         out_gen[4] = tmp[4]  # path stays the same
         out_gen[5] = tmp[5]
         out_gen[6] = tir.const(0, dtype=gen.dtype)  # unused, leave it as 0
-        out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len)  # increment counter
+        with irb.if_scope(out_len % 4 != 0):
+            out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len % 4)
+        with irb.else_scope():
+            out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len)  # increment counter

Review comment:
       @tkonolige I changed the increment to `out_len % 4` here because I already update the `tmp[7]` above with:
   ```python
   tmp[7] = tmp[7] + tir.Cast(gen.dtype, out_len // 4 * 4)
   ```
   The reason for the separate update is that I think we need to update the key before the second `threefry`.
   
   However, I update the second update from `tir.Cast(gen.dtype, out_len % 4)` to `tir.Cast(gen.dtype, 4)`, because 4 random numbers are actually generated (though partially discarded).




-- 
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org