Как использовать Dask для запуска кода Python на GPU? - PullRequest
Новая
       19

Как использовать Dask для запуска кода Python на GPU?

1 голос
/ 18 мая 2019

У меня есть некоторый код, который использует 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, если это возможно?

1 Ответ

2 голосов
/ 18 мая 2019

Возможно, вы захотите просмотреть документацию Dask по графическим процессорам

Но, поскольку я использую CUDA, поместить submit перед функцией не работает.

Нет особой причины, почему так должно быть. All Dask запускает вашу функцию на другом компьютере. Это никак не меняет и не изменяет вашу функцию.

...