Support batched SVD #4628

leofang · 2021-02-05T06:59:45Z

Close #3470. Part of #3062. Another shot based on #3247 ~~to avoid manual looping.~~
UPDATE: See the below discussions in this PR.

Prepare for adopting Array API (data-apis/array-api#114).

Code path divergence:

On CUDA: Call gesvdj if m,n <= 32, else call gesvd in a C loop
On ROCm: Call gesvd as it has no limitation in the input size

TODO:

leofang · 2021-02-05T17:30:38Z

...*gesvdjBatched() only support a stack of matrices with dimension < (32, 32) on either side... @toslunar Looks like we have to do a manual looping anyway for general cases?

leofang · 2021-02-05T17:46:00Z

I should have read carefully the discussion in #2337 first 😅 Looks like I was duplicating the work...

leofang · 2021-02-05T21:36:26Z

...*gesvdjBatched() only support a stack of matrices with dimension < (32, 32) on either side...

I am testing rocsolver_<t>gesvd_batched and it doesn't seem to have this limitation, so I'll try to test & support it in this PR.

Though the result is still incorrect...

leofang · 2021-02-06T04:59:42Z

cupy/cusolver.py

                   u.data.ptr, ldu, v.data.ptr, ldv, params, batch_size)
    work = _cupy.empty(lwork, dtype=a.dtype)
