avoid extra func call
@@ -182,7 +182,8 @@ def device_reduce(ndarray x, op, tuple out_axis, out=None,
def device_segmented_reduce(ndarray x, op, tuple reduce_axis,
- tuple out_axis, out=None, bint keepdims=False):
+ tuple out_axis, out=None, bint keepdims=False,
+ Py_ssize_t contiguous_size=0):
# if import at the top level, a segfault would happen when import cupy!
from cupy._creation.ranges import arange
@@ -195,7 +196,6 @@ def device_segmented_reduce(ndarray x, op, tuple reduce_axis,
cdef void* offset_start_ptr
cdef int dtype_id, n_segments, op_code
cdef size_t ws_size
- cdef Py_ssize_t contiguous_size
cdef tuple out_shape
cdef Stream_t s
@@ -213,7 +213,6 @@ def device_segmented_reduce(ndarray x, op, tuple reduce_axis,
raise RuntimeError('input is neither C- nor F- contiguous.')
# prepare input
- contiguous_size = _preprocess_array(x.shape, reduce_axis, out_axis, order)
out_shape = _get_output_shape(x, out_axis, keepdims)
x_ptr = <void*>x.data.ptr
y = ndarray(out_shape, dtype=x.dtype, order=order)
@@ -396,18 +395,18 @@ cdef bint can_use_device_reduce(ndarray x, int op, tuple out_axis, dtype=None):
and x.size <= 0x7fffffff) # until we resolve cupy/cupy#3309
-cdef bint can_use_device_segmented_reduce(
+cdef (bint, Py_ssize_t) can_use_device_segmented_reduce(
ndarray x, int op, tuple reduce_axis, tuple out_axis,
dtype=None, str order='C'):
if not _cub_reduce_dtype_compatible(x.dtype, op, dtype):
- return False
+ return (False, 0)
if not _cub_device_segmented_reduce_axis_compatible(
reduce_axis, x.ndim, order):
- return False
+ return (False, 0)
# until we resolve cupy/cupy#3309
cdef Py_ssize_t contiguous_size = _preprocess_array(
x.shape, reduce_axis, out_axis, order)
- return contiguous_size <= 0x7fffffff
+ return contiguous_size <= 0x7fffffff, contiguous_size
cdef _cub_support_dtype(bint sum_mode, int dev_id):
@@ -462,9 +461,10 @@ cpdef cub_reduction(
"""
# if import at the top level, a segfault would happen when import cupy!
from cupy.core._reduction import _get_axis
- cdef bint enforce_numpy_API = False
+ cdef bint enforce_numpy_API = False, is_ok
cdef str order
cdef tuple reduce_axis, out_axis
+ cdef Py_ssize_t contiguous_size
if op in (CUPY_CUB_ARGMIN, CUPY_CUB_ARGMAX):
# For argmin and argmax, NumPy does not allow a tuple for axis.
@@ -507,10 +507,11 @@ cpdef cub_reduction(
# segmented reduction not currently implemented for argmax, argmin
return None
- if can_use_device_segmented_reduce(arr, op, reduce_axis, out_axis,
- dtype, order):
+ is_ok, contiguous_size = can_use_device_segmented_reduce(
+ arr, op, reduce_axis, out_axis, dtype, order)
+ if is_ok and contiguous_size > 0:
return device_segmented_reduce(arr, op, reduce_axis, out_axis,
- out, keepdims)
+ out, keepdims, contiguous_size)
return None
