Repository: cupy/cupy
Pull request: Support cublasXgetrsBatched and add cupy.cublas.batched_gesv
Closes https://github.com/cupy/cupy/issues/3931

Problem statement:
Performance issue : cupy is 800 times slower than numpy in linalg.solve()
Dear cupy team,

I found that the cupy.linalg.solve routine is much slower than the corresponding numpy routine (about 800 times, which makes is unusable), in a very simple and massively parallel instance. Here I solve 10000 linear systems of size 2x2 in parallel.

```console
import numpy as np
import cupy as cp
import time

n=10000 
a = cp.empty((n,2,2),dtype=np.float32)
v = cp.empty((n,2),dtype=np.float32)

a[:,0,0]=a[:,1,1]=1; a[:,1,0]=a[:,0,1]=0
v[0]=0; v[1]=1

t0=time.time()
np.linalg.solve(a,v)
t1=time.time()
np.linalg.solve(a.get(),v.get())
t2=time.time()

t_cupy,t_numpy=t1-t0,t2-t1
print(f"t_cupy={t_cupy}, t_numpy={t_numpy}, ratio={t_cupy/t_numpy}")
```

Output : 
```console
t_cupy=2.396995782852173, t_numpy=0.0030279159545898438, ratio=791.6322047244095
```

OS : Windows
CuPy version : 7.8
Cuda version : 10.2
Graphics card : Nvidia 1080 GTX

Any chance to solve this performance issue ?

-- Full output of cp.show_config() --
CuPy Version          : 7.8.0
CUDA Root             : None
CUDA Build Version    : 10020
CUDA Driver Version   : 11000
CUDA Runtime Version  : 10020
cuBLAS Version        : 10202
cuFFT Version         : 10102
cuRAND Version        : 10102
cuSOLVER Version      : (10, 3, 0)
cuSPARSE Version      : 10301
NVRTC Version         : (10, 2)
cuDNN Build Version   : 7605
cuDNN Version         : 7605
NCCL Build Version    : None
NCCL Runtime Version  : None
CUB Version           : None
cuTENSOR Version      : None

Hints from issue discussion:
It is not surprising to see such slow performance when solving multiple small linear matrix equations, because the current `cupy.linalg.solve` implementation solves the multiple linear matrix equations one by one.
https://github.com/cupy/cupy/blob/e2b5a3c19f8aec1c4acfda239265ba4e54e5bd60/cupy/linalg/solve.py#L65-L68

Probably, your use case can be significantly improved by using [cublasXgetrfBatched()](https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getrfbatched) and [cublasXgettrsBatched(](https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getrsbatched)).
Ok, thanks for the workaround.
It would be nice feature if cupy handled it directly though.