-    info = _cupy.empty(1, dtype=_numpy.int32)
-    solver(handle, jobz, m, n, a.data.ptr, lda, s.data.ptr,
+    info = _cupy.empty(batch_size, dtype=_numpy.int32)


This is a bug: both cuSOLVER and rocSOLVER need the info array to be of the batch size. I thought we fixed it earlier...?

Backported this fix to v8 (#4747).

...but runs almost as fast as a pure Python loop...

avoid slicing overhead, check info in the end, etc

leofang · 2021-02-08T04:23:57Z

On CUDA: the manual looping follows what's done in JAX: https://github.com/google/jax/blob/8bf3f032989caddf05c702acc1ff5353abbe72d2/jaxlib/cusolver.cc#L937-L947
It is made possible here by cythonizing cupy/cusolver.py so that we can speed up the loop using Cython's constructs. ~~But the caveat here is that by calling the <t>gesvd() wrapper we release and acquire the gil in every iteration, which isn't optimal.~~ UPDATE: this is fixed now, see #4628 (comment).

leofang · 2021-02-17T05:45:32Z

Jenkins, test this please

emcastillo · 2021-02-17T05:53:00Z

@anaruse can you please take a look at this? 😇

chainer-ci · 2021-02-17T07:09:01Z

Jenkins CI test (for commit 3515fb9, target branch master) succeeded!

toslunar · 2021-02-17T11:32:17Z

cupy_backends/cupy_lapack.h

+        A += m * n;
+        S += k;
+        U += m * m;
+        VT += n * n;


Is it assumed that (lda, ldu, ldvt) == (m, m, n)?
Keeping the signature of gesvd is nice, but I'd like to see the wrapper is trivially correct.

#4628 (comment)

Yes, when the wrapper is called this is ensured. How about removing lda etc from the signature and just using plain m, n?

How about removing lda etc from the signature and just using plain m, n?

Done in ce32d03.

toslunar · 2021-02-17T11:37:19Z

tests/cupy_tests/linalg_tests/test_decomposition.py

+                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)


I'm fine with calling matmul regardless of shape.

tests/cupy_tests/linalg_tests/test_decomposition.py

toslunar · 2021-02-17T11:57:16Z

cupy/linalg/_decomposition.py

+        # copy (via possible type casting) is done in _gesvd_batched
+        out = _gesvd_batched(a, a_dtype, full_matrices, compute_uv, False)
+    if compute_uv:
+        u, s, v = out


When I read the code first, I missed _gesvd_batched returns v instead of vt. Adding a comment to the docstring of _gesvd_batched seems sufficient for now.

Good catch, @toslunar! There was actually a bug related to this, which I fixed in 62e11ac. Yes, it returns V instead of V^H (transpose conjugate), in order to align with _gesvdj_batched.

cupy/linalg/_decomposition.py

leofang · 2021-02-17T12:51:02Z

Check again

Jenkins, test this please

chainer-ci · 2021-02-17T17:15:35Z

Jenkins CI test (for commit 3515fb9, target branch master) succeeded!

Co-authored-by: Toshiki Kataoka <kataoka@preferred.jp>

leofang · 2021-02-17T20:42:03Z

Jenkins, test this please

chainer-ci · 2021-02-18T00:57:48Z

Jenkins CI test (for commit 5a0376a, target branch master) succeeded!

leofang · 2021-02-18T05:46:06Z

Jenkins, test this please

toslunar

LGTM

chainer-ci · 2021-02-18T09:02:22Z

Jenkins CI test (for commit 5a0376a, target branch master) failed with status FAILURE.

toslunar · 2021-02-18T11:23:02Z

The failed test on Jenkins is known to be flaky (#4673).

leofang · 2021-02-18T13:17:10Z

Thanks, @toslunar!

leofang added 4 commits February 5, 2021 01:57

[WIP] support batched svd (fdFD work)

a419753

[WIP] clean up

a611bca

[WIP] support all dtypes

ce38b72

[WIP] support empty input

6a36d33

leofang added 2 commits February 5, 2021 14:16

[WIP] rocSOLVER

007dc83

[WIP] rocSOLVER

7d83b3e

leofang added 4 commits February 5, 2021 16:37

[WIP] Fix segfault!

d683f58

Though the result is still incorrect...

[WIP] all fp32 tests passed on ROCm!

42ea1fe

[WIP] all dtypes added

ae0535a

style

324dec0

leofang commented Feb 6, 2021

View reviewed changes

leofang added 7 commits February 6, 2021 23:17

[WIP] cythonize cupy/cusolver.py

4b92185

[WIP] this works...

e80e220

...but runs almost as fast as a pure Python loop...

[WIP] slightly reduce CPU overhead via cimport

5e56e39

[WIP] optimizations

4fe6c61

avoid slicing overhead, check info in the end, etc

[WIP] use _ndarray_init

cd3115f

[WIP] all typed

a972ab8

add comments

6fd5ff5

This comment has been minimized.

Sign in to view

toslunar self-assigned this Feb 8, 2021

This comment has been minimized.

Sign in to view

leofang added 4 commits February 8, 2021 01:50

[WIP] release gil before entering the loop

5ee40b0

[WIP] move specialization to C++

bdc83c4

[WIP] clean up and fix

f215947

[WIP] add tests for loop-based gesvd

802b38b

TEST: skip CUDA 10.0

3515fb9

toslunar requested changes Feb 17, 2021

View reviewed changes

leofang and others added 3 commits February 17, 2021 13:53

update wrapper signature; raise on CUDA 10.0

ce32d03

fix a bug; expand tests; add comments

62e11ac

apply review

5a0376a

Co-authored-by: Toshiki Kataoka <kataoka@preferred.jp>

toslunar approved these changes Feb 18, 2021

View reviewed changes

toslunar added this to the v9.0.0b3 milestone Feb 18, 2021

toslunar merged commit bdc6072 into cupy:master Feb 18, 2021

leofang deleted the batched_svd branch February 18, 2021 13:16

This was referenced Feb 18, 2021

Add cudnn and cutensor for CUDA 11.2 in install_library #4680

Merged

Batched SVD test failures #4684

Closed

Revert _get_arch #4682

Merged

emcastillo mentioned this pull request Feb 19, 2021

Revert "Support batched SVD" #4687

Closed

leofang restored the batched_svd branch February 19, 2021 16:24

leofang mentioned this pull request Feb 25, 2021

[backport] Fix _gesvdj_batched info array size #4747

Merged

leofang deleted the batched_svd branch February 25, 2021 03:20

This was referenced Feb 25, 2021

Cupy.linalg.pinv very very Low performance #4435

Open

PTDS Benchmarks rapidsai/dask-cuda#517

Open

IvanYashchuk mentioned this pull request Mar 30, 2021

syevjBatched from cuSOLVER sometimes doesn't work correctly #4847

Open

Uh oh!

Support batched SVD #4628

Support batched SVD #4628

Uh oh!

Conversation

leofang commented Feb 5, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

leofang commented Feb 5, 2021

Uh oh!

leofang commented Feb 5, 2021

Uh oh!

leofang commented Feb 5, 2021

Uh oh!

leofang Feb 6, 2021

Choose a reason for hiding this comment

Uh oh!

leofang Feb 25, 2021

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

This comment has been minimized.

leofang commented Feb 8, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment has been minimized.

leofang commented Feb 17, 2021

Uh oh!

emcastillo commented Feb 17, 2021

Uh oh!

chainer-ci commented Feb 17, 2021

Uh oh!

toslunar Feb 17, 2021

Choose a reason for hiding this comment

Uh oh!

leofang Feb 17, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leofang Feb 17, 2021

Choose a reason for hiding this comment

Uh oh!

toslunar Feb 17, 2021

Choose a reason for hiding this comment

Uh oh!

Uh oh!

toslunar Feb 17, 2021

Choose a reason for hiding this comment

Uh oh!

leofang Feb 17, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

leofang commented Feb 17, 2021

Uh oh!

chainer-ci commented Feb 17, 2021

Uh oh!

leofang commented Feb 17, 2021

Uh oh!

chainer-ci commented Feb 18, 2021

Uh oh!

leofang commented Feb 18, 2021

Uh oh!

toslunar left a comment

Choose a reason for hiding this comment

Uh oh!

chainer-ci commented Feb 18, 2021

Uh oh!

toslunar commented Feb 18, 2021

Uh oh!

leofang commented Feb 18, 2021

Uh oh!

Uh oh!

leofang commented Feb 5, 2021 •

edited

Loading

leofang commented Feb 8, 2021 •

edited

Loading

leofang Feb 17, 2021 •

edited

Loading

leofang Feb 17, 2021 •

edited

Loading