From 2e4e45772b9b47f8e4b4b169ffb19a64ab341922 Mon Sep 17 00:00:00 2001 From: zilinzhu Date: Thu, 20 May 2021 11:45:46 +0800 Subject: [PATCH 1/3] support generating data of any shape in threefry_generate --- python/tvm/relay/op/random/kernel.py | 3 +-- python/tvm/topi/random/kernel.py | 16 +++++++++++----- tests/python/relay/test_prng.py | 6 ++---- tests/python/topi/python/test_topi_prng.py | 6 ++++++ 4 files changed, 20 insertions(+), 11 deletions(-) diff --git a/python/tvm/relay/op/random/kernel.py b/python/tvm/relay/op/random/kernel.py index fc1248e85678..6c82cc154eb6 100644 --- a/python/tvm/relay/op/random/kernel.py +++ b/python/tvm/relay/op/random/kernel.py @@ -77,8 +77,7 @@ def threefry_generate(key, shape): this function.** shape : Sequence[int] - Desired outputs shape of random numbers. **Currently the total - number of elements must be a multiple of 4.** + Desired outputs shape of random numbers. Returns ------- diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 1be4c86e63c0..6b9abe3f5d79 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -216,7 +216,7 @@ def threefry_generate(gen, out_shape): not be reused in another function, otherwise random numbers will be repeated. out_shape : Sequence[int] - Output shape of the random numbers. Product of all dimensions must be a multiple of 4. + Output shape of the random numbers. Returns ------- @@ -229,9 +229,6 @@ def threefry_generate(gen, out_shape): out_len = tir.const(1) for s in out_shape: out_len *= s - assert ( - out_len.value % 4 == 0 - ), f"Threefry can only generate arrays who's size is a multiple of 4 ({out_len} was provided)." assert ( out_len.value <= 2 ** 64 - 1 ), f"Can only generate up to 2^64 random numbers, but {out_len} were requested." @@ -297,6 +294,12 @@ def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr): # Compute random values _threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4) + with irb.if_scope(out_len % 4 != 0): + remaining = irb.allocate(gen.dtype, 4, name="remaining", scope="global") + tmp[7] = tmp[7] + tir.Cast(gen.dtype, out_len // 4 * 4) # increment counter + _threefry(irb, tmp, 0, tmp, 4, remaining, 0, 1) + with irb.for_range(0, out_len % 4, dtype="uint64", name="i") as i: + out_array[out_len // 4 * 4 + i] = remaining[i] # Update generator state out_gen[0] = tmp[0] # key stays the same @@ -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 out_gen[8] = tmp[8] # path unchanged, so no update here out_gen[9] = tmp[9] diff --git a/tests/python/relay/test_prng.py b/tests/python/relay/test_prng.py index 01515a93546b..be97160778f7 100644 --- a/tests/python/relay/test_prng.py +++ b/tests/python/relay/test_prng.py @@ -137,12 +137,10 @@ def test_threefry_split_infer_fail(): @tvm.testing.requires_llvm -@pytest.mark.xfail(raises=tvm.error.TVMError) -def test_threefry_generate_incorrect_out_size(): +def test_threefry_generate_out_size(): key = tvm.relay.random.threefry_key(1) - # xfail: output size should be multiple of 4 key, rand1 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (5,)), 2) - out1, out2 = tvm.relay.create_executor( + out = tvm.relay.create_executor( "vm", tvm.IRModule.from_expr(tvm.relay.Function([], rand1)), target=tvm.target.Target("llvm"), diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py index 1be32e6ea1d1..db72763ca704 100644 --- a/tests/python/topi/python/test_topi_prng.py +++ b/tests/python/topi/python/test_topi_prng.py @@ -112,6 +112,12 @@ def test_threefry_generate(target, dev): # check that gen out does not equal input assert (a != gen).any(), "Output generator should be different from input generator" + # check that we can generate data whose total number of elements is not a multiple of 4. + a, rands = threefry_generate(target, dev, gen, (7,)) + assert ( + rands.shape[0] == 7 and len(rands.shape) == 2 + ), "Output shape should match requested shape" + # test enough generates to go over generate limit gen = np.array( [0, 0, 0, 0, 0, 0, 0, 2 ** 64 - 2, 1 << 63, 0], dtype="uint64" From 4f8548eae161afd935153a2c38f1eb7d75707bd4 Mon Sep 17 00:00:00 2001 From: zilinzhu Date: Fri, 21 May 2021 09:52:33 +0800 Subject: [PATCH 2/3] update upon reviews --- python/tvm/topi/random/kernel.py | 11 ++++++----- tests/python/relay/test_prng.py | 18 ++++++++++++++++++ tests/python/topi/python/test_topi_prng.py | 2 +- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 6b9abe3f5d79..03ee27003c12 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -293,8 +293,9 @@ def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr): _shift_right(irb, gen[8], gen[9], tmp, 8, tmp, 9) # Compute random values - _threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4) - with irb.if_scope(out_len % 4 != 0): + if out_len.value >= 4: + _threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4) + if out_len.value % 4 != 0: remaining = irb.allocate(gen.dtype, 4, name="remaining", scope="global") tmp[7] = tmp[7] + tir.Cast(gen.dtype, out_len // 4 * 4) # increment counter _threefry(irb, tmp, 0, tmp, 4, remaining, 0, 1) @@ -309,9 +310,9 @@ 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 - with irb.if_scope(out_len % 4 != 0): - out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len % 4) - with irb.else_scope(): + if out_len.value % 4 != 0: + out_gen[7] = tmp[7] + tir.Cast(gen.dtype, 4) # increment counter for the remaining + else: out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len) # increment counter out_gen[8] = tmp[8] # path unchanged, so no update here out_gen[9] = tmp[9] diff --git a/tests/python/relay/test_prng.py b/tests/python/relay/test_prng.py index be97160778f7..f460cc56fb2b 100644 --- a/tests/python/relay/test_prng.py +++ b/tests/python/relay/test_prng.py @@ -79,6 +79,23 @@ def test_threefry_sequential_generate(target, dev): ).any(), "Sequential generates should not have the same output" +@tvm.testing.parametrize_targets +def test_threefry_sequential_generate_remaining(target, dev): + key = tvm.relay.random.threefry_key(1) + key, rand1 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (3,)), 2) + _, rand2 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (3,)), 2) + out1, out2 = tvm.relay.create_executor( + "vm", + tvm.IRModule.from_expr(tvm.relay.Function([], tvm.relay.Tuple((rand1, rand2)))), + target=target, + device=dev, + ).evaluate()() + + assert ( + out1.asnumpy() != out2.asnumpy() + ).any(), "Sequential generates should not have the same output" + + def test_threefry_generate_infer(): oshape = (12,) key_type = tvm.relay.TensorType([10], dtype="uint64") @@ -152,3 +169,4 @@ def test_threefry_generate_out_size(): test_threefry_repeatability(tvm.target.Target("llvm"), tvm.device("cpu")) test_threefry_split(tvm.target.Target("llvm"), tvm.device("cpu")) test_threefry_sequential_generate(tvm.target.Target("llvm"), tvm.device("cpu")) + test_threefry_sequential_generate_remaining(tvm.target.Target("llvm"), tvm.device("cpu")) diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py index db72763ca704..b9ac51419772 100644 --- a/tests/python/topi/python/test_topi_prng.py +++ b/tests/python/topi/python/test_topi_prng.py @@ -115,7 +115,7 @@ def test_threefry_generate(target, dev): # check that we can generate data whose total number of elements is not a multiple of 4. a, rands = threefry_generate(target, dev, gen, (7,)) assert ( - rands.shape[0] == 7 and len(rands.shape) == 2 + rands.shape[0] == 7 and len(rands.shape) == 1 ), "Output shape should match requested shape" # test enough generates to go over generate limit From c0dc06231fcff526949a1c143b83d0797e8035f3 Mon Sep 17 00:00:00 2001 From: zilinzhu Date: Mon, 24 May 2021 10:54:05 +0800 Subject: [PATCH 3/3] update upon review --- python/tvm/topi/random/kernel.py | 7 +++++-- tests/python/relay/test_prng.py | 6 +++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index 03ee27003c12..8b6bb114b181 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -311,7 +311,10 @@ def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr): out_gen[5] = tmp[5] out_gen[6] = tir.const(0, dtype=gen.dtype) # unused, leave it as 0 if out_len.value % 4 != 0: - out_gen[7] = tmp[7] + tir.Cast(gen.dtype, 4) # increment counter for the remaining + # increment counter for the remaining + # as we will generate 4 random numbers for the remaining, increase 4 here. + # the main increment was done before the second _threefry. + out_gen[7] = tmp[7] + tir.Cast(gen.dtype, 4) else: out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len) # increment counter out_gen[8] = tmp[8] # path unchanged, so no update here @@ -497,7 +500,7 @@ def uniform(gen, low, high, out_shape, out_dtype): less than high. out_shape : Sequence[int] - Output shape of the random numbers. Product of all dimensions must be a multiple of 4. + Output shape of the random numbers. out_dtype : str The output dtype. diff --git a/tests/python/relay/test_prng.py b/tests/python/relay/test_prng.py index f460cc56fb2b..f79c79329b67 100644 --- a/tests/python/relay/test_prng.py +++ b/tests/python/relay/test_prng.py @@ -82,8 +82,8 @@ def test_threefry_sequential_generate(target, dev): @tvm.testing.parametrize_targets def test_threefry_sequential_generate_remaining(target, dev): key = tvm.relay.random.threefry_key(1) - key, rand1 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (3,)), 2) - _, rand2 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (3,)), 2) + key, rand1 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (7,)), 2) + _, rand2 = tvm.relay.TupleWrapper(tvm.relay.random.threefry_generate(key, (7,)), 2) out1, out2 = tvm.relay.create_executor( "vm", tvm.IRModule.from_expr(tvm.relay.Function([], tvm.relay.Tuple((rand1, rand2)))), @@ -92,7 +92,7 @@ def test_threefry_sequential_generate_remaining(target, dev): ).evaluate()() assert ( - out1.asnumpy() != out2.asnumpy() + out1.asnumpy()[-3:] != out2.asnumpy()[-3:] ).any(), "Sequential generates should not have the same output"