GPU Support

[jakirkham]

Would be useful to support these operations on GPUs as well. One way to do this might be to dispatch using CuPy.

[jakirkham]

Here's what convolution would look like if someone rolled their own today using CuPy 7.0.0 and Dask 2.9.1. There are a few rough edges, but we could improve on this over time both in dask-image and upstream.

import numpy
import numpy as np

import cupy
import cupy as cp

import cupyx
import cupyx.scipy
import cupyx.scipy.ndimage
import cupyx.scipy.ndimage.filters

import dask
import dask.array
import dask.array as da

import dask_image
import dask_image.ndfilters
import dask_image.ndfilters._utils

s = (100, 110)
a = da.from_array(cp.arange(int(np.prod(s)), dtype=cp.float32).reshape(s), chunks=10)
w = cp.ones(a.ndim * (3,), dtype=cp.float32)

origin = dask_image.ndfilters._utils._get_origin(w.shape, 0)
depth = dask_image.ndfilters._utils._get_depth(w.shape, origin)
depth, boundary = dask_image.ndfilters._utils._get_depth_boundary(a.ndim, depth, "none")



r = a.map_overlap(cupyx.scipy.ndimage.filters.convolve,
                  depth=depth,
                  boundary=boundary,
                  dtype=a.dtype,
                  weights=w,
                  mode="reflect",
                  cval=0.0,
                  origin=origin)
r._meta = cp.empty(r.ndim * (0,), dtype=r.dtype)

r.compute()

[GenevieveBuckley]

Linking to this blog post from Matt about mixing dask and cupy arrays, since I found it helpful to read and it's relevant for anyone who comes across this issue: https://matthewrocklin.com/blog/work/2019/01/03/dask-array-gpus-first-steps

[jakirkham]

Also we may be able to use __array_function__ to dispatch over different operations (like convolve here). Some discussion along these lines in issue ( scipy/scipy#10204 ).

[GenevieveBuckley]

I haven't had a chance to look much into cupyimg (BSD-3 licensed), but that could also be useful to look at here: https://github.com/mritools/cupyimg

[jakirkham]

Yeah I’ve seen that. There’s also some work to add equivalent functions to CuPy.

The long term solution would be to use __array_function__ to dispatch. Maybe a short to medium term solution would be to provide our own dispatching mechanism.

[jakirkham]

I've been thinking about implementing some sort of dispatching solution. Maybe that gives us something useful in the near term to medium term. I'll try to put a POC together when I have a chance so we can take a look.

[GenevieveBuckley]

@jakirkham is it worth us having a quick discussion about the still to be decided dispatch mechanism? (Been trying and failing to get my thoughts coherent enough for text for a couple of months now)

I'm interested in putting in some of my own effort into helping things along. I've been doing a lot of reading around NEP 18, 37, etc. but I'm still pretty unclear about what things should look like for us here. Presumably you're a lot clearer on this in your own mind.

[jakirkham]

Yeah I think that is a great idea! Was thinking about suggesting that as well. Am on US Pacific timezone now so should be easier to come up with a time that works. Should we try to schedule something for next week?

[GenevieveBuckley]

Sounds good @jakirkham
I've just sent you an email, so we don't clog up the thread here with conversation about scheduling.

[GenevieveBuckley]

So it seems GPU support for our ndfourier and ndmeasure modules will be significantly more complicated to implement. The extra difficulty here is because these modules don't have the same implementation pattern using map_overlap like we have in ndfilters etc.

One of the main blockers as I understand it is that there isn't currently any support for switching array backends for functions like dask.array.ones(), dask.array.arange() etc. These wrap numpy directly (see here for an example) and get used a lot - both in our code and in dask functions like dask.fft.fftfreq.

Worth noting: there is an open issue at dask/dask#6118 suggesting some problems with dask fourier transforms of cupy arrays might be solved with cupy release 8.0.0 this September. However, this won't fix all our problems because dask-image uses dask.fft.fftfreq which also relies on dask arange, which suffers the problem described in the previous paragraph.

We also use some numpy features that don't have equivalents in cupy - in particular we use the numpy.ndindex iterator for label comprehensions (which all of the ndmeasure module relies on), but there is no cupy.ndindex that exists.

[jakirkham]

Thanks for looking into these 😄

Yeah the general array creation problem is a known issue, which is what Peter has been working on with PR ( numpy/numpy#16935 ). I think we will still need some updates to Dask. Though there have been a few people poking at that problem too.

There may be other ways to work around this shorter term. Happy to think about that as well.

Given what NumPy's ndindex is doing, I don't think we need a CuPy implementation per se. Though feel free to point out the issues you are seeing and would be happy to take a look 🙂