您发布的示例似乎来自(或被抄袭)一本名为“Python Parallel Programming Cookbook”的书,直到五分钟前我才听说过。老实说,如果我是那本书的作者,我会为包含这样一个 hacky、破碎的例子而感到羞耻。
以下是对您发布的内容及其输出的小修改:
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy
a = numpy.random.randn(5,5)
a = a.astype(numpy.float32)
a_gpu = cuda.mem_alloc(a.nbytes)
cuda.memcpy_htod(a_gpu, a)
mod = SourceModule("""
__global__ void doubleMatrix(float *a)
{
int idx = threadIdx.x + threadIdx.y*4;
a[idx] *= 2.f;
}
""")
func = mod.get_function("doubleMatrix")
func(a_gpu, block=(5,5,1))
a_doubled = numpy.empty_like(a)
cuda.memcpy_dtoh(a_doubled, a_gpu)
print a_doubled - 2.0*a
[警告:Python 2 语法]
In [2]: %run matdouble.py
[[ 0. 0. 0. 0. 0. ]
[ 0. 0. 0. 0. 0. ]
[ 0. 0. 0. 0. 0. ]
[ 0. 0. 0. 0. 0. ]
[ 0. -0.62060976 0.49836278 -1.60820103 1.71903515]]
即代码没有按预期工作,这可能是您混淆的根源。
this very recent answer 中描述了寻址存储在线性内存中的多维数组(如 numpy 数组)的正确方法。任何明智的程序员都会在您的示例中编写类似这样的内核:
__global__ void doubleMatrix(float *a, int lda)
{
int idx = threadIdx.x + threadIdx.y * lda;
a[idx] *= 2.f;
}
这样数组的前导维度作为参数传递给内核(在这种情况下应该是 5,而不是 4)。这样做会得到以下结果:
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy
a = numpy.random.randn(5,5)
a = a.astype(numpy.float32)
a_gpu = cuda.mem_alloc(a.nbytes)
cuda.memcpy_htod(a_gpu, a)
mod = SourceModule("""
__global__ void doubleMatrix(float *a, int lda)
{
int idx = threadIdx.x + threadIdx.y * lda;
a[idx] *= 2.f;
}
""")
func = mod.get_function("doubleMatrix")
lda = numpy.int32(a.shape[-1])
func(a_gpu, lda, block=(5,5,1))
a_doubled = numpy.empty_like(a)
cuda.memcpy_dtoh(a_doubled, a_gpu)
print a_doubled - 2.0*a
产生预期结果:
In [3]: %run matdouble.py
[[ 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0.]
[ 0. 0. 0. 0. 0.]]
对于 GPU 的函数调用,为什么将块定义为 5、5、1。
因为它是一个 5x5 元素的数组,所以在我的理解中
大小应该是 5,5 而不是 5,5,1。
在 CUDA 中,所有块都隐含地具有三个维度。 (5,5) 的块大小与 (5,5,1) 的块大小相同。最后一个维度可以忽略,因为它是一个(即块中的所有线程都有threadIdx.z = 1)。您不应该陷入的陷阱是将 CUDA 块或网格的尺寸与输入数组的尺寸混为一谈。有时让它们相同很方便,但同样没有必要甚至不建议这样做。本示例中正确编写的 BLAS 样式的内核(假设行主要存储顺序)可能如下所示:
__global__ void doubleMatrix(float *a, int m, int n, int lda)
{
int col = threadIdx.x + blockIdx.x * blockDim.x;
int row = threadIdx.y + blockDim.y * blockDim.y;
for(; row < m; row += blockDim.y * gridDim.y) {
for(; col < n; col += blockDim.x * gridDim.x) {
int idx = col + row * lda;
a[idx] *= 2.f;
}
}
}
[注:在浏览器中编写,未编译或测试]
这里任何合法的块和网格维度将正确处理任何大小的输入数组元素总数将适合有符号的 32 位整数。如果你运行太多线程,有些线程什么也不会做。如果你运行的线程太少,一些会处理多个数组元素。如果您运行与输入数组具有相同维度的网格,则每个线程将只处理一个输入,正如您正在研究的示例中的意图。如果你想了解如何选择最合适的块和网格大小,我建议从here开始。