Add support for cupyx sparse to dask.array.dot by anaruse · Pull Request #6846 · dask/dask

anaruse · 2020-11-17T04:54:12Z

This addresses #6820 and makes it available to run dask.array.dot with cupyx sparse (and scipy sparse as well).

import dask.array as da

import cupy as xp
from cupyx.scipy.sparse import csr_matrix
# import numpy as xp
# from scipy.sparse import csr_matrix

matrix_type = csr_matrix
print('matrix_type: {}'.format(matrix_type))

dtype = 'f'
x = xp.arange(24, dtype=dtype).reshape(4, 6)
w = xp.arange(18, dtype=dtype).reshape(6, 3)
ref = xp.dot(x, w)
print('ref:\n{}'.format(ref))

dx = da.from_array(x, chunks=(2, 3), asarray=False, fancy=False)
dx = dx.map_blocks(matrix_type, dtype=dtype)

dw = da.from_array(w, chunks=(3, 1), asarray=False, fancy=False)
dw = dw.map_blocks(matrix_type, dtype=dtype)

ret = da.dot(dx, dw).compute()
print('type(ret): {}'.format(type(ret)))
print('ret:\n{}'.format(ret.todense()))

Click to see the output with current master

matrix_type: <class 'cupyx.scipy.sparse.csr.csr_matrix'>
ref:
[[ 165.  180.  195.]
 [ 435.  486.  537.]
 [ 705.  792.  879.]
 [ 975. 1098. 1221.]]
Traceback (most recent call last):
  File "test-cupyx-sparse.py", line 23, in <module>
    ret = da.dot(dx, dw).compute()
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/base.py", line 224, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/base.py", line 512, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/threaded.py", line 84, in get
    **kwargs
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/local.py", line 486, in get_async
    raise_exception(exc, tb)
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/local.py", line 316, in reraise
    raise exc
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/local.py", line 222, in execute_task
    result = _execute_task(task, data)
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/optimization.py", line 963, in __call__
    return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 151, in get
    result = _execute_task(task, cache)
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 121, in <genexpr>
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 115, in _execute_task
    return [_execute_task(a, cache) for a in arg]
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 115, in <listcomp>
    return [_execute_task(a, cache) for a in arg]
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/utils.py", line 29, in apply
    return func(*args, **kwargs)
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/dask-2.30.0+57.g3e264491-py3.7.egg/dask/array/routines.py", line 229, in _tensordot
    x = tensordot(a, b, axes=axes)
  File "<__array_function__ internals>", line 6, in tensordot
  File "/home/anaruse/.pyenv/versions/miniconda3-4.7.12/lib/python3.7/site-packages/numpy/core/numeric.py", line 1075, in tensordot
    if as_[axes_a[k]] != bs[axes_b[k]]:
IndexError: tuple index out of range

Click to see the output with this PR

matrix_type: <class 'cupyx.scipy.sparse.csr.csr_matrix'>
ref:
[[ 165.  180.  195.]
 [ 435.  486.  537.]
 [ 705.  792.  879.]
 [ 975. 1098. 1221.]]
type(ret): <class 'cupyx.scipy.sparse.coo.coo_matrix'>
ret:
[[ 165.  180.  195.]
 [ 435.  486.  537.]
 [ 705.  792.  879.]
 [ 975. 1098. 1221.]]

jsignell · 2020-11-17T15:42:07Z

Pinging @jakirkham to review if you get a chance.

mrocklin · 2020-11-17T15:52:55Z

Or @pentschev

…

On Tue, Nov 17, 2020 at 7:42 AM Julia Signell ***@***.***> wrote: Pinging @jakirkham <https://github.com/jakirkham> to review if you get a chance. — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#6846 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTAJTSJNBV5UYFEX2UDSQKKWBANCNFSM4TYCFKWQ> .

pentschev

The changes look good to me overall, I've added some suggestions though.

pentschev · 2020-11-17T17:39:02Z

dask/array/backends.py


    concatenate_lookup.register(scipy.sparse.spmatrix, _concatenate)
+
+    def _tensordot(a, b, axes):


It seems like this implementation is exactly the same as the one above, is that correct? I would suggest we have the implementation on a separate function and have both lookups calling that function, this will simplify future maintenance.

That is correct. I will fix it to use the same implementation for CuPy sparse and SciPy sparse.

pentschev · 2020-11-17T17:45:51Z

dask/array/routines.py

-    for a in sorted(axes[0]):
-        ind.insert(a, None)
-    x = x[tuple(ind)]
+    if len(axes[0]) != 1:


Are changes in this file a general Dask issue/limitation or are these specific for CuPy and SciPy? Dask can also use PyData's Sparse, so perhaps adding a test to cover these changes for CuPy in https://github.com/dask/dask/blob/master/dask/array/tests/test_cupy.py and PyData Sparse in https://github.com/dask/dask/blob/master/dask/array/tests/test_sparse.py would be good.

This change affects not only CuPy and SciPy but also others, so I will add tests to check that.

I don't think we need to add a test for PyData Sparse, as there already seems to be a test for tensordot.

dask/dask/array/tests/test_sparse.py

Lines 97 to 111 in bf17b72

def test_tensordot():

x = da.random.random((2, 3, 4), chunks=(1, 2, 2))

x[x < 0.8] = 0

y = da.random.random((4, 3, 2), chunks=(2, 2, 1))

y[y < 0.8] = 0

xx = x.map_blocks(sparse.COO.from_numpy)

yy = y.map_blocks(sparse.COO.from_numpy)

assert_eq(da.tensordot(x, y, axes=(2, 0)), da.tensordot(xx, yy, axes=(2, 0)))

assert_eq(da.tensordot(x, y, axes=(1, 1)), da.tensordot(xx, yy, axes=(1, 1)))

assert_eq(

da.tensordot(x, y, axes=((1, 2), (1, 0))),

da.tensordot(xx, yy, axes=((1, 2), (1, 0))),

)

anaruse · 2020-11-18T05:22:02Z

Thanks for your comment, @pentschev ! I've fixed the PR based on your suggestions, so could you take a look?

pentschev

Changes look good to me, thanks @anaruse ! :)

From my side, this is good to be merged, could someone with merge rights take a final look, maybe @mrocklin ?

martindurant · 2020-11-26T20:45:12Z

Thanks @anaruse !

tomwhite · 2020-11-30T11:07:32Z

Unfortunately this introduces memory issues for general Dask users due to the use of concatenate=True (there's some related discussion in the context of matmul in #6874). Is the change to use concatenate=True necessary for cupyx sparse, or is it an unrelated change?

anaruse · 2020-11-30T11:35:19Z

It is a necessary change for cupyx (and scipy) sparse...
Because cupyx sparse only support 2D arrays, you can't adopt the method of reducing the dimension(s) with sum call after dot , which is being discussed in #6874.

tomwhite · 2020-11-30T11:39:35Z

Thanks for the clarification @anaruse - I think we need some way of avoiding concatenate=True in the non-cupyx case though due to the memory scalability issues it introduces.

ravwojdyla · 2020-11-30T12:42:57Z

We could have a special case for sparse array (based on type) that introduces the contraction in _tensordot only for sparse arrays? But I wonder if there is a more generic solution 🤔

jakirkham · 2020-11-30T18:12:20Z

Could someone please file a new issue with an (ideally not too large) example showing the issue?

tomwhite · 2020-12-01T09:51:37Z

I filed #6916 with an example for this issue.

This is a quick fix for issue #6916 It removes the memory problem for non-sparse array types. Note that tensordot with sparse arrays still needs a proper fix for the memory issues observed after #6846

anaruse added 3 commits November 17, 2020 13:26

Add tensordot_lookup.register in register_scipy_sparse

27bea9e

Use concatenate=True if tensordot has a single reduction axis

f7f7b1b

Add tensordot for CuPy sparse

d66af64

jrbourbeau added the array label Nov 17, 2020

pentschev reviewed Nov 17, 2020

View reviewed changes

anaruse added 2 commits November 18, 2020 11:32

Use same _tensordot implementation for CuPy and SciPy

32e07e5

Add a test for dask.array.dot with CuPy sparse

15eae3d

pentschev approved these changes Nov 18, 2020

View reviewed changes

martindurant merged commit fbccc4e into dask:master Nov 26, 2020

tomwhite mentioned this pull request Dec 1, 2020

Unbounded memory usage in tensordot #6916

Closed

stephenworsley mentioned this pull request Dec 2, 2020

Sparse array approach SciTools/iris-esmf-regrid#23

Closed

ravwojdyla mentioned this pull request Jan 20, 2021

Validate the new matmul works and scales well on cupy and sparse #7050

Open

pentschev mentioned this pull request Feb 15, 2021

dask.array.dot doesn't work with cupyx sparse based dask arrays #6820

Closed

GenevieveBuckley mentioned this pull request Aug 3, 2021

Quick fix for unbounded memory usage in tensordot #7980

Merged

3 tasks

daxiongshu mentioned this pull request Oct 11, 2021

[WIP] fix OOM error of dask-glm with cupy on GPU #8247

Open

3 tasks


		concatenate_lookup.register(scipy.sparse.spmatrix, _concatenate)

		def _tensordot(a, b, axes):

	def test_tensordot():
	x = da.random.random((2, 3, 4), chunks=(1, 2, 2))
	x[x < 0.8] = 0
	y = da.random.random((4, 3, 2), chunks=(2, 2, 1))
	y[y < 0.8] = 0

	xx = x.map_blocks(sparse.COO.from_numpy)
	yy = y.map_blocks(sparse.COO.from_numpy)

	assert_eq(da.tensordot(x, y, axes=(2, 0)), da.tensordot(xx, yy, axes=(2, 0)))
	assert_eq(da.tensordot(x, y, axes=(1, 1)), da.tensordot(xx, yy, axes=(1, 1)))
	assert_eq(
	da.tensordot(x, y, axes=((1, 2), (1, 0))),
	da.tensordot(xx, yy, axes=((1, 2), (1, 0))),
	)

Uh oh!

Conversation

anaruse commented Nov 17, 2020

Uh oh!

jsignell commented Nov 17, 2020

Uh oh!

mrocklin commented Nov 17, 2020 via email

Uh oh!

pentschev left a comment

Choose a reason for hiding this comment

Uh oh!

pentschev Nov 17, 2020

Choose a reason for hiding this comment

Uh oh!

anaruse Nov 18, 2020

Choose a reason for hiding this comment

Uh oh!

pentschev Nov 17, 2020

Choose a reason for hiding this comment

Uh oh!

anaruse Nov 18, 2020

Choose a reason for hiding this comment

Uh oh!

anaruse Nov 18, 2020

Choose a reason for hiding this comment

Uh oh!

anaruse commented Nov 18, 2020

Uh oh!

pentschev left a comment

Choose a reason for hiding this comment

Uh oh!

martindurant commented Nov 26, 2020

Uh oh!

tomwhite commented Nov 30, 2020

Uh oh!

anaruse commented Nov 30, 2020

Uh oh!

tomwhite commented Nov 30, 2020

Uh oh!

ravwojdyla commented Nov 30, 2020

Uh oh!

jakirkham commented Nov 30, 2020

Uh oh!

tomwhite commented Dec 1, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants