我有一个稍微修改的例子从numba官方文件如下:
from numba import njit
import numpy as np
@njit
def do_sum(A, lb, ub):
n = len(A)
acc = 0.0
for i in range(n):
a = 0.0 if np.isnan(A[i]) else A[i]
acc += abs(max(min(a, ub[i]), lb[i]))
return acc
基本上,任务是对一个numpy数组进行绝对有界求和,这个数组可能有np.nan
。A
,lb
和ub
都是长度相同的1d
数组。长度通常在数千或数万的范围内。而且,函数预计会被调用数百万次-这就是为什么我想优化它。
由于A
可能包含np.nan
,所以我似乎不能使用@njit(fastmath=True)
,但是,当A
不包含np.nan
时,我的基准测试结果显示,使用@njit(fastmath=True)
明显比@njit
快。
我的问题是,在这两者之间是否存在一些最佳点,这样我就可以使代码与fastmath
一起工作,并获得比上面的do_sum
实现更快的速度?
或者说,无论什么样的方法都可以让do_sum
更快,这将是非常受欢迎的。
为简单起见,我们可以假设lb
和ub
不包含nan
值,但如果解决方案也能处理它们,那就更好了。
(By顺便说一下,这个代码示例与官方numba文档中说明parallel=True
的示例非常相似,但当我尝试在上面的do_sum
中添加parallel=True
时,我得到了显著的减慢。不知道为什么会这样。)
1条答案
按热度按时间rdlzhqv91#
我做了很多测试,成功的是并行化
测试代码,如果其他人想尝试优化: