Repository: cupy/cupy

Issue:
Support batched SVD (`cupy.linalg.svd`)
https://github.com/cupy/cupy/pull/2337#issuecomment-630577091
> It is supposed that now batched cupy.cusolver.gesvdj that uses cusolverDn<t>gesvdjBatched is merged in #3247 and remaining is to fix cupy.linalg.svd to support batched svd using cupy.cusolver layer.


Discussion:
Python Array API will require this capability: https://github.com/data-apis/array-api/pull/114
Looks like #3247 provides the needed support, I wonder why it wasn't included in `cupy.linalg.svd`? @toslunar Do you remember? 
I'd like to support the batched case, and now I suppose it requires a warning in the docstring.

I was uncertain if it's acceptable to use different algorithms depending on shapes.  We will see two types of the difference:

- **sign** (because of the arbitrariness of SVD): It may damage careless users.  For example, recomputing SVDs (for debugging or for memory efficiency) will not always be a good idea.
- **precision**: It shouldn't be a big deal.  Users who want to choose a suitable algorithm for their input can call `cupy.cusolver` functions.

Thanks for reply, @toslunar. I am thinking to call `gesvdjBatched` in `cupy.linalg.svd` when the input has `ndim>=3`, and I am a bit confused from what you said: it does not seem to be possible to set which algorithm to use in `cusolverDn<t>gesvdjBatched()`. Do I miss something?
Sorry.  I meant `gesvd` and `gesvdj` are different in algorithm.  `cupy.linalg.svd` uses `cusolverDn<t>gesvd()` for a 2-dim input, but `cusolverDn<t>gesvdBatched()` does not exist.

Ah OK I understand now, thanks. I think adding a warning in the docstring is sufficient to unblock things. 
> I was uncertain if it's acceptable to use different algorithms depending on shapes.

I think it's acceptable as it's common to change the underlying algorithm for different kinds of inputs, in this case the array dimension. Would you like me to send a PR or you'd prefer to do this yourself @toslunar? 🙂

Resolving patch:
diff --git a/cupy/linalg/_decomposition.py b/cupy/linalg/_decomposition.py
--- a/cupy/linalg/_decomposition.py
+++ b/cupy/linalg/_decomposition.py
@@ -4,8 +4,10 @@
 from cupy_backends.cuda.api import runtime
 from cupy_backends.cuda.libs import cublas
 from cupy_backends.cuda.libs import cusolver
+from cupy.core import internal
 from cupy.cuda import device
 from cupy.cusolver import check_availability
+from cupy.cusolver import _gesvdj_batched, _gesvd_batched
 from cupy.linalg import _util
 
 
@@ -134,7 +136,7 @@ def _potrf_batched(a):
     n = x.shape[-1]
     ldx = x.strides[-2] // x.dtype.itemsize
     handle = device.get_cusolver_handle()
