Repository: cupy/cupy
Pull request: Speed up `cupy.vdot`
Close #3672. Use CUB device reduction to accelerate it. If not available, fall back to the original kernel.

Before: 
```
vdot                :    CPU:   17.272 us   +/- 0.938 (min:   16.300 / max:   39.416) us     GPU-0:  640.646 us   +/-18.809 (min:  630.496 / max:  774.144) us
```
After:
```
vdot                :    CPU:   57.839 us   +/- 6.615 (min:   55.680 / max:  216.692) us     GPU-0:   80.302 us   +/-18.042 (min:   75.776 / max:  232.096) us
```
For some reason CPU overhead is unexpectedly increased...Not sure what I overlooked 😞


Problem statement:
`cupy.vdot` is way slower than naive implementation 
I am on CUDA 10.0 + GTX 2080 Ti + master branch, and this is what I see (with `CUPY_ACCELERATORS=cub`):
```python
>>> import cupy as cp
>>> from cupyx.time import repeat
>>>
>>> a = cp.random.random(1000000)
>>> b = cp.random.random(1000000)
>>> def my_vdot(a, b):
...     return cp.sum(a * b)
... 
>>> print(repeat(cp.vdot, (a, b)))
vdot                :    CPU:   17.272 us   +/- 0.938 (min:   16.300 / max:   39.416) us     GPU-0:  640.646 us   +/-18.809 (min:  630.496 / max:  774.144) us
>>> print(repeat(my_vdot, (a, b)))
my_vdot             :    CPU:   26.312 us   +/- 3.240 (min:   25.204 / max:   97.959) us     GPU-0:   74.100 us   +/-12.585 (min:   70.560 / max:  274.432) us
>>> my_vdot(a, b)  == cp.vdot(a, b)
array(True)
```
I could probably look into this myself, just opening a ticket in case I forget and someone can pick up 😅

Hints from issue discussion:
Without using CUB the naive implementation would win only by a margin:
```python
>>> cp.core._accelerator.set_routine_accelerators([])
>>> cp.core._accelerator.set_reduction_accelerators([])
>>> print(repeat(my_vdot, (a, b)))
my_vdot             :    CPU:   23.317 us   +/- 4.025 (min:   22.159 / max:  223.056) us     GPU-0:  634.641 us   +/-11.096 (min:  626.688 / max:  780.128) us
```
(`cp.vdot` took 640.65 us, see above).
@leofang Your `my_vdot` compute `a + b` into contiguous memory space and then call reduction operation, whereas `cupy.vdot` does not so. Your implementation seems faster in many cases, but sometimes the current implementation is faster in some memory layout.
```py
import cupy
import cupyx.time


tensordot_kernel = cupy.ReductionKernel(
    'S x, T y', 'U out',
    'static_cast<U>(x) * static_cast<U>(y)',
    'a + b', 'out = a', '0',
    'tensordot_kernel')


def tensordot_naive(a, b, out):
    return (a.ravel() * b.ravel()).sum(out=out)


x = cupy.arange(2 ** 20, dtype='float32')
y = cupy.arange(2 ** 20, dtype='float32')
out = cupy.empty((), dtype='float32')

print(cupyx.time.repeat(tensordot_kernel, (x, y, out), max_duration=1))
print(cupyx.time.repeat(tensordot_naive, (x, y, out), max_duration=1))

print(cupyx.time.repeat(tensordot_kernel, (x, x, out), max_duration=1))
print(cupyx.time.repeat(tensordot_naive, (x, x, out), max_duration=1))
```
```
tensordot_kernel    :    CPU:   11.923 us   +/-12.392 (min:   10.815 / max:  437.032) us     GPU-0:  836.028 us   +/-25.847 (min:  820.384 / max: 1331.936) us
tensordot_naive     :    CPU:   29.429 us   +/- 2.777 (min:   28.100 / max:   83.901) us     GPU-0:  576.043 us   +/- 3.707 (min:  567.008 / max:  624.992) us
tensordot_kernel    :    CPU:   11.374 us   +/- 0.452 (min:   10.863 / max:   18.317) us     GPU-0:  438.565 us   +/- 1.325 (min:  429.984 / max:  444.256) us
tensordot_naive     :    CPU:   29.346 us   +/- 5.537 (min:   28.081 / max:  262.398) us     GPU-0:  537.016 us   +/- 5.473 (min:  533.664 / max:  765.152) us
```
The above `tensordot_naive` requires an additional temporary memory space, but I personally agree with the change in #3678. How do you think? @kmaehashi @emcastillo 
Is `CUPY_ACCELERATORS=cub` set in the example above?

