在 cuda 上运行 python 代码 — numba 是一个支持 cuda 的 python (2)

我们为什么要加入 block 这个结构，这是因为对于每一个 block 都共享一个缓存，也就是位于同一个block 的所有 thread 都共享一块内存,而且对这块内存访问要远远快于访问全局内存。

A = np.array([
    [1,2,1],
    [3,1,2],
    [1,2,0]
])

B = np.array([
    [1,1,0],
    [2,1,1],
    [3,2,1]
])

我们简单复习一下线性代数矩阵相乘，对于上图中 A 和 B 相乘，矩阵相乘需要满足一定前提，也就是 A 矩阵列需要和 B 矩阵行数保持一致，对于相乘得到矩阵 C 的 (i,j) 元素等于用 A 矩阵 i 行所有元素和 B 矩阵第 j 列所有元素两两相乘后在相加求和。

那么在 cuda 中我们是如何通过并行计算来加速矩阵乘法的呢?

从上面图不难看出我们定义一个 grid 形状和输出 C 形状相同，然后每一个位置 kernel 通过将 A 矩阵对应行元素和 B 矩阵相应的列元素对应位置元素相乘后相加得到 C 矩阵元素值，所有这些操作都是在一个线程里独立完成的。

@cuda.jit
def matmul(A,B,C):
    1,2 = cuda.grid(2)
    if 1 < 4 and 2 < 4:
        tmp = 0
        for k in range(4)
            tmp = A[1,k]*B[k,2]
        C[1,2] = tmp

这样做也就是意味着我们每计算一个 C 元素需要从全局共享中将 A 矩阵和 B 矩阵的 4 个值分别加载，接下来我们就来介绍一种更有效方法来加载数据，也就是通过共享内存。

共享内存可以帮助我们更好理解为什么要设计 block 这个结构。我们将矩阵结果划分若干 block 每一个 block 都是方形的矩阵。每一个 block 都有自己共享内存。

当我们 block 创建之后 cuda.shared.array 创建在共享内存中创建数组，数组大小和 block 的形状保持一致，在 block 中所有线程都可以访问到这块共享。

为每一个 block 创建 2 共享内存分别用于存储 A 和 B 矩阵上对应位置上数据

例如在上图中，sA 存储数据是 A 矩阵上橘黄色部分的数据，同样 sB 存储的是 B 矩阵橘黄色的部分的数据。

sA = cuda.shared.array(shape=(TPB,TPB),dtype=float32)
sB = cuda.shared.array(shape=(TPB,TPB),dtype=float32)

x, y = cuda.grid(2) 
tx = cuda.threadIdx.x 
ty = cuda.threadIdx.y 
bpg = cuda.gridDim.x 

这里的 TPB 是每一个每一个 block 所具有线程数，cuda.gridDim 告诉我们 grid 有多少个 block。接下来我们循环 grid 每一个 block 然后对于每一个遍历到 block 我们将矩阵 A 和 B 对应位置元素放入到共享内存上。

for i in range(bpg): 
    sA[tx, ty] = A[x, ty + i * TPB] 
    sB[tx, ty] = B[tx + i * TPB, y]

将数据放入 sA 和 sB 后，可以对两个共享矩阵进行相乘得到一个中间结果。

接下来继续遍历到 A 矩阵的一下行和 B 矩阵下一列，将数据存储到 block 对应 sA 和 sB 并且对两个共享矩阵进行相乘得到一个矩阵。

for j in range(TPB):
      tmp += sA[tx, j] * sB[j, ty]

我们将两个对应位置相乘的矩阵进行相加就得到想要结果，这个大家可以自己动手试一试。