使用numba cuda 加速Python运算

习惯了cuda c，可能会认为cuda和c才是黄金档搭。

Python作为一种开发效率比较高的脚本语言，有助于我们快速实现某种功能。

但是Python执行效率极其之慢。

这种情况下，用cuda的高并发特性，来提升Python执行速度，是一种很好的选择。

１.随机数生成

随机数生成是一项很重要的功能。

当Python自带的random，np.random在cuda函数中无法直接使用时，这是一个非常头疼的事。

有一个方法是将随机数/序列提前在cuda函数外实现好，再传递给cuda核函数使用，但是这就要占用cuda的显存，同时还要考虑加载数据的时间。

幸好的事numba提供了numba.cuda.random，可以便于我们生成随机数。

numba random官方网页中提供了一个示例，通过均匀分布来实现pi的计算。

由于numba.cuda.random.xoroshiro128p_normal_float64默认生成$N(0,1)$分布序列。

这里提供一个使用numba.cuda.random来生成复合高斯分布(如均值为100，方差为30的)的随机数：

有$N(\mu ,sigma)$分布的序列转成$N(0,1)$，标准化公式为：

$y=\frac{x-\mu }{\sqrt{\delta }}$

故从有$N(0,1)$分布的序列转成$N(\mu ,sigma)$分布，为：

$y=\sqrt{\delta }\cdot x+\mu$

代码如下：

from numba import cuda
from numba.cuda.random import create_xoroshiro128p_states, xoroshiro128p_normal_float64

import numpy as np
import math

@cuda.jit
def random_gen(rng_states,  out):
    """Find the maximum value in values and store in result[0]"""
    
    thread_id = cuda.grid(1)
    print("thread_id",thread_id)
    out[thread_id]=xoroshiro128p_normal_float64(rng_states, thread_id)
    out[thread_id]=int(out[thread_id]*math.sqrt(30)+100)
    
threads_per_block = 16
blocks = 16
rng_states = create_xoroshiro128p_states(threads_per_block * blocks, seed=1)
out = np.zeros((threads_per_block * blocks), dtype=np.float32)
out_d = cuda.to_device(out)
random_gen[blocks, threads_per_block](rng_states, out_d)
out = out_d.copy_to_host()
print('\n', out)

产生如下序列：

 [ 92. 100.  97. 101.  95. 103. 101. 105.  92. 101. 100.  97.  91.  90.
  97. 104. 100.  98.  97. 102. ...]

用numpy可求得均值和方差分别为：

99.609375  30.902099609375

生成整数随机序列，可以通过均匀分布，再经过适当放缩、平移实现，如采用(0,1)均匀分布实现[0,100]整数的均匀采样：

int(100*xoroshiro128p_uniform_float64(rng_states, col))

注意：rng_states 中的个数不宜设置过大，否则也可能会大大影响gpu运行速度。

2.numba njit与numba cuda性能的比较

以二维卷积运算为例。关于卷积运算，可参考离散信号的卷积与相关。