> Le 3 sept. 2020 à 11:09, Akira Naruse <notifications@github.com> a écrit :
> 
> 
> It is not surprising to see such slow performance when solving multiple small linear matrix equations, because the current cupy.linalg.solve implementation solves the multiple linear matrix equations one by one.
> https://github.com/cupy/cupy/blob/e2b5a3c19f8aec1c4acfda239265ba4e54e5bd60/cupy/linalg/solve.py#L65-L68 <https://github.com/cupy/cupy/blob/e2b5a3c19f8aec1c4acfda239265ba4e54e5bd60/cupy/linalg/solve.py#L65-L68>
> Probably, your use case can be significantly improved by using cublasXgetrfBatched() <https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getrfbatched> and cublasXgettrsBatched( <https://docs.nvidia.com/cuda/cublas/index.html#cublas-lt-t-gt-getrsbatched>).
> 
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub <https://github.com/cupy/cupy/issues/3931#issuecomment-686359901>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAYOZPTOVY3K5XHK27KPKYTSD5MN7ANCNFSM4QUNJZ4Q>.
>

Commit message: Add error checks

Patch to review:
diff --git a/cupy/cublas.py b/cupy/cublas.py
new file mode 100644
index 00000000000..04b5cc58ad2
--- /dev/null
+++ b/cupy/cublas.py
@@ -0,0 +1,78 @@
+import numpy
+
+import cupy
+from cupy_backends.cuda.libs import cublas
+from cupy.cuda import device
+from cupy.linalg import util
+
+
+def batched_gesv(a, b):
+    """Solves multiple linear matrix equations using cublas<t>getr[fs]Batched().
+
+    Computes the solution to system of linear equation ``ax = b``.
+
+    Args:
+        a (cupy.ndarray): The matrix with dimension ``(..., M, M)``.
+        b (cupy.ndarray): The matrix with dimension ``(..., M)`` or
+            ``(..., M, K)``.
+
+    Returns:
+        cupy.ndarray:
+            The matrix with dimension ``(..., M)`` or ``(..., M, K)``.
+    """
+    util._assert_cupy_array(a, b)
+    util._assert_nd_squareness(a)
+
+    if not ((a.ndim == b.ndim or a.ndim == b.ndim + 1) and
+            a.shape[:-1] == b.shape[:a.ndim - 1]):
+        raise ValueError(
+            'a must have (..., M, M) shape and b must have (..., M) '
+            'or (..., M, K)')
+
+    dtype = numpy.promote_types(a.dtype.char, 'f')
+    if dtype == 'f':
+        t = 's'
+    elif dtype == 'd':
+        t = 'd'
+    elif dtype == 'F':
+        t = 'c'
+    elif dtype == 'D':
+        t = 'z'
+    else:
+        raise TypeError('invalid dtype')
+    getrf = getattr(cublas, t + 'getrfBatched')
+    getrs = getattr(cublas, t + 'getrsBatched')
+
+    bs = numpy.prod(a.shape[:-2]) if a.ndim > 2 else 1
+    n = a.shape[-1]
+    nrhs = b.shape[-1] if a.ndim == b.ndim else 1
+    b_shape = b.shape
+    a_data_ptr = a.data.ptr
+    b_data_ptr = b.data.ptr
+    a = cupy.ascontiguousarray(a.reshape(bs, n, n).transpose(0, 2, 1),
+                               dtype=dtype)
+    b = cupy.ascontiguousarray(b.reshape(bs, n, nrhs).transpose(0, 2, 1),
+                               dtype=dtype)
+    if a.data.ptr == a_data_ptr:
+        a = a.copy()
+    if b.data.ptr == b_data_ptr:
+        b = b.copy()
+
+    handle = device.get_cublas_handle()
+    lda = n
+    a_step = lda * n * a.itemsize
+    a_array = cupy.arange(a.data.ptr, a.data.ptr + a_step * bs, a_step,
+                          dtype=cupy.uintp)
+    ldb = n
+    b_step = ldb * nrhs * b.itemsize
+    b_array = cupy.arange(b.data.ptr, b.data.ptr + b_step * bs, b_step,
+                          dtype=cupy.uintp)
+    pivot = cupy.empty((bs, n), dtype=numpy.int32)
+    info = cupy.empty((bs,), dtype=numpy.int32)
+    # LU factorization (A = L * U)
+    getrf(handle, n, a_array.data.ptr, lda, pivot.data.ptr, info.data.ptr, bs)
+    util._check_cublas_info_array_if_synchronization_allowed(getrf, info)
+    # Solves Ax = b
+    getrs(handle, cublas.CUBLAS_OP_N, n, nrhs, a_array.data.ptr, lda,
+          pivot.data.ptr, b_array.data.ptr, ldb, bs)
+    return b.transpose(0, 2, 1).reshape(b_shape)
diff --git a/cupy/linalg/solve.py b/cupy/linalg/solve.py
index 78f996bda3c..52ec3bcc93a 100644
--- a/cupy/linalg/solve.py
+++ b/cupy/linalg/solve.py
@@ -8,6 +8,7 @@
 from cupy.cuda import device
 from cupy.linalg import decomposition
 from cupy.linalg import util
+from cupy.cublas import batched_gesv
 
 
 def solve(a, b):
@@ -54,6 +55,9 @@ def solve(a, b):
     if a.ndim == 2:
         return cupy.cusolver.gesv(a, b)
 
+    if a.shape[-1] <= 256:
+        return batched_gesv(a, b)
+
     x = cupy.empty_like(b)
     shape = a.shape[:-2]
     for i in range(numpy.prod(shape)):
diff --git a/cupy_backends/cuda/cupy_cuda.h b/cupy_backends/cuda/cupy_cuda.h
index 4928baae37a..807ddfa8cd2 100644
--- a/cupy_backends/cuda/cupy_cuda.h
+++ b/cupy_backends/cuda/cupy_cuda.h
@@ -782,6 +782,22 @@ cublasStatus_t cublasZgetrfBatched(...) {
     return CUBLAS_STATUS_SUCCESS;
 }
 
+cublasStatus_t cublasSgetrsBatched(...) {
+    return CUBLAS_STATUS_SUCCESS;
+}
+
+cublasStatus_t cublasDgetrsBatched(...) {
+    return CUBLAS_STATUS_SUCCESS;
+}
+
+cublasStatus_t cublasCgetrsBatched(...) {
+    return CUBLAS_STATUS_SUCCESS;
+}
+
+cublasStatus_t cublasZgetrsBatched(...) {
+    return CUBLAS_STATUS_SUCCESS;
+}
+
 cublasStatus_t cublasSgetriBatched(...) {
     return CUBLAS_STATUS_SUCCESS;
 }
diff --git a/cupy_backends/cuda/libs/cublas.pxd b/cupy_backends/cuda/libs/cublas.pxd
index fba04526e30..2d16affb1be 100644
--- a/cupy_backends/cuda/libs/cublas.pxd
+++ b/cupy_backends/cuda/libs/cublas.pxd
@@ -249,6 +249,19 @@ cpdef cgetrfBatched(intptr_t handle, int n, size_t Aarray, int lda,
 cpdef zgetrfBatched(intptr_t handle, int n, size_t Aarray, int lda,
                     size_t PivotArray, size_t infoArray, int batchSize)
 
+cpdef int sgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize)
+cpdef int dgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize)
+cpdef int cgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize)
+cpdef int zgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize)
+
 cpdef sgetriBatched(intptr_t handle, int n, size_t Aarray, int lda,
                     size_t PivotArray, size_t Carray, int ldc,
                     size_t infoArray, int batchSize)
