pyCuda - dostęp do platformy nVidia CUDA w Pythonie

17 November 2008 Comments

pyCuda to Pythonowe API dla bibliotek platformy nVidia CUDA - służącej do obliczeń w GPU karty graficznej. Opis platformy CUDA dostępny jest w Bibliotece Linuksa. By cieszyć się działającym modułem pyCUDA potrzebujemy odpowiednią kartę graficzną obsługującą tą platformę, pakiet SDK CUDA, a także moduły numpy, lapack i pytools (ze strony pycuda). pyCuda można zainstalować ręcznie, bądź też za pomocą easy_install. W razie wątpliwości należy sięgnąć do dokumentacji.

Podstawy pyCuda

By móc dokonywać obliczeń za pomocą PyCuda na początek trzeba moduł załadować i stworzyć "kontekst" CUDA:

import pycuda.driver as cuda

cuda.init()
assert cuda.Device.count() >= 1

dev = cuda.Device(0)
ctx = dev.make_context()

Kolejnym etapem będzie zazwyczaj przesłanie danych do karty graficznej. W przypadku PyCuda najczęściej przesyłane będą dane z tablic/macierzy numpy. Dla przykładu stwórzmy macierz 4x4 złożoną z losowych liczb:

import numpy
a = numpy.random.randn(4,4)

Domyślnie numpy używa liczb o podwójnej precyzji, których CUDA obecnie nie obsługuje. Należy zastosować liczby o pojedynczej precyzji:

a = a.astype(numpy.float32)

A gdy mamy macierz musimy mieć miejsce, do którego ją prześlemy. Należy zaalokować na to pamięć w karcie graficznej:

a_gpu = cuda.mem_alloc(a.size * a.dtype.itemsize)

I przesłać dane:

cuda.memcpy_htod(a_gpu, a)

Jako prosty przykład napiszemy kod, który podwoi wszystkie liczby w a_gpu. By tego dokonać należy napisać odpowiadający tej operacji kod w CUDA C i przekazać go do konstruktora

pycuda.driver.SourceModule:
mod = cuda.SourceModule("""
  __global__ void doublify(float *a)
  {
    int idx = threadIdx.x + threadIdx.y*4;
    a[idx] *= 2;
  }
  """)

Jeżeli nie ma błędów, to kod zostanie skompilowany i załadowany w karcie graficznej. Po załadowaniu funkcji można uzyskać do niej dostęp i ją wykonać na naszej macierzy:

func = mod.get_function("doublify")
func(a_gpu, block=(4,4,1))

Po dokonaniu obliczeń pobieramy dane z karty i wyświetlamy, wraz z oryginałem:

a_doubled = numpy.empty_like(a)
cuda.memcpy_dtoh(a_doubled, a_gpu)
print a_doubled
print a

Przykładowy wynik to:

[[ 0.51360393  1.40589952  2.25009012  3.02563429]
 [-0.75841576 -1.18757617  2.72269917  3.12156057]
 [ 0.28826082 -2.92448163  1.21624792  2.86353827]
 [ 1.57651746  0.63500965  2.21570683 -0.44537592]]
[[ 0.25680196  0.70294976  1.12504506  1.51281714]
 [-0.37920788 -0.59378809  1.36134958  1.56078029]
 [ 0.14413041 -1.46224082  0.60812396  1.43176913]
 [ 0.78825873  0.31750482  1.10785341 -0.22268796]]

Podany kod można uprościć korzystając z dodatkowych elementów pyCuda, a dokładnie gpuarray:

import pycuda.gpuarray as gpuarray
import pycuda.driver as cuda
import numpy

cuda.init()
assert cuda.Device.count() >= 1

dev = cuda.Device(0)
ctx = dev.make_context()

a_gpu = gpuarray.to_gpu(numpy.random.randn(4,4).astype(numpy.float32))
a_doubled = (2*a_gpu).get()
print a_doubled
print a_gpu

RkBlog

Podstawy Pythona, 17 November 2008