对大小为1111*1111的方阵进行卷积运算(卷积核为11*11）。

import numpy as np
import time
k=1111
spot=np.random.randint(0,10,(k,k)).astype(np.float64)
kernel=np.random.randint(0,10,(11,11)).astype(np.float64)

2.1 scipy实现

from scipy import signal 
t1=time.time()
spot_conv = signal.convolve2d(spot, kernel, mode='same')
t2=time.time()
print(t2-t1)

耗时约为0.31s。

2.2 个人编程实现

所用代码同离散信号的卷积与相关。

def Convolve2D(img, kernel,pad=None):    
    kernel=np.flipud(np.fliplr(kernel)) #卷积核反转
    h_k,w_k = kernel.shape   
    h_img,w_img = img.shape
    if pad is None:
        pad=(h_k-1)//2
    if pad != 0:         
        img_pad = np.zeros((img.shape[0] + pad*2, img.shape[1] + pad*2),dtype=np.float64)      
        img_pad[int(pad):int(-1 * pad), int(pad):int(-1 * pad)] = img      
    else:
        img_pad = img      
    
    out = np.zeros((h_img, w_img),dtype=np.float64)
    for y in prange(img.shape[1]):          
            for x in prange(img.shape[0]):            
                    out[x, y] = np.sum(kernel * img_pad[x: x + h_k, y: y + w_k])   
 
    return out
t3=time.time()
spot_conv2=Convolve2D(spot, kernel)
t4=time.time()
print(t4-t3)

耗时约为2s。

2.3 numba njit实现

import numba
from numba import njit,prange

@numba.njit(parallel=True) 
def Convolve2D(img, kernel,pad=None):    
    kernel=np.flipud(np.fliplr(kernel)) #卷积核反转
    h_k,w_k = kernel.shape   
    h_img,w_img = img.shape
    if pad is None:
        pad=(h_k-1)//2
    if pad != 0:         
        img_pad = np.zeros((img.shape[0] + pad*2, img.shape[1] + pad*2),dtype=np.float64)      
        img_pad[int(pad):int(-1 * pad), int(pad):int(-1 * pad)] = img      
    else:
        img_pad = img      
    
    out = np.zeros((h_img, w_img),dtype=np.float64)
    for y in prange(img.shape[1]):          
            for x in prange(img.shape[0]):            
                    out[x, y] = np.sum(kernel * img_pad[x: x + h_k, y: y + w_k])   
    return out

t3=time.time()
for i in range(10):
    spot_conv2=Convolve2D(spot, kernel)
t4=time.time()

print((t4-t3)/10)

耗时约为0.07s。已经相当不错了。

这里不要把prange写成range，否则会报错：Segmentation fault。

2.4 numba cuda实现

import numba
from numba import cuda
@cuda.jit
def conv_cuda(img,img_conv,kernel,h,w, ks):
    threadId = cuda.grid(1)
    row, col = threadId//w, threadId % w
    
    if (row < h and col < w):
        res = 0.0
        for i in range(ks):
            for j in range(ks):
                curRow = row - ks // 2 + i
                curCol = col - ks // 2 + j
                if (0 <= curRow < h and 0 <= curCol < w):
                    value = img[curRow*w+curCol] 
                    res += value * kernel[i*ks+j]
        
        cuda.syncthreads()
        img_conv[ row*w+col] = res 
        #img_conv变量是必须的，
        #如果直接将值赋给img[row*w+col]，由于同步问题其他线程就有可能使用的是卷积后的数，
        #而不是原始数据


h,w=img.shape
img_d = cuda.to_device(np.ascontiguousarray(img.reshape(-1)))
img_conv_d = cuda.to_device(np.ascontiguousarray(np.zeros(h*w)))
kernel=np.flipud(np.fliplr(kernel))
kernel_d = cuda.to_device(np.ascontiguousarray(kernel.reshape(-1)))
ks=kernel.shape[0]

t5=time.time()
nthreads = 1024
nblocks = (h*w // nthreads) + 1
for i in range(10):
    conv_cuda[nblocks, nthreads](img_d,img_conv_d,kernel_d,h,w, ks) 
    img_conv3=img_conv_d.copy_to_host()
    img_conv3=img_conv3.reshape(h,w)
t6=time.time()

print((t6-t5)/10)

最终结果为0.03s。

变量在从host复制到device前，最好使用np.ascontiguousarray转为连续存储。

如果是多张图片进行卷积运算，设置设置二维的线程数。

@cuda.jit
def imgs_conv_cuda(imgs,imgs_conv, imgs_num, psf, h,w,ks):
    imgId= cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    imgPixel_Id =cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
    
    #cuda.grid(2)
    
    row, col = imgPixel_Id//w, imgPixel_Id % w

    if imgId < imgs_num:
        if (row < h and col < w):
            intens = 0.0
            for i in range(ks):
                for j in range(ks):
                    curRow = row - ks // 2 + i
                    curCol = col - ks // 2 + j
                    if (0 <= curRow < h and 0 <= curCol < w):
                        value = imgs[imgId, curRow*w+curCol]
                        intens += value * psf[i*ks+j]            
            imgs_conv[imgId, row*w+col] = intens

获取线程两个维度id号的两种方法：

    imgId= cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    imgPixel_Id =cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
    imgId,imgPixel_Id =cuda.grid(2)

参考文献

[1] https://numba.readthedocs.io/en/stable/
[2] 基于 Numba 的 CUDA Python 编程简介
[3] https://numba.pydata.org/numba-doc/latest/cuda/random.html
[4] An introduction to CUDA in Python (Part 3)