check if any accelerators is in use
@@ -32,6 +32,7 @@ cimport cython # NOQA
from libcpp cimport vector
from libc.stdint cimport int64_t
+from cupy.core cimport _accelerator
from cupy.core cimport _carray
from cupy.core cimport _dtype
from cupy.core._dtype cimport get_dtype
@@ -2760,11 +2761,10 @@ cpdef ndarray matmul(ndarray a, ndarray b, ndarray out=None):
cdef int _cuda_runtime_version = -1
-# TODO(leofang): revive this when ReductionKernel supports CUB block reduction
-# cdef _tensordot_core_mul_sum = ReductionKernel(
-# 'S x, T y', 'U out',
-# 'static_cast<U>(x) * static_cast<U>(y)',
-# 'a + b', 'out = a', '0', '_tensordot_core_mul_sum')
+cdef _tensordot_core_mul_sum = ReductionKernel(
+ 'S x, T y', 'U out',
+ 'static_cast<U>(x) * static_cast<U>(y)',
+ 'a + b', 'out = a', '0', '_tensordot_core_mul_sum')
cpdef ndarray tensordot_core(
@@ -2812,9 +2812,15 @@ cpdef ndarray tensordot_core(
out = _ndarray_init(ret_shape, dtype)
if m == 1 and n == 1:
- # TODO(leofang): switch back to _tensordot_core_mul_sum when
- # ReductionKernel supports CUB block reduction
- out = (a.ravel() * b.ravel()).sum(out=_manipulation._reshape(out, ()))
+ if len(_accelerator._routine_accelerators) > 0:
+ # fast path using CUB or cuTENSOR; if not applicable, it'd fall
+ # back to CuPy's reduction kernel
+ out = (a.ravel() * b.ravel()).sum(
+ out=_manipulation._reshape(out, ()))
+ else:
+ _tensordot_core_mul_sum(
+ a.ravel(), b.ravel(),
+ out=_manipulation._reshape(out, ()))
if out is not ret:
elementwise_copy(out, ret)
return ret
