Repository: cupy/cupy
Pull request: Improve cupy.random.randint
Closes https://github.com/cupy/cupy/issues/4120

Problem statement:
cupy.random.randint is slow
Hi all. It seems, `numpy.random.randint` is notably faster than `cupy.random.randint`. Is this expected? 

I've also attained a dramatic speedup, sampling uniform floats with `cupy` and rounding them to `int`. 

```python
import cupy as cp
import numpy as np
import time

def runWithTime(f, iters):
    start = time.time()
    for i in range(iters):
        f()
    print(time.time() - start)

size = 100000
high = 1000000
iters = 50000

runWithTime(lambda: cp.random.uniform(low=0, high=high, size=size).astype('int', copy=False), iters)
runWithTime(lambda: np.random.randint(0, high, size=size), iters)
runWithTime(lambda: cp.random.randint(0, high, size=size), iters) 
```

prints 

```
2.910490036010742
23.180258750915527
36.04059100151062
```

Is there a reason why we cannot rewrite `cupy.random.randint`, relying `cupy.random.uniform`?

* Conditions 
```
CuPy Version          : 6.0.0
CUDA Root             : /usr/local/cuda
CUDA Build Version    : 10000
CUDA Driver Version   : 10010
CUDA Runtime Version  : 10000
cuDNN Build Version   : 7301
cuDNN Version         : 7605
NCCL Build Version    : 1000
NCCL Runtime Version  : (unknown)
Linux 4.4.180-102-default 
Intel(R) Xeon(R) Gold 5118 CPU @ 2.30GHz
Tesla P100
```

Hints from issue discussion:
> Is there a reason why we cannot rewrite cupy.random.randint, relying cupy.random.uniform?

Perhaps it is because casting the output of `xp.random.uniform` to `int` doesn't make the results completely uniform?
As commented it is not appropriate to generate `randint` results from `uniform`, but will keep this issue open as there is obviously a room to improve performance in the current randint implementation.
The current algorithm aligns with numpy.  However the algorithm, rejection sampling, is not very efficient with GPU.
numpy: https://github.com/numpy/numpy/blob/a72b89c7c2e30f5df5cf27f68b6afd45361934fd/numpy/random/src/distributions/distributions.c#L1072-L1081
cupy: https://github.com/cupy/cupy/blob/9a8cdf80b212f31278d7168851b1dac0c6d75f06/cupy/random/_generator.py#L689-L710


@anaruse For a continuous random variable `xi ~ U[0, high)` (where `high` is natural) pdf is `f(x)=1/high` and a chance to hit a unit interval is `P(i <= xi < i+1) = 1/high`. But `int(xi) == i` iff. `i <= xi < i+1` for any natural `i < high`  and hence `P(int(xi) == i) = 1/high`. This proves uniformity of the discrete r. v. `int(xi)` over `0,1,...,high-1`.

OTOH, we are discussing the pseudorandom generators here, so strictly speaking, the numbers `uniform` and `randint` yield are neither independent, nor uniform. Would you be so kind to explain, what you mean by "completely uniform"?

@toslunar thank you for the excellent point. This indeed seems to seriously harm the performance. I've repeated the experiment with `high = 2**20 - 1` (which implies absence of rejections) and  got this
```
2.212979316711426 # uniform+cast
20.105228424072266 # numpy.random.randint
11.889317512512207 # cupy.random.randint
```
Performance is drastically improved (we even beat `numpy`). I think, performance could be improved, if we sampled `size32 / p` for `p` being the chance of rejection in the line 696. That is, I suggest to sample the _expected_ number of integers needed to complete the procedure, or maybe even some more. Was this tradeoff already discussed? 


@akopich  For simplicity, let's assume that the return range of `random.uniform` is [0.5, 1.0) and its data type is float64 (64-bit floating point).  float64 has 52 fraction bits, so the number of possible value in this range is 2^52. Given that this is assigned to N integers, it cannot be assigned strictly equally unless N is a power of 2. This is what I mean by "not complete uniform". Of course, I don't think this is a practical problem unless N is very large, but in a very strict sense, it's not uniform distribution.

@anaruse Could you please clarify a little further? Is your argument based on the fact, that the uniform distribution is modeled by float64-values of equal probability for intervals of equal exponent field, like [0.5, 1)? `randint` seems to be implemented only for int32 and it is possible to distribute 2^52 equiprobable objects into 2^32 bins _almost_  uniformly. And in general (not on [0.5,1)), we have even more (float64-valued) objects, which are not equiprobable. 

Also, I've noticed, `cupy.random.uniform` might return `high`.
```python
import cupy as cp

high = int(2**32) - 1
low = int(2**32) - 3
size = 100000000
print(cp.max(cp.random.uniform(low=low, high=high, size=size)) == high)
```
prints `True`. Is this expected behavior? 

Speaking of "_almost_". How do you test the random generators `cupy` provides?

Commit message: Bug fix and code clean-up

Patch to review:
diff --git a/cupy/random/_generator.py b/cupy/random/_generator.py
index 4481dfb6a10..0a7d5bc0a8d 100644
--- a/cupy/random/_generator.py
+++ b/cupy/random/_generator.py
@@ -672,50 +672,72 @@ def _interval(self, mx, size):
                 'mx must be non-negative (actual: {})'.format(mx))
         elif mx <= _UINT32_MAX:
             dtype = numpy.uint32
+            upper_limit = _UINT32_MAX - (1 << 32) % (mx + 1)
         elif mx <= _UINT64_MAX:
             dtype = numpy.uint64
+            upper_limit = _UINT64_MAX - (1 << 64) % (mx + 1)
         else:
             raise ValueError(
                 'mx must be within uint64 range (actual: {})'.format(mx))
 
