PyCuda初步认识

PyCuda可以再python中嵌入一个cuda的kernel,下面来尝试一下。

使用之前,需要初始化PyCuda。

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule

接着我们使用numpy,再将其中的数据转移到device的显存中。

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

先转成单精度浮点数。

a = a.astype(numpy.float32)

在device上分配内存。

a_gpu = cuda.mem_alloc(a.nbytes)

最终,把数据转移过去。

cuda.memcpy_htod(a_gpu, a)

接下来,我们定义一个CUDA的kernel。注意,这个CUDA C的kernel是以字符串给出的,将会被编译为能在GPU上可执行的文件。

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

接下来,我们获取这个kernel函数,然后在我们的a_gpu上launch这个kernel。

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

最终,我们将数据取回到我们的内存中并打印出来。

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

最终的结果像这样。

[[ 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]]

S12041-PyCUDA-Simpler-GPU-Programming-Python.pdf

PyCUDA_crashcourse2_split.pdf

参考资料