Skip to content

Commit

Permalink
[PRNG] Add check to PRNG to make sure that unsigned integer arithmeti…
Browse files Browse the repository at this point in the history
…c is wrapping
  • Loading branch information
tkonolige committed Jan 14, 2021
1 parent d7a9a7c commit e6bd9d3
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions python/tvm/topi/random/kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit e6bd9d3

Please sign in to comment.