def test_cub_max(self, xp, dtype, axis):
a = testing.shaped_random(self.shape, xp, dtype, order=self.order)
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 self.backend == 'device':
func_name = 'cupy._core._routines_statistics.cub.'
if len(axis) == len(self.shape):
func_name += 'device_reduce'
else:
func_name += 'device_segmented_reduce'
with testing.AssertFunctionIsCalled(func_name, return_value=ret):
a.max(axis=axis)
elif self.backend == 'block':
# this is the only function we can mock; the rest is cdef'd
func_name = 'cupy._core._cub_reduction.'
func_name += '_SimpleCubReductionKernel_get_cached_function'
func = _cub_reduction._SimpleCubReductionKernel_get_cached_function
if len(axis) == len(self.shape):
times_called = 2 # two passes
else:
times_called = 1 # one pass
if a.size == 0:
times_called = 0 # _reduction.pyx has an early return path
with testing.AssertFunctionIsCalled(func_name,
wraps=func,
times_called=times_called):
a.max(axis=axis)
elif self.backend == 'fallback':
pass
# ...then perform the actual computation
return a.max(axis=axis)
def test_cub_argmax(self, xp, dtype):
_skip_cuda90(dtype)
a = testing.shaped_random(self.shape, xp, dtype)
if self.order == 'C':
a = xp.ascontiguousarray(a)
else:
a = xp.asfortranarray(a)
if xp is numpy:
return a.argmax(axis=self.axis)
# xp is cupy, first ensure we really use CUB
ret = cupy.empty(()) # Cython checks return type, need to fool it
if self.backend == 'device':
func_name = 'cupy._core._routines_statistics.cub.'
func_name += 'device_reduce'
with testing.AssertFunctionIsCalled(func_name, return_value=ret):
a.argmax(axis=self.axis)
elif self.backend == 'block':
# this is the only function we can mock; the rest is cdef'd
func_name = 'cupy._core._cub_reduction.'
func_name += '_SimpleCubReductionKernel_get_cached_function'
func = _cub_reduction._SimpleCubReductionKernel_get_cached_function
if self.axis is not None and len(self.shape) > 1:
times_called = 1 # one pass
else:
times_called = 2 # two passes
with testing.AssertFunctionIsCalled(func_name,
wraps=func,
times_called=times_called):
a.argmax(axis=self.axis)
# ...then perform the actual computation
return a.argmax(axis=self.axis)
def test_cub_prod(self, xp, dtype, axis):
a = testing.shaped_random(self.shape, xp, dtype)
if self.order in ('c', 'C'):
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 self.backend == 'device':
func_name = 'cupy.core._routines_math.cub.'
if len(axis) == len(self.shape):
func_name += 'device_reduce'
else:
func_name += 'device_segmented_reduce'
with testing.AssertFunctionIsCalled(func_name, return_value=ret):
a.prod(axis=axis)
elif self.backend == 'block':
# this is the only function we can mock; the rest is cdef'd
func_name = 'cupy.core._cub_reduction.'
func_name += '_SimpleCubReductionKernel_get_cached_function'
func = _cub_reduction._SimpleCubReductionKernel_get_cached_function
if len(axis) == len(self.shape):
times_called = 2 # two passes
else:
times_called = 1 # one pass
with testing.AssertFunctionIsCalled(func_name,
wraps=func,
times_called=times_called):
a.prod(axis=axis)
# ...then perform the actual computation
return a.prod(axis=axis)
def test_spy_ndarray(self):
orig = cupy.ndarray
with testing.AssertFunctionIsCalled('cupy.ndarray',
wraps=cupy.ndarray):
a = cupy.ndarray((2, 3), numpy.float32)
assert cupy.ndarray is orig
assert isinstance(a, cupy.ndarray)
def test_fail_called_twice(self):
orig = cupy.ndarray
with pytest.raises(AssertionError):
with testing.AssertFunctionIsCalled('cupy.ndarray'):
cupy.ndarray((2, 3), numpy.float32)
cupy.ndarray((2, 3), numpy.float32)
assert cupy.ndarray is orig
def test_inner_error(self):
orig = cupy.ndarray
with pytest.raises(numpy.AxisError):
with testing.AssertFunctionIsCalled('cupy.ndarray'):
cupy.ndarray((2, 3), numpy.float32)
raise numpy.AxisError('foo')
assert cupy.ndarray is orig
def test_can_use_accelerator_set_unset(self):
# ensure we use CUB block reduction and not CUB device reduction
old_routine_accelerators = _accelerator.get_routine_accelerators()
_accelerator.set_routine_accelerators([])
a = cupy.random.random((10, 10))
# this is the only function we can mock; the rest is cdef'd
func = ''.join(('cupy.core._cub_reduction.',
'_SimpleCubReductionKernel_get_cached_function'))
with testing.AssertFunctionIsCalled(func):
a.sum()
with testing.AssertFunctionIsCalled(func):
a.sum(axis=1)
with testing.AssertFunctionIsCalled(func, times_called=0):
a.sum(axis=0)
_accelerator.set_routine_accelerators(old_routine_accelerators)
def test_histogram(self, xp, dtype):
x = testing.shaped_arange((10, ), xp, dtype)
if xp is numpy:
return xp.histogram(x)
# xp is cupy, first ensure we really use CUB
cub_func = 'cupy._statistics.histogram.cub.device_histogram'
with testing.AssertFunctionIsCalled(cub_func):
xp.histogram(x)
# ...then perform the actual computation
return xp.histogram(x)
def test_optimize2(self):
# Ensure the CUB optimizer is not run when the CUB kernel is not used.
func = 'cupy.core._cub_reduction._get_cub_optimized_params'
times_called = 2 if ('cub' in self.backend) else 0
# Setting "wraps" is necessary to avoid errors being silently ignored.
with testing.AssertFunctionIsCalled(
func,
times_called=times_called,
wraps=cupy.core._cub_reduction._get_cub_optimized_params):
with cupyx.optimizing.optimize():
self.x.sum()
with cupyx.optimizing.optimize():
self.x.sum(axis=1)
with cupyx.optimizing.optimize():
self.x.sum(axis=0) # CUB optimizer not used
def test_optimize1(self):
# Ensure the optimizer is run 3 times for all backends.
func = 'cupyx.optimizing._optimize._optimize'
times_called = 3
# Setting "wraps" is necessary to avoid compilation errors.
with testing.AssertFunctionIsCalled(
func,
times_called=times_called,
wraps=cupyx.optimizing._optimize._optimize):
with cupyx.optimizing.optimize():
self.x.sum()
with cupyx.optimizing.optimize():
self.x.sum(axis=1)
with cupyx.optimizing.optimize():
self.x.sum(axis=0) # CUB falls back to the simple reduction
def test_cub_max(self, xp, dtype, axis):
a = testing.shaped_random(self.shape, xp, dtype, order=self.order)
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)
def test_cub_argmax(self, xp, dtype):
a = testing.shaped_random(self.shape, xp, dtype)
if self.order == 'C':
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()
def test_cutensor_sum(self, xp, dtype, axis):
a = testing.shaped_random(self.shape, xp, dtype)
if self.order in ('c', 'C'):
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 cuTENSOR
ret = cupy.empty(()) # Cython checks return type, need to fool it
func = 'cupy.cutensor._try_reduction_routine'
with testing.AssertFunctionIsCalled(func, return_value=ret):
a.sum(axis=axis)
# ...then perform the actual computation
return a.sum(axis=axis)
def test_cub_cumsum(self, xp, dtype):
a = testing.shaped_random(self.shape, xp, dtype)
if self.order in ('c', 'C'):
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()
def test_cub_prod(self, xp, dtype, axis):
assert cupy.cuda.cub_enabled
a = testing.shaped_random(self.shape, xp, dtype)
if self.order in ('c', 'C'):
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)
def test_cub_cumprod(self, xp, dtype):
if self.backend == 'block':
raise unittest.SkipTest('does not support')
a = testing.shaped_random(self.shape, xp, dtype)
if self.order in ('c', 'C'):
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 test_fail_not_called(self):
orig = cupy.ndarray
with pytest.raises(AssertionError):
with testing.AssertFunctionIsCalled('cupy.ndarray'):
pass
assert cupy.ndarray is orig
def test_times_called(self):
orig = cupy.ndarray
with testing.AssertFunctionIsCalled('cupy.ndarray', times_called=2):
cupy.ndarray((2, 3), numpy.float32)
cupy.ndarray((2, 3), numpy.float32)
assert cupy.ndarray is orig
