日期: 2021 年 6 月 20 日

numba可以对numpy的array使用gpu进行并行计算，从而进行加速。

import numpy as np
from numba import cuda

自定义numba cuda函数
@cuda.jit
def mat_add(A, B, C):
    j, i = cuda.grid(2)
    C[i, j] = A[i, j] + B[i, j]

a = np.arange(12).reshape(3, 4)
b = np.arange(12, 24).reshape(3, 4)
c = cuda.device_array_like(a_in)

a_dev = cuda.to_device(a)
b_dev = cuda.to_device(b)

blocks = (1, 1)
threads_per_block = (4, 3)

mat_add[blocks, threads_per_block](a_dev, b_dev, c)
print(c.copy_to_host())

其它

stride_j, stride_i = cuda.gridsize(2)

使用共享内存
tmp = cuda.shared.array((32, 32), numba.types.int32)
对共享内存的读写要注意进行同步
cuda.syncthreads()

pytorch现在提供给numba一个接口，可以使用numba直接对pytorch的gpu tensor进行操作，从而可以比较方便地实现自己想要的一些运算。

import torch
import numba.cuda

a = torch.arange(12).reshape(3,4).cuda()
a_nb = numba.cuda.as_cuda_array(a)
# 这里a_nb和a的数据是同一片内存空间的，因此a_nb对数据修改，a的数据也会同样修改，反之亦然。
# 得到a_nb后，就可以用自定义的并行运算函数进行操作了