-    batch_size = cupy.core.internal.prod(x.shape[:-2])
+    batch_size = internal.prod(x.shape[:-2])
     dev_info = cupy.empty(batch_size, dtype=numpy.int32)
 
     potrfBatched(
@@ -352,6 +354,64 @@ def qr(a, mode='reduced'):
     return q, _util._triu(r)
 
 
+def _svd_batched(a, a_dtype, full_matrices, compute_uv):
+    batch_shape = a.shape[:-2]
+    batch_size = internal.prod(batch_shape)
+    n, m = a.shape[-2:]
+    s_dtype = a_dtype.lower()
+
+    # first handle any 0-size inputs
+    if batch_size == 0:
+        k = min(m, n)
+        s = cupy.empty(batch_shape + (k,), s_dtype)
+        if compute_uv:
+            if full_matrices:
+                u = cupy.empty(batch_shape + (n, n), dtype=a_dtype)
+                vt = cupy.empty(batch_shape + (m, m), dtype=a_dtype)
+            else:
+                u = cupy.empty(batch_shape + (n, k), dtype=a_dtype)
+                vt = cupy.empty(batch_shape + (k, m), dtype=a_dtype)
+            return u, s, vt
+        else:
+            return s
+    elif m == 0 or n == 0:
+        s = cupy.empty(batch_shape + (0,), s_dtype)
+        if compute_uv:
+            if full_matrices:
+                u = cupy.empty(batch_shape + (n, n), dtype=a_dtype)
+                u[...] = cupy.identity(n, dtype=a_dtype)
+                vt = cupy.empty(batch_shape + (m, m), dtype=a_dtype)
+                vt[...] = cupy.identity(m, dtype=a_dtype)
+            else:
+                u = cupy.empty(batch_shape + (n, 0), dtype=a_dtype)
+                vt = cupy.empty(batch_shape + (0, m), dtype=a_dtype)
+            return u, s, vt
+        else:
+            return s
+
+    # ...then delegate real computation to cuSOLVER
+    a = a.reshape(-1, *(a.shape[-2:]))
+    if runtime.is_hip or (m <= 32 and n <= 32):
+        # copy is done in _gesvdj_batched, so let's try not to do it here
+        a = a.astype(a_dtype, order='C', copy=False)
+        out = _gesvdj_batched(a, full_matrices, compute_uv, False)
+    else:
+        # manually loop over cusolverDn<t>gesvd()
+        # copy (via possible type casting) is done in _gesvd_batched
+        # note: _gesvd_batched returns V, not V^H
+        out = _gesvd_batched(a, a_dtype, full_matrices, compute_uv, False)
+    if compute_uv:
+        u, s, v = out
+        u = u.reshape(*batch_shape, *(u.shape[-2:]))
+        s = s.reshape(*batch_shape, *(s.shape[-1:]))
+        v = v.reshape(*batch_shape, *(v.shape[-2:]))
+        return u, s, v.swapaxes(-2, -1).conj()
+    else:
+        s = out
+        s = s.reshape(*batch_shape, *(s.shape[-1:]))
+        return s
+
+
 def svd(a, full_matrices=True, compute_uv=True):
     """Singular Value Decomposition.
 
@@ -360,11 +420,11 @@ def svd(a, full_matrices=True, compute_uv=True):
     singular values.
 
     Args:
-        a (cupy.ndarray): The input matrix with dimension ``(M, N)``.
+        a (cupy.ndarray): The input matrix with dimension ``(..., M, N)``.
         full_matrices (bool): If True, it returns u and v with dimensions
-            ``(M, M)`` and ``(N, N)``. Otherwise, the dimensions of u and v
-            are respectively ``(M, K)`` and ``(K, N)``, where
-            ``K = min(M, N)``.
+            ``(..., M, M)`` and ``(..., N, N)``. Otherwise, the dimensions
+            of u and v are ``(..., M, K)`` and ``(..., K, N)``, respectively,
+            where ``K = min(M, N)``.
         compute_uv (bool): If ``False``, it only returns singular values.
 
     Returns:
@@ -378,11 +438,17 @@ def svd(a, full_matrices=True, compute_uv=True):
         configuration to a value that is not `ignore` in
         :func:`cupyx.errstate` or :func:`cupyx.seterr`.
 
+    .. note::
+        On CUDA, when ``a.ndim > 2`` and the matrix dimensions <= 32, a fast
+        code path based on Jacobian method (``gesvdj``) is taken. Otherwise,
+        a QR method (``gesvd``) is used.
+
+        On ROCm, there is no such a fast code path that switches the underlying
+        algorithm.
+
     .. seealso:: :func:`numpy.linalg.svd`
     """
-    # TODO(Saito): Current implementation only accepts two-dimensional arrays
     _util._assert_cupy_array(a)
-    _util._assert_rank2(a)
 
     # Cast to float32 or float64
     a_dtype = numpy.promote_types(a.dtype.char, 'f').char
@@ -396,10 +462,11 @@ def svd(a, full_matrices=True, compute_uv=True):
         a_dtype = 'D'
         s_dtype = 'd'
 
+    if a.ndim > 2:
+        return _svd_batched(a, a_dtype, full_matrices, compute_uv)
+
     # Remark 1: gesvd only supports m >= n (WHAT?)
-    # Remark 2: gesvd only supports jobu = 'A' and jobvt = 'A'
-    # Remark 3: gesvd returns matrix U and V^H
-    # Remark 4: Remark 2 is removed since cuda 8.0 (new!)
+    # Remark 2: gesvd returns matrix U and V^H
     n, m = a.shape
 
     if m == 0 or n == 0:
diff --git a/cupy/linalg/_util.py b/cupy/linalg/_util.py
--- a/cupy/linalg/_util.py
+++ b/cupy/linalg/_util.py
@@ -37,12 +37,17 @@ def _check_cusolver_dev_info_if_synchronization_allowed(routine, dev_info):
     if config_linalg == 'ignore':
         return
 
+    try:
+        name = routine.__name__
+    except AttributeError:
+        name = routine  # routine is a str
+
     assert config_linalg == 'raise'
     if (dev_info != 0).any():
         raise linalg.LinAlgError(
             'Error reported by {} in cuSOLVER. devInfo = {}. Please refer'
             ' to the cuSOLVER documentation.'.format(
-                routine.__name__, dev_info))
+                name, dev_info))
 
 
 def _check_cublas_info_array_if_synchronization_allowed(routine, info_array):
