Add GPU support with cupy and dask

[mwakok]

I have added support for GPU processing for CUDA-enabled GPUs through cupy. Calling ddm.processing.dmm will automatically determine whether cupy is installed properly.

Installation

In order to use cupy, you will need the following requirement: https://docs.cupy.dev/en/stable/install.html#requirements

I have excluded cupy from the environment.yml as not everyone has a GPU and the cuda-toolkit is ~1Gb. In order to install it, add the cuda option to the pip installation of the package.

pip install -e .[dev,cuda]

I used these references:

• https://blog.dask.org/2019/01/03/dask-array-gpus-first-steps
• https://stackoverflow.com/questions/56776385/how-to-create-a-dask-array-from-cupy-array

TODO
In order to make use of the dask dashboard, you will need to install dask-cuda. However, this package is currently only available for Linux.

[mwakok]

Benchmark

Dataset
Simulated dataset with 1000 frames of 512x512, totalling 512 MB.

CPU
Intel i7-8665U @ 1.9GHz (TU Delft laptop)

• Chunk size = (5,512,512)
• Calling dask.compute on entire tau range
• Hyperthreading enabled (2 threads per core)
• Boost enabled @ 3GHz (not stable during calculation)

Result: Calculating ddm for taus=np.arange(1,100) takes ~270 seconds

GPU
Nvidia Geforce GTX 1070 with 8GB of dedicated GPU memory

Benchmark 1:

• Chunk size = (50,512,512)
• Calling dask.compute on every tau separately

Result: Calculating ddm for taus=np.arange(1,100) takes ~40 seconds
Uses ~1.5Gb of GPU memory

Benchmark 2:

• Chunk size = (50,512,512)
• Calling dask.compute on entire tau range

Result: Calculating ddm for taus=np.arange(1,100) takes ~15 seconds
Uses ~4.5Gb of GPU memory