From 89b3f138328ec9e32b6501557c50831877543e4a Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Thu, 14 Jan 2021 15:43:32 -0800 Subject: [PATCH 1/5] [PRNG] Add check to PRNG to make sure that unsigned integer arithmetic is wrapping --- python/tvm/topi/random/kernel.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 576fd9254a79..52c0d0896654 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -135,7 +135,7 @@ def _threefry( assert key_buf.dtype == counter_buf.dtype, "threefry key and counter must be the same dtype" def mix(a, b, rotation): - x = a + b # TODO should be wrapping + x = a + b # wrapping y = x ^ ((b << rotation) | (b >> (iwidth - rotation))) return [x, y] @@ -167,7 +167,7 @@ def key_schedule(s, i): with irb.for_range(0, out_shape, name="l") as l: # pylint: disable=invalid-name for i in range(nrounds // 4): for j in range(nwords): - out_buf[out_offset + l * nwords + j] += key_schedule(i, j) # TODO wrapping + out_buf[out_offset + l * nwords + j] += key_schedule(i, j) # wrapping for k in range(4): for j in range(nwords // 2): ( @@ -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 integer 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. From 7d8237a22baf6203dc1ec921085781a7b7cc477b Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Fri, 15 Jan 2021 09:49:51 -0800 Subject: [PATCH 2/5] Add threefry_test_wrapping: a manual test for wrapping unsigned arithmetic. --- python/tvm/topi/random/kernel.py | 41 ++++++++++++++++++++++ tests/python/topi/python/test_topi_prng.py | 8 +++++ 2 files changed, 49 insertions(+) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 52c0d0896654..b21a1726faeb 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -20,6 +20,8 @@ from ... import tir from ...tir import ir_builder +import numpy as np + # Threefry PRNG with splitting based on # - J. K. Salmon, M. A. Moraes, R. O. Dror and D. E. Shaw, "Parallel random numbers: As easy as 1, @@ -201,6 +203,13 @@ def threefry_generate(gen, out_shape): then a new generator is created by applying Threefry to the current key, path, and counter. This new generator will have a reset counter. + Warning + ------- + Threeyfry requires that unsigned integer arithmetic wraps on overflow. Currently TVM has no + guarantee of this, so threefry contains an internal assert to check wrapping behavior. This + assert may or may not run depending on your platform, so it is recommended you run + :py:func:`threefry_test_wrapping` to verify wrapping behavior. + Parameters ---------- gen : Tensor[10, uint64] @@ -420,3 +429,35 @@ def gen_ir(gen_ptr, out_left_ptr, out_right_ptr): name="threefry_split", tag="threefry_split", ) + + +def threefry_test_wrapping(target, ctx): + """Test that unsigned arithmetic wraps on overflow. + + Parameters + ---------- + target : tvm.target.Target + Target to run against + ctx : tvm.runtime.TVMContext + Context to run the test on + + Returns + ------- + is_wrapping : bool + Whether or not unsigned integer arithmetic is wrapping for this target, context pair. True + indicates that threefry will work on this platform. + """ + + def ir(out_ptr): + irb = ir_builder.create() + out = irb.buffer_ptr(out_ptr) + out[0] = tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64") + tvm.tir.const(1, "uint64") + return irb.get() + + out = tvm.tir.decl_buffer((1,), dtype="uint64") + f = tvm.te.extern([out.shape], [], lambda ins, outs: ir(outs[0]), dtype="uint64") + s = tvm.te.create_schedule([f.op]) + p = tvm.te.placeholder((1,), "uint64") + out_ary = tvm.nd.array(np.zeros((1,), "uint64")) + tvm.build(s, [p])(out_ary) + return out_ary.asnumpy()[0] == 0 diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py index 43b0494ee6f5..649e5410c147 100644 --- a/tests/python/topi/python/test_topi_prng.py +++ b/tests/python/topi/python/test_topi_prng.py @@ -111,6 +111,14 @@ def test_threefry_generate(target, ctx): ).any(), "Overflowing counter with no space left in path should change state" +@tvm.testing.parametrize_targets +def test_threefry_wrapping(target, ctx): + assert tvm.topi.random.threefry_test_wrapping( + target, ctx + ), f"{target} does not suppport wrapping unsigned integer arithmetic" + + if __name__ == "__main__": test_threefry_split(tvm.target.Target("llvm"), tvm.context("cpu")) test_threefry_generate(tvm.target.Target("llvm"), tvm.context("cpu")) + test_threefry_wrapping(tvm.target.Target("llvm"), tvm.context("cpu")) From 8fa10896ad7097a4a9efd33bf4af07828139678e Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Fri, 15 Jan 2021 10:52:37 -0800 Subject: [PATCH 3/5] fix test to actually run on the target --- python/tvm/topi/random/kernel.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index b21a1726faeb..17ccac587d56 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -17,11 +17,10 @@ """Pseudorandom number kernels.""" import tvm import tvm.topi +import numpy as np from ... import tir from ...tir import ir_builder -import numpy as np - # Threefry PRNG with splitting based on # - J. K. Salmon, M. A. Moraes, R. O. Dror and D. E. Shaw, "Parallel random numbers: As easy as 1, @@ -447,17 +446,22 @@ def threefry_test_wrapping(target, ctx): Whether or not unsigned integer arithmetic is wrapping for this target, context pair. True indicates that threefry will work on this platform. """ + if isinstance(target, str): + target = tvm.target.Target(target) - def ir(out_ptr): + def gen_ir(out_ptr): irb = ir_builder.create() out = irb.buffer_ptr(out_ptr) + if "gpu" in target.keys: + tx = tvm.te.thread_axis("threadIdx.x") + irb.scope_attr(tx, "thread_extent", 1) out[0] = tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64") + tvm.tir.const(1, "uint64") return irb.get() out = tvm.tir.decl_buffer((1,), dtype="uint64") - f = tvm.te.extern([out.shape], [], lambda ins, outs: ir(outs[0]), dtype="uint64") + f = tvm.te.extern([out.shape], [], lambda ins, outs: gen_ir(outs[0]), dtype="uint64", out_buffers=[out]) s = tvm.te.create_schedule([f.op]) p = tvm.te.placeholder((1,), "uint64") - out_ary = tvm.nd.array(np.zeros((1,), "uint64")) - tvm.build(s, [p])(out_ary) + out_ary = tvm.nd.array(np.ones((1,), "uint64"), ctx) + tvm.build(s, [f], target=target)(out_ary) return out_ary.asnumpy()[0] == 0 From d7a2051eea00e07ce9ed4bb308f75071a0f4f414 Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Fri, 15 Jan 2021 11:32:45 -0800 Subject: [PATCH 4/5] formatting --- python/tvm/topi/random/kernel.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 17ccac587d56..72d4edf60b13 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -459,7 +459,9 @@ def gen_ir(out_ptr): return irb.get() out = tvm.tir.decl_buffer((1,), dtype="uint64") - f = tvm.te.extern([out.shape], [], lambda ins, outs: gen_ir(outs[0]), dtype="uint64", out_buffers=[out]) + f = tvm.te.extern( + [out.shape], [], lambda ins, outs: gen_ir(outs[0]), dtype="uint64", out_buffers=[out] + ) s = tvm.te.create_schedule([f.op]) p = tvm.te.placeholder((1,), "uint64") out_ary = tvm.nd.array(np.ones((1,), "uint64"), ctx) From 5271704fbd3f392869537a270296ab805000d3ad Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Fri, 15 Jan 2021 11:39:05 -0800 Subject: [PATCH 5/5] lint --- python/tvm/topi/random/kernel.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 72d4edf60b13..b21db3778744 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -453,8 +453,8 @@ def gen_ir(out_ptr): irb = ir_builder.create() out = irb.buffer_ptr(out_ptr) if "gpu" in target.keys: - tx = tvm.te.thread_axis("threadIdx.x") - irb.scope_attr(tx, "thread_extent", 1) + thread_x = tvm.te.thread_axis("threadIdx.x") + irb.scope_attr(thread_x, "thread_extent", 1) out[0] = tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64") + tvm.tir.const(1, "uint64") return irb.get() @@ -463,7 +463,6 @@ def gen_ir(out_ptr): [out.shape], [], lambda ins, outs: gen_ir(outs[0]), dtype="uint64", out_buffers=[out] ) s = tvm.te.create_schedule([f.op]) - p = tvm.te.placeholder((1,), "uint64") out_ary = tvm.nd.array(np.ones((1,), "uint64"), ctx) tvm.build(s, [f], target=target)(out_ary) return out_ary.asnumpy()[0] == 0