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

一起养成写作习惯！这是我参与「掘金日新计划 · 4 月更文挑战」的第19天，点击查看活动详情。

不过这里存在不同步问题，由于每一个 thread 速度不一致，也就是共享内存不同步的问题，也就是当我们在计算时用到共享块，这是可能共享块加载更新了数据，和计算和其对应共享数据不同步问题，为了解决这个问题我们添加语句来解决问题。

cuda.syncthreads()
复制代码

首先第一个同步我们会放在将数据读取到共享内存 sA 和 sB 上调用一下 cuda.syncthreads()来进行一次同步。也就是需要确保线程同步，在将数据加载到共享内存之后，然后是在计算矩阵乘积，如果不进行这些线程同步，就有可能在一个线程同时更新共享数组的时候，另一个线程仍然在使用这个共享进行乘法，从而产生错误的结果。

from numba import cuda
复制代码

调用这个命令来查看 GPU 的一些信息。

cuda.detect()
复制代码

首先我们在 cpu 上创建一个数组

arr_cpu = np.random.randint(0,10,size=(2000,2000))
复制代码

接下来把在 cpu 上创建好的数组转到 GPU 上，调用 cuda 的· to_device 方法，这里简单说一下，我们 GPU 上进行编程，通常把 cpu 和内存称为 host ，相对于与 host 的 GPU 以及显存我们都称为 device，所以在命名习惯上，我们通常 array 加上 d 或者 ad 这样前缀来表示 device

d_arr = cuda.to_device(arr_cpu)
复制代码

接下来实现矩阵乘法运算，也就是对矩阵 A 和 B 相乘后，结果保存在矩阵 C 中

@cuda.jit
def matmul(A, B, C):
  i, j = cuda.grid(2)
  if i < C.shape[0] and j < C.shape[1]:
    tmp = 0
    for k in range(A.shape[1]):
      tmp += A[i,k] * B[k,j]
    C[i,j] = tmp
复制代码

首先我们创建一个 grid ，然后我们需要让这个 grid 大小输出矩阵矩阵大小一致，这样我们得到 thread 数量是和 C 矩阵元素数量一致，每一个 thread 用于计算矩阵 ij 位置的值，这个值就是用 A 矩阵 i 行元素与 B 矩阵第 j 列元素对应元素相乘得到数相加求和得到结果作为 C 矩阵 ij 位置的值

cp.random.seed(12)
A = cp.random.uniform(1,10,size=(2000,2000),dtype=np.float64)
B = cp.random.uniform(1,10,size=(2000,2000),dtype=np.float64)
C = cp.zeros((2000,2000),dtype=np.float64)
复制代码

这里创建 A 和 B 都是 2000 维的方阵，每一个值都是一个 1 到 10 之间随机数，然后所以 C 矩阵也是 2000 维方阵

C
array([[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., 0., 0., ..., 0., 0., 0.], 
    [0., 0., 0., ..., 0., 0., 0.]])
复制代码

首先我们需要定义要给 grid 然后将划分 16x16 的块(block) 每一个 block 为 125x125 大小矩阵，所以一共有 2000x2000 计算单元

threadsperblock = (16,16)
blockspergrid_x = int(np.ceil(C.shape[0] / threadsperblock[0]))
blockspergrid_y = int(np.ceil(C.shape[1] / threadsperblock[1]))
blockspergrid = (blockspergrid_x,blockspergrid_y)
print(blockspergrid)
print(f"The kernel will be executed up to element {threadsperblock[0]*blockspergrid_x}")
复制代码

(125, 125) 
The kernel will be executed up to element 2000
复制代码

matmul[blockspergrid,threadsperblock](A,B,C)
复制代码

下面方法实现了一个快速的矩阵的乘法运算，这里引入共享内存，来减少对全局内存的访问次数，这样好处加快了运算。为了避免线程之间不同步问题，分别在更新共享内存和利用共享内存进行计算处添加同步方法来解决。

@cuda.jit
def fast_matmul(A,B,C):
  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

  if x >= C.shape[0] and y >= C.shape[1]:
    return
  tmp = 0.
  for i in range(bpg):
    sA[tx, ty] = A[x, ty + i * TPB]
    sB[tx, ty] = B[tx + i * TPB, y]

    cuda.syncthreads()

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

    cuda.syncthreads()
  
  C[x,y] = tmp
复制代码

SIZE = 4000
A = cp.random.uniform(1,10,size=(SIZE,SIZE),dtype=np.float32)
B = cp.random.uniform(1,10,size=(SIZE,SIZE),dtype=np.float32)
C_fast = cp.zeros((SIZE,SIZE),dtype=np.float32)
C_slow = cp.zeros((SIZE,SIZE),dtype=np.float32)
复制代码

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

猜你喜欢