如何使用 Dask 在 GPU 上运行 python 代码？答案

【问题标题】：How to use Dask to run python code on the GPU?如何使用 Dask 在 GPU 上运行 python 代码？
【发布时间】：2019-05-17 23:15:10
【问题描述】：

我有一些使用 Numba cuda.jit 的代码，以便我在 gpu 上运行，如果可能的话，我想在它上面添加 dask。

示例代码

#!/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()

如果我的函数 addingNumbersCUDA 没有使用任何 CUDA，我只需将 client.submit 放在我的函数前面（以及之后的收集），它就会起作用。但是，由于我使用的是 CUDA，所以将 submit 放在函数前面是行不通的。 dask 文档说您可以针对 gpu，但不清楚如何在实践中实际设置它。如果可能，我将如何设置我的函数以将 dask 与目标 gpu 和 cuda.jit 一起使用？

【问题讨论】：

标签： python dask dask-distributed

【解决方案1】：

你可能想看看Dask's documentation on GPUs

但是，由于我使用的是 CUDA，所以将提交放在函数前面是行不通的。

没有什么特别的理由会出现这种情况。所有 Dask 都会在另一台计算机上运行您的功能。它不会以任何方式改变或修改您的功能。

【讨论】：

我已经为 gpu 设置了我的 dask worker，我的函数设置为 "x = client.submit(addingNumbersCUDA[blockspergrid, threadsperblock], big_array, big_array2, save_array, resources={'GPU': 1})”，但是当我收集它时，返回“无”。你知道这是为什么吗？
你的函数没有返回值。看起来您正在存储到 out 参数作为副作用。我建议构建一个返回值的函数。
我不能专门从 numba cuda 函数返回值。 “内核不能显式返回值；所有结果数据都必须写入传递给函数的数组。”是否可以从 dask 获取存储在我的数组中的结果？
您可以围绕您的 Numba 函数创建另一个函数来处理突变和就地行为。 Dask 不支持就地操作。通常，您的新函数会分配一个输出数组，调用您的就地操作 numba 函数，然后正常返回输出数组。