使用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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

产生如下序列：

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

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

99.609375  30.902099609375
1

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

int(100*xoroshiro128p_uniform_float64(rng_states, col))
1

参考文献

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