diff --git a/cupy_backends/cuda/libs/cublas.pyx b/cupy_backends/cuda/libs/cublas.pyx
index b64e880f391..c26cd15025e 100644
--- a/cupy_backends/cuda/libs/cublas.pyx
+++ b/cupy_backends/cuda/libs/cublas.pyx
@@ -221,6 +221,24 @@ cdef extern from '../cupy_cuda.h' nogil:
     int cublasZgetrfBatched(
         Handle handle, int n, cuDoubleComplex **Aarray, int lda,
         int *PivotArray, int *infoArray, int batchSize)
+
+    int cublasSgetrsBatched(
+        Handle handle, Operation trans, int n, int nrhs,
+        const float **Aarray, int lda, const int *devIpiv,
+        float **Barray, int ldb, int *info, int batchSize)
+    int cublasDgetrsBatched(
+        Handle handle, Operation trans, int n, int nrhs,
+        const double **Aarray, int lda, const int *devIpiv,
+        double **Barray, int ldb, int *info, int batchSize)
+    int cublasCgetrsBatched(
+        Handle handle, Operation trans, int n, int nrhs,
+        const cuComplex **Aarray, int lda, const int *devIpiv,
+        cuComplex **Barray, int ldb, int *info, int batchSize)
+    int cublasZgetrsBatched(
+        Handle handle, Operation trans, int n, int nrhs,
+        const cuDoubleComplex **Aarray, int lda, const int *devIpiv,
+        cuDoubleComplex **Barray, int ldb, int *info, int batchSize)
+
     int cublasSgetriBatched(
         Handle handle, int n, const float **Aarray, int lda,
         int *PivotArray, float *Carray[], int ldc, int *infoArray,
@@ -1004,6 +1022,59 @@ cpdef zgetrfBatched(intptr_t handle, int n, size_t Aarray, int lda,
     check_status(status)
 
 
+cpdef int sgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize):
+    cdef int info
+    _setStream(handle)
+    with nogil:
+        status = cublasSgetrsBatched(
+            <Handle>handle, <Operation>trans, n, nrhs,
+            <const float**>Aarray, lda, <const int*>devIpiv,
+            <float**>Barray, ldb, &info, batchSize)
+    check_status(status)
+    return info
+
+cpdef int dgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize):
+    cdef int info
+    _setStream(handle)
+    with nogil:
+        status = cublasDgetrsBatched(
+            <Handle>handle, <Operation>trans, n, nrhs,
+            <const double**>Aarray, lda, <const int*>devIpiv,
+            <double**>Barray, ldb, &info, batchSize)
+    check_status(status)
+    return info
+
+cpdef int cgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize):
+    cdef int info
+    _setStream(handle)
+    with nogil:
+        status = cublasCgetrsBatched(
+            <Handle>handle, <Operation>trans, n, nrhs,
+            <const cuComplex**>Aarray, lda, <const int*>devIpiv,
+            <cuComplex**>Barray, ldb, &info, batchSize)
+    check_status(status)
+    return info
+
+cpdef int zgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize):
+    cdef int info
+    _setStream(handle)
+    with nogil:
+        status = cublasZgetrsBatched(
+            <Handle>handle, <Operation>trans, n, nrhs,
+            <const cuDoubleComplex**>Aarray, lda, <const int*>devIpiv,
+            <cuDoubleComplex**>Barray, ldb, &info, batchSize)
+    check_status(status)
+    return info
+
+
 cpdef sgetriBatched(
         intptr_t handle, int n, size_t Aarray, int lda, size_t PivotArray,
         size_t Carray, int ldc, size_t infoArray, int batchSize):
