Repository: apache/tvm

Issue:
[BUG][Threefry] Seg fault when running threefry_generate with some out_shape
Segmentation fault occurs when the `out_shape` of `threefry_generate` setted to some cases.

### To Reproduce
#### environment:
Ubuntu 18.04, llvm version 8.0.1
#### script:
```python
import tvm
import tvm.relay
import tvm.testing
import tvm.topi
import numpy as np


def threefry_generate(target, ctx, gen, size):
    gen_placeholder = tvm.te.placeholder(gen.shape, name="gen", dtype="uint64")
    left_placeholder, right_placeholder = tvm.topi.random.threefry_generate(gen_placeholder, size)
    s = tvm.topi.generic.schedule_extern([left_placeholder, right_placeholder])
    f = tvm.build(s, [gen_placeholder, left_placeholder, right_placeholder])
    out_gen = tvm.nd.array(np.zeros(gen.shape, dtype="uint64"))
    rands = tvm.nd.array(np.zeros(size, dtype="uint64"))
    f(tvm.nd.array(gen), out_gen, rands)
    return out_gen.asnumpy(), rands.asnumpy()


@tvm.testing.parametrize_targets
def test_threefry_generate(target, ctx):
    gen = tvm.relay.random.threefry_key(0).data.asnumpy()
    shape = (1, 3, 224, 224)

    a, rands = threefry_generate(target, ctx, gen, shape)


if __name__ == "__main__":
    test_threefry_generate(tvm.target.Target("llvm"), tvm.context("cpu"))

```
#### error outputs:
```
malloc(): memory corruption
Aborted (core dumped)
```
#### error cases:
`out_shape = (1, 3, 224, 224), (1, 3, 32, 32), (1, 3, 16, 16)`



Discussion:
cc @tkonolige 

Resolving patch:
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -141,7 +141,7 @@ def mix(a, b, rotation):
         return [x, y]
 
     # temporary buffer for holding the results of _PERMUTATIONS
-    tmp = irb.allocate(out_buf.dtype, out_shape, name="tmp", scope="global")
+    tmp = irb.allocate(out_buf.dtype, out_shape * nwords, name="tmp", scope="global")
     tmp_offset = 0
 
     # Initialize entire key. It is composed of the original key with one


Test patch:
diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py
--- a/tests/python/topi/python/test_topi_prng.py
+++ b/tests/python/topi/python/test_topi_prng.py
@@ -87,9 +87,9 @@ def test_threefry_generate(target, ctx):
     gen = tvm.relay.random.threefry_key(0).data.asnumpy()
 
     # check that we can generate some data
-    a, rands = threefry_generate(target, ctx, gen, (100,))
+    a, rands = threefry_generate(target, ctx, gen, (2048,))
     assert (
-        rands.shape[0] == 100 and len(rands.shape) == 1
+        rands.shape[0] == 2048 and len(rands.shape) == 1
     ), "Output shape should match requested shape"
 
     # check that gen out does not equal input
@@ -99,13 +99,13 @@ def test_threefry_generate(target, ctx):
     gen = np.array(
         [0, 0, 0, 0, 0, 0, 0, 2 ** 64 - 2, 1 << 63, 0], dtype="uint64"
     )  # make counter large
-    a, rands = threefry_generate(target, ctx, gen, (100,))
+    a, rands = threefry_generate(target, ctx, gen, (2048,))
     assert gen[4] != a[4], "Overflow of counter should trigger path change"
-    assert a[7] == 100, "Overflow of counter should still update counter"
+    assert a[7] == 2048, "Overflow of counter should still update counter"
 
     # check generate with path at length limit
     gen = np.array([0, 0, 0, 0, 0, 0, 0, 2 ** 64 - 2, 0, 0], dtype="uint64")  # make counter large
-    a, rands = threefry_generate(target, ctx, gen, (100,))
+    a, rands = threefry_generate(target, ctx, gen, (2048,))
     assert (
         gen[0:4] != a[0:4]
     ).any(), "Overflowing counter with no space left in path should change state"