Numba加速器
我把一段matlab代码重写迁移到python之后,运行速度慢了近100倍。经过代码分析,找到最耗时间的是python里的for循环。这可能与python语言本身的运行效率有关系,python是一门解释型的高级语言,在python被操作系统理解之前,要先转化为pyc字节码给python虚拟机,通过虚拟机与硬件操作系统交互。相比于一些更底层的语言,python的设计使得它更容易被程序员理解和使用,但trade-off是对机器来说更不容易理解,运行效率降低了。
Numba就是为了解决python计算效率低的痛点开发的JIT(just-in-time)编译器。简单来说,Numba把一部分python代码编译成机器语言,只需在第一次调用时编译一次,之后再调用都能以机器语言的速度运行。使用Numba之后,我的同一段的代码运算速度和matlab相近甚至更快了。
Numba安装
用conda直接装
1 | $ conda install numba |
在CPU上使用Numba
使用Numba只需要把需要加速的那部分代码封装成function,在function前加一个@jit
装饰器(decorator)就行了。是的,就这么简单。
1 | from numba import jit |
需要注意的是,并不是所有的python对象Numba都支持,尽量只加速最耗时的那部分代码,这样既能避免代码中包含Numba不支持的python类型导致报错,又可以给代码瘦身提高编译速度。详见官网的说明。
编译有两个模式可选:object mode和nopython mode。后者默认你对要加速的代码充分了解,不包含Numba不支持的python类型,直接编译成机器语言。否则会直接报错。而object mode则不会,遇到不能编的就退出JIT,按python原生的来。所以想要发挥Numba最好的性能,尽量使用nopython mode。
1 |
|
Cpu parallel
Vectorize
Numba+GPU并行
GPU与并行
Numba使用GPU
- 简单并行
- 多维并行
并行适合处理互相独立的运算,如果是对同一个数的循环累加,用单个线程代替一次循环可能会出问题。比如,我要并行化的原代码是一个两层嵌套循环,其中外层相互独立,内层是对同一个数的累加
1 | def calx(x, ss2, model_size, a, zwrap, smooth_steps): |
假如我把两层循环都并行化,这里我令block的维度是2*smooth_steps,用block里的一个thread计算一次内循环;令grid的维度是model_size,用grid里的每个block计算一次外循环。
1 |
|
输出cpu的calx函数和gpu的gpu_calx函数计算结果的前三个数,发现两者不相同,有1e-2的差异
1 | calx_result[0:3]= |
或许可以这样来解释:用并行的线程取代累加循环时,每个线程并不关心其他线程是否计算完成。假设线程a从内存中取出x_device[i - ss2, 0]执行加法运算时,可能同时有线程b也在执行相同的操作,而两者取出的x_device[i - ss2, 0]是同一值。这样两个线程执行完加法返回值的时候,其实相当于只执行了一次加法运算,线程a或b的其中一个运算是无效的。所以造成了最后结果的差异。
于是,我只能放弃并行内层的累加循环
1 |
|
然鹅,并行后的代码实际运行时间并不比加了@jit
的cpu函数快。原因可能是计算本身的dimension不够大,gpu并行节省的时间还打不过gpu和cpu之间数据传输花费的时间。并不是gpu并行就一定比cpu快,这里有一个数据计算量和调用线程、数据传输时间之间的trade-off,cpu虽然不如gpu核数多,但往往优化做的很好,所以小数据量的计算在cpu上可能更好。
-复杂函数的并行
刚刚并行化的是最里层的一个简单函数,只涉及数组元素的算术运算。除此之外,我还想要把最外层的一个独立for循环并行化,这个for循环计算的是独立的40个集合预报成员。但问题是,这层for循环是相当外层的循环了,在它内部还有多层函数嵌套和一些numpy方法。
1 |
|
为了解决多层函数嵌套的编译问题,我在每个内层函数前加上@cuda.jit(device=True)
装饰,在这篇文章有提到device function:
On the other hand, a
device function
can only be invoked from inside a device (by a kernel or another device function). The plus point is, you can return a value from adevice function
. So, you can use this return value of the function to compute something inside akernel function
or adevice function
.
但是在@cuda.jit
装饰的函数内不能使用numpy方法是硬伤,numba官方不支持在cuda并行编译的程序里使用numpy methods
Unsupported numpy features:
- array creation APIs.
- array methods.
- functions that returns a new array.
所以只能,放弃了!!!Orzzzzz…
- 内存
Cuda.to_device
Shared memory
参考