Repository: cupy/cupy
Pull request: Use larger type to represent index range in `cupy.take`
Fix #3017 .

`cupy.take()` currently returns unexpected results when the data size of an array to take is the same with `n + 1`, where `n` is the maximum integer the type of indices can represent. For example, `n + 1 = 256` for the indices of type `uint8`.

This PR fixes the issue to use larger type to represent index range in the kernel of `cupy.take`.


Problem statement:
cupy.take() does not function correctly when indices are unsigned integers or negative
```
CuPy Version          : 7.1.1
CUDA Root             : /usr/local/cuda
CUDA Build Version    : 10020
CUDA Driver Version   : 10020
CUDA Runtime Version  : 10020
cuBLAS Version        : 10202
cuFFT Version         : 10102
cuRAND Version        : 10102
cuSOLVER Version      : (10, 3, 0)
cuSPARSE Version      : 10301
NVRTC Version         : (10, 2)
cuDNN Build Version   : 7605
cuDNN Version         : 7605
NCCL Build Version    : 2506
NCCL Runtime Version  : 2506
```

**Code to reproduce**
```
a = cp.ones(256)
indices = cp.array([150], dtype=cp.uint8)
a.take(indices)
```
This should result in an array with the value `1` in it (since the original `a` was only ones) however this gets `0` or other random values instead. Seems to be treated it as the value -106 instead of 150 allowing one to read data outside of the array (also note that apparently negative indices are also allowed and also grab garbage data).

Hints from issue discussion:
Thanks for the reporting. I reproduced the issue and will look into its detail.
I’ve found that the issue is not because the index to take is interpreted as `int8` but the type of `index_range` in the take kernel is specified to `uint8` as well as that of `indeces`, which can not represent 256 that is the size of `a` and, causes taking a value somewhere unexpected.

https://github.com/cupy/cupy/blob/bb3205ceb85996328e0c95a7bec385430b38c454/cupy/core/_routines_indexing.pyx#L366-L368

If we make the size of `a` 255 instead of 256, the following code works fine:

```Python
>>> a = cp.ones(255)
>>> indices = cp.array([150], dtype=cp.uint8)
>>> a.take(indices)
array([1.])
```

Commit message: Use larger type to represent index range in `cupy.take`

Patch to review:
diff --git a/cupy/core/_routines_indexing.pyx b/cupy/core/_routines_indexing.pyx
index efbd8c834af..4bd5e947b8e 100644
--- a/cupy/core/_routines_indexing.pyx
+++ b/cupy/core/_routines_indexing.pyx
@@ -363,7 +363,8 @@ out = a[out_i];
 
 
 _take_kernel = ElementwiseKernel(
-    'raw T a, S indices, uint32 ldim, uint32 cdim, uint32 rdim, S index_range',
+    'raw T a, S indices, uint32 ldim, uint32 cdim, uint32 rdim, '
+    'int64 index_range',
     'T out', _take_kernel_core, 'cupy_take')
 
 
diff --git a/tests/cupy_tests/indexing_tests/test_indexing.py b/tests/cupy_tests/indexing_tests/test_indexing.py
index b36e70b394f..cef63d1571d 100644
--- a/tests/cupy_tests/indexing_tests/test_indexing.py
+++ b/tests/cupy_tests/indexing_tests/test_indexing.py
@@ -1,5 +1,7 @@
 import unittest
 
+import numpy
+
 import cupy
 from cupy import testing
 
@@ -29,6 +31,15 @@ def test_take_no_axis(self, xp):
         b = xp.array([[10, 5], [3, 20]])
         return a.take(b)
 
+    # see cupy#3017
+    @testing.for_dtypes('bBhH')
+    @testing.numpy_cupy_array_equal()
+    def test_take_index_range(self, xp, dtype):
+        iinfo = numpy.iinfo(dtype)
+        a = testing.shaped_arange((iinfo.max + 1,), xp)
+        b = testing.shaped_random((1,), xp, dtype=dtype, scale=iinfo.max)
+        return a.take(b)
+
     @testing.numpy_cupy_array_equal()
     def test_take_along_axis(self, xp):
         a = testing.shaped_random((2, 4, 3), xp, dtype='float32')


Code review comments:
diff_hunk: @@ -29,6 +31,15 @@ def test_take_no_axis(self, xp):
         b = xp.array([[10, 5], [3, 20]])
         return a.take(b)
 
+    # see cupy#3017
+    @testing.for_dtypes('bBhH')
path: tests/cupy_tests/indexing_tests/test_indexing.py
text: @user1:
`for_int_dtypes`?
Char codes should in general be avoided because they depend on environment.

@author:
Here I didn't use `for_int_dtypes` because I thought that it allocated too much memory for larger integer types, but, yes, now I think I can broadcast a tiny array to test this.


@author:
I fixed the test to use `for_int_dtypes` with broadcasting arrays to keep off actual memory allocation. Tests with 64-bit integer types are skipped because they are too large to use for dimensions. Would you take a look at it again?

diff_hunk: @@ -29,6 +31,15 @@ def test_take_no_axis(self, xp):
         b = xp.array([[10, 5], [3, 20]])
         return a.take(b)
 
+    # see cupy#3017
+    @testing.for_dtypes('bBhH')
+    @testing.numpy_cupy_array_equal()
+    def test_take_index_range(self, xp, dtype):
+        iinfo = numpy.iinfo(dtype)
+        a = testing.shaped_arange((iinfo.max + 1,), xp)
+        b = testing.shaped_random((1,), xp, dtype=dtype, scale=iinfo.max)
path: tests/cupy_tests/indexing_tests/test_indexing.py
text: @user1:
Let's use a concrete value instead of `shaped_random`.

@author:
Fixed!