[Relay][PRNG] Support generating data of any shape in threefry_generate

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

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."
@@ -296,7 +293,14 @@ 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)
+        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)
+            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 +310,13 @@ 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
+        if out_len.value % 4 != 0:
+            # 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
         out_gen[9] = tmp[9]
 
@@ -490,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.

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, (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)))),
+        target=target,
+        device=dev,
+    ).evaluate()()
+
+    assert (
+        out1.asnumpy()[-3:] != out2.asnumpy()[-3:]
+    ).any(), "Sequential generates should not have the same output"
+
+
 def test_threefry_generate_infer():
     oshape = (12,)
     key_type = tvm.relay.TensorType([10], dtype="uint64")
@@ -137,12 +154,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"),
@@ -154,3 +169,4 @@ def test_threefry_generate_incorrect_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"))

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) == 1
+    ), "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"