欢迎光临散文网 会员登陆 & 注册

python使用numba库实现gpu加速

2023-02-11 15:08 作者:高频交易dragon  | 我要投稿


  Numba是一个针对Python的开源JIT编译器,由Anaconda公司主导开发,可以对Python原生代码进行CPU和GPU加速。Numba对NumPy数组和函数非常友好。


        使用Numba非常方便,只需要在Python原生函数上增加一个装饰器(Decorator)。Numba会将这些函数使用即时编译JIT方式编译成机器码,这些代码将以近乎机器码的速度运行。

2.GPU加速: 

与传统的Python CPU代码不同的是:


使用from numba import cuda引入cuda库

在GPU函数上添加@cuda.jit装饰符,表示该函数是一个在GPU设备上运行的函数,GPU函数又被称为核函数。

主函数调用GPU核函数时,需要添加如[1, 2]这样的执行配置,这个配置是在告知GPU以多大的并行粒度同时进行计算。gpu_func[1, 2]()表示开启一个block块,每个block块开启2个线程并行地执行gpu_func函数,函数将被并行地执行2次。

GPU核函数的启动方式是异步的:启动GPU函数后,CPU不会等待GPU函数执行完毕才执行下一行代码。必要时,需要调用cuda.synchronize(),告知CPU等待GPU执行完核函数后,再进行CPU端后续计算。这个过程被称为同步。

 CUDA将核函数所定义的运算称为线程(Thread),多个线程组成一个块(Block),多个块组成网格(Grid)。这样一个Grid可以定义成千上万个线程,也就解决了并行执行上万次操作的问题。例如,把前面的程序改为并行执行8次:可以用2个Block,每个Block中有4个Thread。原来的代码可以改为gpu_func[2, 4](),其中方括号中第一个数字表示整个Grid有多少个Block,方括号中第二个数字表示一个Block有多少个Thread。

2个bolck,4个线程


  CUDA提供了一系列内置变量,以记录Thread和Block的大小及索引下标。以[2, 4]这样的配置为例:

blockDim.x变量表示Block的大小是4,即每个Block有4个Thread,

threadIdx.x变量是一个从0到blockDim.x - 1(4-1=3)的索引下标,记录这是第几个Thread;gridDim.x变量表示Grid的大小是2,即每个Grid有2个Block,blockIdx.x变量是一个从0到gridDim.x - 1(2-1=1)的索引下标,记录这是第几个Block。

某个Thread在整个Grid中的位置编号为:threadIdx.x + blockIdx.x * blockDim.x

线程1的位置
注意数据存放的格式维度
显卡内存图

图7:GPU的内存模型。GPU里面的内存分为三种:per thread local memory, per block shared memory,和global memory。在实际编程的时候,需要考虑多用shared memory,少用global memory,因为shared比global的访问和存取速度快很多。

3 GPU的Hello world: 矩阵相乘

能够让人理解GPU编程的 Hello world 程序,便是矩阵相乘。这个纽约大学的教程非常棒,详细讲解了如何编写GPU程序进行矩阵相乘。我当时学习Numba和CUDA,这个教程发挥了很大的作用。


3.1介绍几个名词

首先,要弄清楚下面6个名词的概念。编写GPU程序,其实是一种CPU和GPU之间的“互动”,所谓的异构开发。


Host。代表CPU。

Device。代表GPU。

Host memory。RAM内存。

Device memory。GPU上的存储。

Kernal function。GPU函数,执行在device上面,调用者是host。

Device function。GPU函数,执行在device上面,调用者是kernal function或者device function


3.2 GPU程序的执行流程

图5可视化了GPU程序的流程:


把数组的数据(numpy.ndarray)从host memory拷贝到device memory上。

配置kernal function的参数,调用kernal function。参数有两种:一种用中括号[ ],一种用小括号( )。中括号的参数是threadperblock和blockpergrid,参考图6。小括号就是那种普通函数的参数。

几千个线程同时调用同一个kernal function,在GPU里面进行计算。kernal function的编写,是一门技术。

把GPU里面的运算结果,从device memory拷贝回host memory。THE END。


3.3 矩阵相乘源代码

这份Python代码,有6个函数,进行 C = A × B \mathbf{C} = \mathbf{A} \times \mathbf{B}C=A×B


cpu_mat_mul(A, B, C), 基于CPU的矩阵相乘,O ( n 3 ) O(n^3)O(n 

3

 )。

cpu_mat_mul_jit(A, B, C),和cpu_mat_mul相比,唯一的区别就是加了装饰器 @jit,这是一种编译优化工具,可以加快几十倍。

mat_mul_naive_kernal(A, B, C),矩阵相乘的GPU函数,使用global memory。

mat_mul_shared_kernal(A, B, C),矩阵相乘的GPU,使用了shared memory。shared memory比global memory访问速度快不少。

host_naive(A, B, C),mat_mul_naive_kernal的“配套设施”,kernal不能独立存在的,前前后后需要一些代码辅助。

host_optimized(A, B, C),mat_mul_shared_kernal的“配套设施”。如果host_naive比cpu_mat_mul快三千倍的话,host_optimized可能快三千五百倍。

Numba also exposes three kinds of GPU memory:

  • global device memory

  • shared memory

  • local memory

Find ways to parallelize sequential code

Minimize data transfers between the host and the device

Adjust kernel launch configuration to maximize device utilization

Ensure global memory accesses are coalesced

Minimize redundant accesses to global memory whenever possible

Avoid different execution paths within the same warp


python使用numba库实现gpu加速的评论 (共 条)

分享到微博请遵守国家法律