diff --git a/tests/cupy_tests/test_cublas.py b/tests/cupy_tests/test_cublas.py
new file mode 100644
index 00000000000..1f58c95db25
--- /dev/null
+++ b/tests/cupy_tests/test_cublas.py
@@ -0,0 +1,67 @@
+import unittest
+
+import numpy
+
+import cupy
+from cupy import cublas
+from cupy import testing
+from cupy.testing import attr
+
+
+@testing.parameterize(*testing.product({
+    'dtype': ['float32', 'float64', 'complex64', 'complex128'],
+    'n': [10, 33, 100],
+    'bs': [None, 1, 10],
+    'nrhs': [None, 1, 10],
+}))
+@attr.gpu
+class TestBatchedGesv(unittest.TestCase):
+    _tol = {'f': 1e-5, 'd': 1e-12}
+
+    def _make_array(self, shape, alpha, beta):
+        a = testing.shaped_random(shape, cupy, dtype=self.r_dtype,
+                                  scale=alpha) + beta
+        return a
+
+    def _make_matrix(self, shape, alpha, beta):
+        a = self._make_array(shape, alpha, beta)
+        if self.dtype.char in 'FD':
+            a = a + 1j * self._make_array(shape, alpha, beta)
+        return a
+
+    def setUp(self):
+        self.dtype = numpy.dtype(self.dtype)
+        if self.dtype.char in 'fF':
+            self.r_dtype = numpy.float32
+        else:
+            self.r_dtype = numpy.float64
+        n = self.n
+        bs = 1 if self.bs is None else self.bs
+        nrhs = 1 if self.nrhs is None else self.nrhs
+        # Diagonally dominant matrix is used as it is stable
+        alpha = 2.0 / n
+        a = self._make_matrix((bs, n, n), alpha, -alpha / 2)
+        diag = cupy.diag(cupy.ones((n,), dtype=self.r_dtype))
+        for i in range(bs):
+            a[i][diag > 0] = 0
+            a[i] += diag
+        x = self._make_matrix((bs, n, nrhs), 0.2, 0.9)
+        b = cupy.matmul(a, x)
+        a_shape = (n, n) if self.bs is None else (bs, n, n)
+        b_shape = [n]
+        if self.bs is not None:
+            b_shape.insert(0, bs)
+        if self.nrhs is not None:
+            b_shape.append(nrhs)
+        self.a = a.reshape(a_shape)
+        self.b = b.reshape(b_shape)
+        self.x_ref = x.reshape(b_shape)
+        if self.r_dtype == numpy.float32:
+            self.tol = self._tol['f']
+        elif self.r_dtype == numpy.float64:
+            self.tol = self._tol['d']
+
+    def test_batched_gesv(self):
+        x = cublas.batched_gesv(self.a, self.b)
+        cupy.testing.assert_allclose(x, self.x_ref,
+                                     rtol=self.tol, atol=self.tol)


