Merge pull request #3467 from leofang/cub_mock

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

cupy/testing/__init__.py

@@ -43,6 +43,7 @@
 from cupy.testing.helper import NumpyAliasBasicTestBase  # NOQA
 from cupy.testing.helper import NumpyAliasValuesTestBase  # NOQA
 from cupy.testing.helper import NumpyError  # NOQA
+from cupy.testing.helper import AssertFunctionIsCalled  # NOQA
 from cupy.testing.helper import shaped_arange  # NOQA
 from cupy.testing.helper import shaped_random  # NOQA
 from cupy.testing.helper import shaped_reverse_arange  # NOQA

cupy/testing/helper.py

@@ -5,6 +5,7 @@
 import random
 import traceback
 import unittest
+from unittest import mock
 import warnings
 
 import numpy
@@ -1257,3 +1258,19 @@ class NumpyAliasValuesTestBase(NumpyAliasTestBase):
 
     def test_values(self):
         assert self.cupy_func(*self.args) == self.numpy_func(*self.args)
+
+
+class AssertFunctionIsCalled:
+
+    def __init__(self, mock_mod, **kwargs):
+        self.patch = mock.patch(mock_mod, **kwargs)
+
+    def __enter__(self):
+        self.handle = self.patch.__enter__()
+        assert self.handle.call_count == 0
+        return self.handle
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        assert self.handle.call_count == 1
+        del self.handle
+        return self.patch.__exit__(exc_type, exc_value, traceback)

tests/cupy_tests/core_tests/test_ndarray_reduction.py

@@ -1,5 +1,7 @@
 import unittest
 
+import numpy
+
 import cupy
 from cupy import testing
 
@@ -228,6 +230,19 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        if len(axis) == len(self.shape):
+            func = 'cupy.core._routines_statistics.cub.device_reduce'
+        else:
+            func = 'cupy.core._routines_statistics.cub.device_segmented_reduce'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.min(axis=axis)
+        # ...then perform the actual computation
         return a.min(axis=axis)
 
     @testing.for_contiguous_axes()
@@ -240,4 +255,17 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        if len(axis) == len(self.shape):
+            func = 'cupy.core._routines_statistics.cub.device_reduce'
+        else:
+            func = 'cupy.core._routines_statistics.cub.device_segmented_reduce'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.max(axis=axis)
+        # ...then perform the actual computation
         return a.max(axis=axis)

tests/cupy_tests/math_tests/test_sumprod.py

@@ -211,6 +211,19 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        if len(axis) == len(self.shape):
+            func = 'cupy.core._routines_math.cub.device_reduce'
+        else:
+            func = 'cupy.core._routines_math.cub.device_segmented_reduce'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.sum(axis=axis)
+        # ...then perform the actual computation
         return a.sum(axis=axis)
 
     @testing.for_contiguous_axes()
@@ -224,11 +237,24 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        if len(axis) == len(self.shape):
+            func = 'cupy.core._routines_math.cub.device_reduce'
+        else:
+            func = 'cupy.core._routines_math.cub.device_segmented_reduce'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.prod(axis=axis)
+        # ...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 +263,21 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        func = 'cupy.core._routines_math.cub.device_scan'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.cumsum()
+        # ...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 +286,21 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        func = 'cupy.core._routines_math.cub.device_scan'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.cumprod()
+        # ...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

@@ -177,6 +177,16 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        func = 'cupy.core._routines_statistics.cub.device_reduce'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.argmin()
+        # ...then perform the actual computation
         return a.argmin()
 
     @testing.for_dtypes('bhilBHILefdFD')
@@ -188,6 +198,16 @@ 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
+        ret = cupy.empty(())  # Cython checks return type, need to fool it
+        func = 'cupy.core._routines_statistics.cub.device_reduce'
+        with testing.AssertFunctionIsCalled(func, return_value=ret):
+            a.argmax()
+        # ...then perform the actual computation
         return a.argmax()
 
 

tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py

@@ -1526,10 +1526,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 +1545,12 @@ 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'
+        with testing.AssertFunctionIsCalled(func):
+            m * x
+        # ...then perform the actual computation
         return m * x