How to use Dask to run python code on the GPU?

Question

I have some code that uses Numba cuda.jit in order for me to run on the gpu, and I would like to layer dask on top of it if possible.

Example Code

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from numba import cuda, njit
import numpy as np
from dask.distributed import Client, LocalCluster


@cuda.jit()
def addingNumbersCUDA (big_array, big_array2, save_array):
    i = cuda.grid(1)
    if i < big_array.shape[0]:
        for j in range (big_array.shape[1]):
            save_array[i][j] = big_array[i][j] * big_array2[i][j]


if __name__ == "__main__":
    cluster = LocalCluster()
    client = Client(cluster)

    big_array = np.random.random_sample((100, 3000))
    big_array2  = np.random.random_sample((100, 3000))
    save_array = np.zeros(shape=(100, 3000))

    arraysize = 100
    threadsperblock = 64
    blockspergrid = (arraysize + (threadsperblock - 1))

    d_big_array = cuda.to_device(big_array)
    d_big_array2 = cuda.to_device(big_array2)
    d_save_array = cuda.to_device(save_array)

    addingNumbersCUDA[blockspergrid, threadsperblock](d_big_array, d_big_array2, d_save_array)

    save_array = d_save_array.copy_to_host()

If my function addingNumbersCUDA didn't use any CUDA I would just put client.submit in front of my function (along with gather after) and it would work. But, since I'm using CUDA putting submit in front of the function doesn't work. The dask documentation says that you can target the gpu, but it's unclear as to how to actually set it up in practice. How would I set up my function to use dask with the gpu targeted and with cuda.jit if possible?

score 2 · Accepted Answer · answered May 18 '19 at 14:37

2

You may want to look through Dask's documentation on GPUs

But, since I'm using CUDA putting submit in front of the function doesn't work.

There is no particular reason why this should be the case. All Dask does it run your function on a different computer. It doesn't change or modify your function in any way.

answered May 18 '19 at 14:37

MRocklin

55,641
23
163
235

I've set up my dask worker for the gpu and my function as "x = client.submit(addingNumbersCUDA[blockspergrid, threadsperblock], big_array, big_array2, save_array, resources={'GPU': 1})", but when I gather it "none" is returned. Do you know why this is? – Bryce Booze May 20 '19 at 16:32
Your function doesn't have a return value. It looks like you're storing to an out parameter as a side effect. I recommend building a function that retuns a value. – MRocklin May 21 '19 at 13:09
I can't specifically return a value from a numba cuda function. "kernels cannot explicitly return a value; all result data must be written to an array passed to the function." Is it possible to obtain the results stored in my array from dask? – Bryce Booze May 21 '19 at 15:41
You might make another function around your Numba function that handles the mutation and in-place behavior. Dask doesn't support in-place operation. Typically your new function would allocate an output array, call your in-place operating numba function, and then return the output array normally. – MRocklin May 23 '19 at 14:54

How to use Dask to run python code on the GPU?

1 Answers1

Linked