Code review comments:
diff_hunk: @@ -54,6 +55,9 @@ def solve(a, b):
     if a.ndim == 2:
         return cupy.cusolver.gesv(a, b)
 
+    if a.shape[-1] <= 256:
path: cupy/linalg/solve.py
text: @user1:
Is this cuBLAS's requirement? Could you add a comment here?

@user1:
Also, if you wanna hijack late (as is done here), you don't have to call `util._assert_cupy_array()` etc in `batched_gesv` and do the array checks twice. Another option is to hijack sooner and leave `batched_gesv` as is.   

@user1:
One more question: Is it possible that all of the tests will now go through `batched_gesv` (maybe because the test size is small) and make the main path below untested?

@author:
> Is this cuBLAS's requirement? Could you add a comment here?

This is not a cuBLAS's requirement, but there is a low performance issue in batched getrf routines if matrices are not small..

@author:
> One more question: Is it possible that all of the tests will now go through batched_gesv (maybe because the test size is small) and make the main path below untested?

Good point! Maybe, we should add tests with large matrices, like a shape of (2, 512, 512) ?

@author:
I added a routine that allows you to change the threshold to determine whether or not to use batched_gesv.

diff_hunk: @@ -1004,6 +1022,59 @@ cpdef zgetrfBatched(intptr_t handle, int n, size_t Aarray, int lda,
     check_status(status)
 
 
+cpdef int sgetrsBatched(intptr_t handle, int trans, int n, int nrhs,
+                        size_t Aarray, int lda, size_t devIpiv,
+                        size_t Barray, int ldb, int batchSize):
+    cdef int info
+    _setStream(handle)
+    with nogil:
+        status = cublasSgetrsBatched(
+            <Handle>handle, <Operation>trans, n, nrhs,
+            <const float**>Aarray, lda, <const int*>devIpiv,
+            <float**>Barray, ldb, &info, batchSize)
+    check_status(status)
+    return info
path: cupy_backends/cuda/libs/cublas.pyx
text: @user1:
I think it's better to follow other cuBLAS routines (like `sgetrfBatched` etc) and pass in the info array instead of returning an int here. This return value is not used by your `batched_gesv` anyway.

@author:
Actually I was wondering which was better. My first implementation followed other routine's convention, but changed it to current design as info is a host array with size 1 and it is easier to handle it with current design.
Anyway, I agree that current design is not consistent with other routines. So, I will change it to follow other routines.

@user1:
Oh oh I see what you're saying! In most routines `infoArray` is a device array (of size `batch_size`), but in this case it's a host array of size 1! It's very inconsistent API design...😞 

Another inconsistent design is that the batched version lives in cuBLAS but the non-batched version lives in cuSOLVER !@#$%...Otherwise we could just merge `cupy.cublas.batched_gesv` with `cupy.cusolver.gesv`...(Or should we do that anyways?)

@author:
There's a historical reason for that: when we first added the Lapack-like routines to the CUDA library, there was no cuSOLVER yet..