在Python中使用CUDA展开简单可并行化的for循环

Question

我在Python中有一个for循环，我想把它展开到GPU上去。我想这应该有个简单的解决办法，但我还没找到。

我们的函数会遍历一个numpy数组中的元素，进行一些数学运算，然后把结果存储到另一个numpy数组里。每次循环都会把一些结果加到这个结果数组中。我们代码可能会有一个很大的简化，像这样：

import numpy as np

a = np.arange(100)
out = np.array([0, 0])
for x in xrange(a.shape[0]):
  out[0] += a[x]
  out[1] += a[x]/2.0

我该如何在Python中展开这样的循环，以便在GPU上运行呢？

Answer 1

首先可以去这个网站看看：http://documen.tician.de/pycuda/，那里有个例子。

import pycuda.autoinit
import pycuda.driver as drv
import numpy

from pycuda.compiler import SourceModule
mod = SourceModule("""
__global__ void multiply_them(float *dest, float *a, float *b)
{
  const int i = threadIdx.x;
  dest[i] = a[i] * b[i];
}
""")

multiply_them = mod.get_function("multiply_them")

a = numpy.random.randn(400).astype(numpy.float32)
b = numpy.random.randn(400).astype(numpy.float32)

dest = numpy.zeros_like(a)
multiply_them(
        drv.Out(dest), drv.In(a), drv.In(b),
        block=(400,1,1), grid=(1,1))

print dest-a*b

你需要把想要并行处理的代码部分放在C语言的代码段里，然后从Python中调用它。

对于你的例子，数据的大小需要大于100，这样才有意义。你需要找到一种方法把数据分成多个块。如果你想加1,000,000个数字，可以把它分成1000个块。在并行处理的代码中分别加每个块的值，然后在Python中把结果加起来。

不过，加法并不是这种并行处理的自然任务。显卡通常是对每个像素执行相同的任务，而你的任务需要对多个像素进行操作。

可能先学习CUDA会更好。相关的讨论可以参考这个链接：理解CUDA的网格维度、块维度和线程组织（简单解释）