-        mask = (1 << mx.bit_length()) - 1
-        mask = cupy.array(mask, dtype=dtype)
-
-        n = functools.reduce(operator.mul, size, 1)
-
-        if n == 0:
+        n_sample = functools.reduce(operator.mul, size, 1)
+        if n_sample == 0:
             return cupy.empty(size, dtype=dtype)
-
-        sample = cupy.empty((n,), dtype=dtype)
-        size32 = sample.view(dtype=numpy.uint32).size
-        n_rem = n  # The number of remaining elements to sample
-        ret = None
-        while n_rem > 0:
-            # Call 32-bit RNG to fill 32-bit or 64-bit `sample`
-            curand.generate(
-                self._generator, sample.data.ptr, size32)
-            # Drop the samples that exceed the upper limit
+        sample = self._curand_generage(n_sample, dtype)
+
+        mx1 = mx + 1
+        if mx1 != (1 << (mx1.bit_length() - 1)):
+            # Get index of samples that exceed the upper limit
+            ng_indices = self._get_indices(sample, upper_limit, False)
+            n_ng = ng_indices.size
+
+            while n_ng > 0:
+                n_supplement = max(n_ng * 2, 1024)
+                supplement = self._curand_generage(n_supplement, dtype)
+
+                # Get index of supplements that are within the upper limit
+                ok_indices = self._get_indices(supplement, upper_limit, True)
+                n_ok = ok_indices.size
+
+                # Replace the values that exceed the upper limit
+                if n_ok >= n_ng:
+                    sample[ng_indices] = supplement[ok_indices[:n_ng]]
+                    n_ng = 0
+                else:
+                    sample[ng_indices[:n_ok]] = supplement[ok_indices]
+                    ng_indices = ng_indices[n_ok:]
+                    n_ng -= n_ok
+            sample %= mx1
+        else:
+            mask = (1 << mx.bit_length()) - 1
             sample &= mask
-            success = sample <= mx
 
-            if ret is None:
-                # If the sampling has finished in the first iteration,
-                # just return the sample.
-                if success.all():
-                    n_rem = 0
-                    ret = sample
-                    break
+        return sample.reshape(size)
 
-                # Allocate the return array.
-                ret = cupy.empty((n,), dtype=dtype)
-
-            n_succ = min(n_rem, int(success.sum()))
-            ret[n - n_rem:n - n_rem + n_succ] = sample[success][:n_succ]
-            n_rem -= n_succ
-
-        assert n_rem == 0
+    def _curand_generage(self, num, dtype):
+        sample = cupy.empty((num,), dtype=dtype)
+        # Call 32-bit RNG to fill 32-bit or 64-bit `sample`
+        size32 = sample.view(dtype=numpy.uint32).size
+        curand.generate(self._generator, sample.data.ptr, size32)
+        return sample
 
-        return ret.reshape(size)
+    def _get_indices(self, sample, upper_limit, cond):
+        dtype = numpy.uint32 if sample.size < 2**32 else numpy.uint64
+        flags = (sample <= upper_limit) if cond else (sample > upper_limit)
+        csum = cupy.cumsum(flags, dtype=dtype)
+        del flags
+        indices = cupy.empty((int(csum[-1]),), dtype=dtype)
+        self._kernel_get_indices(csum, indices, size=csum.size)
+        return indices
+
+    _kernel_get_indices = core.ElementwiseKernel(
+        'raw U csum', 'raw U indices',
+        '''
+        int j = 0;
+        if (i > 0) { j = csum[i-1]; }
+        if (csum[i] > j) { indices[j] = i; }
+        ''',
+        'cupy_get_indices')
 
     def seed(self, seed=None):
         """Resets the state of the random number generator with a seed.


Code review comments:
diff_hunk: @@ -672,50 +672,72 @@ def _interval(self, mx, size):
                 'mx must be non-negative (actual: {})'.format(mx))
         elif mx <= _UINT32_MAX:
             dtype = numpy.uint32
+            upper_limit = _UINT32_MAX - (1 << 32) % (mx + 1)
         elif mx <= _UINT64_MAX:
             dtype = numpy.uint64
+            upper_limit = _UINT64_MAX - (1 << 64) % (mx + 1)
         else:
             raise ValueError(
                 'mx must be within uint64 range (actual: {})'.format(mx))
 
-        mask = (1 << mx.bit_length()) - 1
-        mask = cupy.array(mask, dtype=dtype)
-
-        n = functools.reduce(operator.mul, size, 1)
-
-        if n == 0:
+        n_sample = functools.reduce(operator.mul, size, 1)
+        if n_sample == 0:
             return cupy.empty(size, dtype=dtype)
-
-        sample = cupy.empty((n,), dtype=dtype)
-        size32 = sample.view(dtype=numpy.uint32).size
-        n_rem = n  # The number of remaining elements to sample
-        ret = None
-        while n_rem > 0:
-            # Call 32-bit RNG to fill 32-bit or 64-bit `sample`
-            curand.generate(
-                self._generator, sample.data.ptr, size32)
-            # Drop the samples that exceed the upper limit
+        sample = self._curand_generage(n_sample, dtype)
+
+        mx1 = mx + 1
+        if mx1 != (1 << (mx1.bit_length() - 1)):
+            # Get index of samples that exceed the upper limit
+            ng_indices = self._get_indices(sample, upper_limit, False)
+            n_ng = ng_indices.size
+
+            while n_ng > 0:
+                n_supplement = max(n_ng * 2, 1024)
+                supplement = self._curand_generage(n_supplement, dtype)
path: cupy/random/_generator.py
text: @user1:
typo `generage` -> `generate`

@author:
😅