Add mock tests to ensure cupy.cuda.cub is used

tests/cupy_tests/core_tests/test_ndarray_reduction.py

@@ -1,4 +1,8 @@
 import unittest
+from unittest import mock
+
+import numpy
+import pytest
 
 import cupy
 from cupy import testing
@@ -228,6 +232,25 @@ def test_cub_min(self, xp, dtype, axis):
             a = xp.ascontiguousarray(a)
         elif self.order in ('f', 'F'):
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.min(axis=axis)
+
+        # xp is cupy, first ensure we really use CUB
+        full_reduction = 'cupy.core._routines_statistics.cub.device_reduce'
+        full_raise = NotImplementedError('gotcha_full')
+        segmented_reduction = ('cupy.core._routines_statistics.cub.'
+                               'device_segmented_reduce')
+        segmented_raise = NotImplementedError('gotcha_segment')
+        with mock.patch(full_reduction, side_effect=full_raise), \
+                mock.patch(segmented_reduction, side_effect=segmented_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.min(axis=axis)
+        if len(axis) == len(self.shape):
+            assert str(e.value) == 'gotcha_full'
+        else:
+            assert str(e.value) == 'gotcha_segment'
+        # ...then perform the actual computation
         return a.min(axis=axis)
 
     @testing.for_contiguous_axes()
@@ -240,4 +263,23 @@ def test_cub_max(self, xp, dtype, axis):
             a = xp.ascontiguousarray(a)
         elif self.order in ('f', 'F'):
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.max(axis=axis)
+
+        # xp is cupy, first ensure we really use CUB
+        full_reduction = 'cupy.core._routines_statistics.cub.device_reduce'
+        full_raise = NotImplementedError('gotcha_full')
+        segmented_reduction = ('cupy.core._routines_statistics.cub.'
+                               'device_segmented_reduce')
+        segmented_raise = NotImplementedError('gotcha_segment')
+        with mock.patch(full_reduction, side_effect=full_raise), \
+                mock.patch(segmented_reduction, side_effect=segmented_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.max(axis=axis)
+        if len(axis) == len(self.shape):
+            assert str(e.value) == 'gotcha_full'
+        else:
+            assert str(e.value) == 'gotcha_segment'
+        # ...then perform the actual computation
         return a.max(axis=axis)

tests/cupy_tests/math_tests/test_sumprod.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest import mock
 
 import numpy
 import pytest
@@ -211,6 +212,25 @@ def test_cub_sum(self, xp, dtype, axis):
             a = xp.ascontiguousarray(a)
         elif self.order in ('f', 'F'):
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.sum(axis=axis)
+
+        # xp is cupy, first ensure we really use CUB
+        full_reduction = 'cupy.core._routines_math.cub.device_reduce'
+        full_raise = NotImplementedError('gotcha_full')
+        segmented_reduction = ('cupy.core._routines_math.cub.'
+                               'device_segmented_reduce')
+        segmented_raise = NotImplementedError('gotcha_segment')
+        with mock.patch(full_reduction, side_effect=full_raise), \
+                mock.patch(segmented_reduction, side_effect=segmented_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.sum(axis=axis)
+        if len(axis) == len(self.shape):
+            assert str(e.value) == 'gotcha_full'
+        else:
+            assert str(e.value) == 'gotcha_segment'
+        # ...then perform the actual computation
         return a.sum(axis=axis)
 
     @testing.for_contiguous_axes()
@@ -224,11 +244,30 @@ def test_cub_prod(self, xp, dtype, axis):
             a = xp.ascontiguousarray(a)
         elif self.order in ('f', 'F'):
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.prod(axis=axis)
+
+        # xp is cupy, first ensure we really use CUB
+        full_reduction = 'cupy.core._routines_math.cub.device_reduce'
+        full_raise = NotImplementedError('gotcha_full')
+        segmented_reduction = ('cupy.core._routines_math.cub.'
+                               'device_segmented_reduce')
+        segmented_raise = NotImplementedError('gotcha_segment')
+        with mock.patch(full_reduction, side_effect=full_raise), \
+                mock.patch(segmented_reduction, side_effect=segmented_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.prod(axis=axis)
+        if len(axis) == len(self.shape):
+            assert str(e.value) == 'gotcha_full'
+        else:
+            assert str(e.value) == 'gotcha_segment'
+        # ...then perform the actual computation
         return a.prod(axis=axis)
 
     # TODO(leofang): test axis after support is added
     # don't test float16 as it's not as accurate?
-    @testing.for_dtypes('bhilBHILfdFD')
+    @testing.for_dtypes('bhilBHILfdF')
     @testing.numpy_cupy_allclose(rtol=1E-4)
     def test_cub_cumsum(self, xp, dtype):
         assert cupy.cuda.cub_enabled
@@ -237,11 +276,23 @@ def test_cub_cumsum(self, xp, dtype):
             a = xp.ascontiguousarray(a)
         elif self.order in ('f', 'F'):
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.cumsum()
+
+        # xp is cupy, first ensure we really use CUB
+        full_scan = 'cupy.core._routines_math.cub.device_scan'
+        full_raise = NotImplementedError('gotcha_full')
+        with mock.patch(full_scan, side_effect=full_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.cumsum()
+        assert str(e.value) == 'gotcha_full'
+        # ...then perform the actual computation
         return a.cumsum()
 
     # TODO(leofang): test axis after support is added
     # don't test float16 as it's not as accurate?
-    @testing.for_dtypes('bhilBHILfdFD')
+    @testing.for_dtypes('bhilBHILfdF')
     @testing.numpy_cupy_allclose(rtol=1E-4)
     def test_cub_cumprod(self, xp, dtype):
         assert cupy.cuda.cub_enabled
@@ -250,7 +301,23 @@ def test_cub_cumprod(self, xp, dtype):
             a = xp.ascontiguousarray(a)
         elif self.order in ('f', 'F'):
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            result = a.cumprod()
+            return self._mitigate_cumprod(xp, dtype, result)
+
+        # xp is cupy, first ensure we really use CUB
+        full_scan = 'cupy.core._routines_math.cub.device_scan'
+        full_raise = NotImplementedError('gotcha_full')
+        with mock.patch(full_scan, side_effect=full_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.cumprod()
+        assert str(e.value) == 'gotcha_full'
+        # ...then perform the actual computation
         result = a.cumprod()
+        return self._mitigate_cumprod(xp, dtype, result)
+
+    def _mitigate_cumprod(self, xp, dtype, result):
         # for testing cumprod against complex arrays, the gotcha is CuPy may
         # produce only Inf at the position where NumPy starts to give NaN. So,
         # an error would be raised during assert_allclose where the positions

tests/cupy_tests/sorting_tests/test_search.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest import mock
 
 import numpy
 import pytest
@@ -177,6 +178,18 @@ def test_cub_argmin(self, xp, dtype):
             a = xp.ascontiguousarray(a)
         else:
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.argmin()
+
+        # xp is cupy, first ensure we really use CUB
+        full_scan = 'cupy.core._routines_statistics.cub.device_reduce'
+        full_raise = NotImplementedError('gotcha_full')
+        with mock.patch(full_scan, side_effect=full_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.argmin()
+        assert str(e.value) == 'gotcha_full'
+        # ...then perform the actual computation
         return a.argmin()
 
     @testing.for_dtypes('bhilBHILefdFD')
@@ -188,6 +201,18 @@ def test_cub_argmax(self, xp, dtype):
             a = xp.ascontiguousarray(a)
         else:
             a = xp.asfortranarray(a)
+
+        if xp is numpy:
+            return a.argmax()
+
+        # xp is cupy, first ensure we really use CUB
+        full_scan = 'cupy.core._routines_statistics.cub.device_reduce'
+        full_raise = NotImplementedError('gotcha_full')
+        with mock.patch(full_scan, side_effect=full_raise), \
+                pytest.raises(NotImplementedError) as e:
+            a.argmax()
+        assert str(e.value) == 'gotcha_full'
+        # ...then perform the actual computation
         return a.argmax()
 
 

tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py

@@ -1,5 +1,6 @@
 import pickle
 import unittest
+from unittest import mock
 
 import numpy
 import pytest
@@ -1526,10 +1527,9 @@ def test_getitem_slice_stop_too_large(self, xp, sp):
         return _make(xp, sp, self.dtype)[None:4]
 
 
+# CUB SpMV works only when the matrix size is nonzero
 @testing.parameterize(*testing.product({
-    'make_method': [
-        '_make', '_make_unordered', '_make_empty', '_make_duplicate',
-        '_make_shape'],
+    'make_method': ['_make', '_make_unordered', '_make_duplicate'],
     'dtype': [numpy.float32, numpy.float64, cupy.complex64, cupy.complex128],
 }))
 @testing.with_requires('scipy')
@@ -1546,4 +1546,15 @@ def test_mul_dense_vector(self, xp, sp):
 
         m = self.make(xp, sp, self.dtype)
         x = xp.arange(4).astype(self.dtype)
+        if xp is numpy:
+            return m * x
+
+        # xp is cupy, first ensure we really use CUB
+        func = 'cupyx.scipy.sparse.csr.device_csrmv'
+        side_effect = NotImplementedError('gotcha')
+        with mock.patch(func, side_effect=side_effect), \
+                pytest.raises(NotImplementedError) as e:
+            m * x
+        assert str(e.value) == 'gotcha'
+        # ...then perform the actual computation
         return m * x