diff --git a/cupy_setup_build.py b/cupy_setup_build.py
--- a/cupy_setup_build.py
+++ b/cupy_setup_build.py
@@ -99,6 +99,7 @@
         'file': cuda_files + [
             'cupy.cuda.nvtx',
             'cupy_backends.cuda.libs.cusolver',
+            'cupy.cusolver',
         ],
         'include': [
             'hip/hip_runtime_api.h',
@@ -153,6 +154,7 @@
         'name': 'cusolver',
         'file': [
             'cupy_backends.cuda.libs.cusolver',
+            'cupy.cusolver',
         ],
         'include': [
             'cusolverDn.h',


Test patch:
diff --git a/cupy/testing/helper.py b/cupy/testing/helper.py
--- a/cupy/testing/helper.py
+++ b/cupy/testing/helper.py
@@ -1404,21 +1404,23 @@ def generate_matrix(
     r"""Returns a matrix with specified singular values.
 
     Generates a random matrix with given singular values.
-    This function generates a random NumPy matrix  that
+    This function generates a random NumPy matrix (or a stack of matrices) that
     has specified singular values. It can be used to generate the inputs for a
     test that can be instable when the input value behaves bad.
     Notation: denote the shape of the generated array by :math:`(B..., M, N)`,
     and :math:`K = min\{M, N\}`. :math:`B...` may be an empty sequence.
+
     Args:
         shape (tuple of int): Shape of the generated array, i.e.,
             :math:`(B..., M, N)`.
-        xp(numpy or cupy): Array module to use.
+        xp (numpy or cupy): Array module to use.
         dtype: Dtype of the generated array.
         singular_values (array-like): Singular values of the generated
             matrices. It must be broadcastable to shape :math:`(B..., K)`.
+
     Returns:
-         numpy.ndarray or cupy.ndarray: A random matrix that has specifiec
-         singular values.
+        numpy.ndarray or cupy.ndarray: A random matrix that has specifiec
+        singular values.
     """
 
     if len(shape) <= 1:
diff --git a/tests/cupy_tests/linalg_tests/test_decomposition.py b/tests/cupy_tests/linalg_tests/test_decomposition.py
--- a/tests/cupy_tests/linalg_tests/test_decomposition.py
+++ b/tests/cupy_tests/linalg_tests/test_decomposition.py
@@ -6,6 +6,7 @@
 import cupy
 from cupy import cusolver
 from cupy import testing
+from cupy.core.internal import prod
 from cupy.testing import condition
 import cupyx
 
@@ -156,31 +157,61 @@ def check_usv(self, shape, dtype):
             assert result_gpu[i].dtype == result_cpu[i].dtype
         u_cpu, s_cpu, vh_cpu = result_cpu
         u_gpu, s_gpu, vh_gpu = result_gpu
-        cupy.testing.assert_allclose(s_gpu, s_cpu, atol=1e-4)
+        cupy.testing.assert_allclose(s_gpu, s_cpu, rtol=1e-5, atol=1e-4)
 
-        k, = s_cpu.shape
         # reconstruct the matrix
-        if self.full_matrices:
-            a_gpu_usv = cupy.dot(u_gpu[:, :k] * s_gpu, vh_gpu[:k, :])
+        k = s_cpu.shape[-1]
+        if len(shape) == 2:
+            if self.full_matrices:
+                a_gpu_usv = cupy.dot(u_gpu[:, :k] * s_gpu, vh_gpu[:k, :])
+            else:
+                a_gpu_usv = cupy.dot(u_gpu * s_gpu, vh_gpu)
         else:
-            a_gpu_usv = cupy.dot(u_gpu * s_gpu, vh_gpu)
-        cupy.testing.assert_allclose(a_gpu, a_gpu_usv, atol=1e-4)
+            if self.full_matrices:
+                a_gpu_usv = cupy.matmul(u_gpu[..., :k] * s_gpu[..., None, :],
+                                        vh_gpu[..., :k, :])
+            else:
+                a_gpu_usv = cupy.matmul(u_gpu*s_gpu[..., None, :], vh_gpu)
+        cupy.testing.assert_allclose(a_gpu, a_gpu_usv, rtol=1e-4, atol=1e-4)
 
         # assert unitary
-        cupy.testing.assert_allclose(
-            cupy.matmul(u_gpu.T.conj(), u_gpu),
-            numpy.eye(u_gpu.shape[1]),
-            atol=1e-4)
-        cupy.testing.assert_allclose(
-            cupy.matmul(vh_gpu, vh_gpu.T.conj()),
-            numpy.eye(vh_gpu.shape[0]),
-            atol=1e-4)
+        if len(shape) == 2:
+            cupy.testing.assert_allclose(
+                cupy.matmul(u_gpu.T.conj(), u_gpu),
+                numpy.eye(u_gpu.shape[1]),
+                atol=1e-4)
+            cupy.testing.assert_allclose(
+                cupy.matmul(vh_gpu, vh_gpu.T.conj()),
+                numpy.eye(vh_gpu.shape[0]),
+                atol=1e-4)
+        else:
+            batch = prod(shape[:-2])
+            u_len = u_gpu.shape[-1]
+            vh_len = vh_gpu.shape[-2]
+
+            if batch == 0:
+                id_u_cpu = numpy.empty(shape[:-2] + (u_len, u_len))
+                id_vh_cpu = numpy.empty(shape[:-2] + (vh_len, vh_len))
+            else:
+                id_u_cpu = [numpy.eye(u_len) for _ in range(batch)]
+                id_u_cpu = numpy.stack(id_u_cpu, axis=0).reshape(
+                    *(shape[:-2]), u_len, u_len)
+                id_vh_cpu = [numpy.eye(vh_len) for _ in range(batch)]
+                id_vh_cpu = numpy.stack(id_vh_cpu, axis=0).reshape(
+                    *(shape[:-2]), vh_len, vh_len)
+
+            cupy.testing.assert_allclose(
+                cupy.matmul(u_gpu.swapaxes(-1, -2).conj(), u_gpu),
+                id_u_cpu, atol=1e-4)
+            cupy.testing.assert_allclose(
+                cupy.matmul(vh_gpu, vh_gpu.swapaxes(-1, -2).conj()),
+                id_vh_cpu, atol=1e-4)
 
     @testing.for_dtypes([
         numpy.int32, numpy.int64, numpy.uint32, numpy.uint64,
         numpy.float32, numpy.float64, numpy.complex64, numpy.complex128,
     ])
-    @testing.numpy_cupy_allclose(atol=1e-4)
+    @testing.numpy_cupy_allclose(rtol=1e-5, atol=1e-4)
     def check_singular(self, shape, xp, dtype):
         array = testing.shaped_random(shape, xp, dtype=dtype, seed=self.seed)
         a = xp.asarray(array, dtype=dtype)
@@ -191,36 +222,126 @@ def check_singular(self, shape, xp, dtype):
         assert (a == a_copy).all()
         return result
 
-    def check_rank2(self, array):
-        with pytest.raises(numpy.linalg.LinAlgError):
-            cupy.linalg.svd(array, full_matrices=self.full_matrices)
-
     @condition.repeat(3, 10)
-    def test_svd(self):
+    def test_svd_rank2(self):
         self.check_usv((3, 7))
         self.check_usv((2, 2))
         self.check_usv((7, 3))
 
     @condition.repeat(3, 10)
-    def test_svd_no_uv(self):
+    def test_svd_rank2_no_uv(self):
         self.check_singular((3, 7))
         self.check_singular((2, 2))
         self.check_singular((7, 3))
 
-    @condition.repeat(3, 10)
-    def test_rank2(self):
-        self.check_rank2(cupy.random.randn(2, 3, 4).astype(numpy.float32))
-        self.check_rank2(cupy.random.randn(1, 2, 3, 4).astype(numpy.float64))
-
     @testing.with_requires('numpy>=1.16')
-    def test_empty_array(self):
+    def test_svd_rank2_empty_array(self):
         self.check_usv((0, 3))
         self.check_usv((3, 0))
         self.check_usv((1, 0))
 
     @testing.with_requires('numpy>=1.16')
     @testing.numpy_cupy_array_equal()
-    def test_empty_array_compute_uv_false(self, xp):
+    def test_svd_rank2_empty_array_compute_uv_false(self, xp):
         array = xp.empty((3, 0))
         return xp.linalg.svd(
             array, full_matrices=self.full_matrices, compute_uv=False)
+
+    @condition.repeat(3, 10)
+    def test_svd_rank3(self):
+        self.check_usv((2, 3, 4))
+        self.check_usv((2, 3, 7))
+        self.check_usv((2, 4, 4))
+        self.check_usv((2, 7, 3))
+        self.check_usv((2, 4, 3))
+        self.check_usv((2, 32, 32))  # still use _gesvdj_batched
+
+    @condition.repeat(3, 10)
+    @pytest.mark.skipif(cupy.cuda.runtime.runtimeGetVersion() == 10000,
+                        reason='cuSOLVER bug')
+    def test_svd_rank3_loop(self):
+        # This tests the loop-based batched gesvd on CUDA (_gesvd_batched)
+        self.check_usv((2, 64, 64))
+        self.check_usv((2, 64, 32))
+        self.check_usv((2, 32, 64))
+
+    @condition.repeat(3, 10)
+    def test_svd_rank3_no_uv(self):
+        self.check_singular((2, 3, 4))
+        self.check_singular((2, 3, 7))
+        self.check_singular((2, 4, 4))
+        self.check_singular((2, 7, 3))
+        self.check_singular((2, 4, 3))
+
+    @condition.repeat(3, 10)
+    @pytest.mark.skipif(cupy.cuda.runtime.runtimeGetVersion() == 10000,
+                        reason='cuSOLVER bug')
+    def test_svd_rank3_no_uv_loop(self):
+        # This tests the loop-based batched gesvd on CUDA (_gesvd_batched)
+        self.check_singular((2, 64, 64))
+        self.check_singular((2, 64, 32))
+        self.check_singular((2, 32, 64))
+
+    @testing.with_requires('numpy>=1.16')
+    def test_svd_rank3_empty_array(self):
+        self.check_usv((0, 3, 4))
+        self.check_usv((3, 0, 4))
+        self.check_usv((3, 4, 0))
+        self.check_usv((3, 0, 0))
+        self.check_usv((0, 3, 0))
+        self.check_usv((0, 0, 3))
+
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_svd_rank3_empty_array_compute_uv_false1(self, xp):
+        array = xp.empty((3, 0, 4))
+        return xp.linalg.svd(
+            array, full_matrices=self.full_matrices, compute_uv=False)
+
+    @testing.with_requires('numpy>=1.16')
+    @testing.numpy_cupy_array_equal()
+    def test_svd_rank3_empty_array_compute_uv_false2(self, xp):
+        array = xp.empty((0, 3, 4))
+        return xp.linalg.svd(
+            array, full_matrices=self.full_matrices, compute_uv=False)
+
+    @condition.repeat(3, 10)
+    def test_svd_rank4(self):
+        self.check_usv((2, 2, 3, 4))
+        self.check_usv((2, 2, 3, 7))
+        self.check_usv((2, 2, 4, 4))
+        self.check_usv((2, 2, 7, 3))
+        self.check_usv((2, 2, 4, 3))
+        self.check_usv((2, 2, 32, 32))  # still use _gesvdj_batched
+
+    @condition.repeat(3, 10)
+    @pytest.mark.skipif(cupy.cuda.runtime.runtimeGetVersion() == 10000,
+                        reason='cuSOLVER bug')
+    def test_svd_rank4_loop(self):
+        # This tests the loop-based batched gesvd on CUDA (_gesvd_batched)
+        self.check_usv((3, 2, 64, 64))
+        self.check_usv((3, 2, 64, 32))
+        self.check_usv((3, 2, 32, 64))
+
+    @condition.repeat(3, 10)
+    def test_svd_rank4_no_uv(self):
+        self.check_singular((2, 2, 3, 4))
+        self.check_singular((2, 2, 3, 7))
+        self.check_singular((2, 2, 4, 4))
+        self.check_singular((2, 2, 7, 3))
+        self.check_singular((2, 2, 4, 3))
+
+    @condition.repeat(3, 10)
+    @pytest.mark.skipif(cupy.cuda.runtime.runtimeGetVersion() == 10000,
+                        reason='cuSOLVER bug')
+    def test_svd_rank4_no_uv_loop(self):
+        # This tests the loop-based batched gesvd on CUDA (_gesvd_batched)
+        self.check_singular((3, 2, 64, 64))
+        self.check_singular((3, 2, 64, 32))
+        self.check_singular((3, 2, 32, 64))
+
+    @testing.with_requires('numpy>=1.16')
+    def test_svd_rank4_empty_array(self):
+        self.check_usv((0, 2, 3, 4))
+        self.check_usv((1, 2, 0, 4))
+        self.check_usv((1, 2, 3, 0))
diff --git a/tests/cupy_tests/linalg_tests/test_solve.py b/tests/cupy_tests/linalg_tests/test_solve.py
--- a/tests/cupy_tests/linalg_tests/test_solve.py
+++ b/tests/cupy_tests/linalg_tests/test_solve.py
@@ -163,9 +163,9 @@ def check_x(self, a_shape, rcond, dtype):
         cupy.testing.assert_allclose(result_cpu, result_gpu, atol=1e-3)
         cupy.testing.assert_array_equal(a_gpu_copy, a_gpu)
 
-    def check_shape(self, a_shape, rcond):
+    def check_shape(self, a_shape):
         a = cupy.random.rand(*a_shape)
-        with self.assertRaises(numpy.linalg.LinAlgError):
+        with self.assertRaises(ValueError):
             cupy.linalg.pinv(a)
 
     def test_pinv(self):
@@ -178,10 +178,10 @@ def test_pinv(self):
         self.check_x((5, 3), rcond=0.6)
 
     def test_invalid_shape(self):
-        self.check_shape((2, 3, 4), rcond=1e-15)
-        self.check_shape((2, 3, 4), rcond=0.5)
-        self.check_shape((4, 3, 2, 1), rcond=1e-14)
-        self.check_shape((4, 3, 2, 1), rcond=0.1)
+        # TODO(leofang): NumPy supports batched pinv. Since we now support
+        # batched svd (which is behind pinv), this should not raise
+        self.check_shape((2, 3, 4))
+        self.check_shape((4, 3, 2, 1))
 
 
 @testing.gpu
diff --git a/tests/cupy_tests/testing_tests/test_helper.py b/tests/cupy_tests/testing_tests/test_helper.py
--- a/tests/cupy_tests/testing_tests/test_helper.py
+++ b/tests/cupy_tests/testing_tests/test_helper.py
@@ -505,7 +505,7 @@ def test_shape_mismatch(self):
                 (2, 2), singular_values=numpy.ones(3))
 
     def test_shape_mismatch_2(self):
-        with self.assertRaises(numpy.linalg.LinAlgError):
+        with self.assertRaises(ValueError):
             testing.generate_matrix(
                 (0, 2, 2), singular_values=numpy.ones(3))