I think that resorting to the naive implementation backed by CUB if all the arrays are contiguous is fine, if not we use the current one.
`CUPY_ACCELERATORS=cub` is not set in the above performance numbers. If set, `tensordot_naive` in any case as far as I checked.
```
tensordot_kernel    :    CPU:   12.353 us   +/-12.194 (min:   11.217 / max:  429.384) us     GPU-0:  838.373 us   +/-26.176 (min:  820.512 / max: 1243.168) us
tensordot_naive     :    CPU:   40.067 us   +/- 1.072 (min:   38.911 / max:   65.483) us     GPU-0:   44.603 us   +/- 1.460 (min:   42.528 / max:   80.096) us
tensordot_kernel    :    CPU:   11.715 us   +/- 0.469 (min:   11.194 / max:   19.401) us     GPU-0:  438.515 us   +/- 1.323 (min:  430.464 / max:  446.176) us
tensordot_naive     :    CPU:   39.927 us   +/- 0.948 (min:   38.747 / max:   59.839) us     GPU-0:   44.409 us   +/- 1.052 (min:   42.496 / max:   64.640) us
```
Interesting...I guess we are finding the balance between register pressure and I/O in global memory.

Commit message: check if any accelerators is in use

Patch to review:
diff --git a/cupy/core/core.pyx b/cupy/core/core.pyx
index ba0e8d18b37..fd65744e0c0 100644
--- a/cupy/core/core.pyx
+++ b/cupy/core/core.pyx
@@ -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
@@ -2774,7 +2775,7 @@ cpdef ndarray tensordot_core(
     cdef Py_ssize_t mode, handle
     cdef bint use_sgemmEx
     cdef float one_fp32, zero_fp32
-    ret_dtype = a.dtype.char
+    cdef str ret_dtype = a.dtype.char
     if ret_dtype != b.dtype.char:
         ret_dtype = numpy.promote_types(ret_dtype, b.dtype).char
 
@@ -2811,8 +2812,15 @@ cpdef ndarray tensordot_core(
             out = _ndarray_init(ret_shape, dtype)
 
     if m == 1 and n == 1:
-        _tensordot_core_mul_sum(
-            a.ravel(), b.ravel(), _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


Code review comments:
diff_hunk: @@ -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
path: cupy/core/core.pyx
text: @user1:
The implementation of this PR fallback to CuPy's default `(a.ravel() * b.ravel()).sum(...)`, not reduction kernel.

@author:
Oh, by this I meant `sum` is backed by CuPy's reduction kernel implemented in `cupy/core/_reduction.pyx`, not the `ReductionKernel` object. Would you like me to make it clearer?

@user1:
How about `CuPy's default sum kernel`?

@author:
Sounds good.

diff_hunk: @@ -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:
path: cupy/core/core.pyx
text: @user1:
We may support another accelerator, so could you write as following way?
```py
    for accelerator in _accelerator._routine_accelerators:
        if accelerator == _accelerator.ACCELERATOR_CUB:
            out = ...
            break
        if accelerator == _accelerator.ACCELERATOR_CUTENSOR:
            out = ...
            break
    else
        out = ...
```

@author:
This logic is handled in `sum()`, isn't it?

@user1:
I mean that we may support another accelerator which is not used cupy.sum in the future.

@author:
Oh OK, wasn't aware of this possibility. Will fix. 

(What accelerators you think will be added in the near future, btw? Just curious 😇)

@user1:
We have no plan for another accelerator, currently.