[PRNG] Add check to PRNG to make sure that unsigned integer arithmetic is wrapping

Author: tkonolige

python/tvm/topi/random/kernel.py

@@ -234,6 +234,18 @@ def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr):
         out_gen = irb.buffer_ptr(out_gen_ptr)
         out_array = irb.buffer_ptr(out_array_ptr)
 
+        # Check that unsigned arithmetic wraps, as it is required to implement threefry correctly.
+        irb.emit(
+            tvm.tir.AssertStmt(
+                tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64") + tvm.tir.const(1, "uint64")
+                == tvm.tir.const(0, "uint64"),
+                tvm.tir.StringImm(
+                    "Unsigned interger arithmetic is not wrapping, but threefry requires wrapping."
+                ),
+                tvm.tir.Evaluate(0),
+            )
+        )
+
         # Create a temporary array to hold the generator state we will use to create the random
         # numbers. We cannot use gen because we may need to update the key + path if there is not
         # enough room in the counter.