Comparing numba njit/vectorize/guvectorize

Question

我一直在测试以下块numba加速:

import numpy as np
import timeit
from numba import njit
import numba 
@numba.guvectorize(["void(float64[:],float64[:],float64[:],float64, float64, float64[:])"],
"(m),(m),(m),(),()->(m)",nopython=True,target="parallel")
def func_diff_calc_numba_v2(X,refY,Y,lower,upper,arr):
fac=1000
for i in range(len(X)):
if X[i] >=lower and X[i] <upper:
diff=Y[i]-refY[i]
arr[i] = diff**2*fac
else:
arr[i] = 0
@numba.vectorize('(float64, float64, float64, float64, float64)',nopython=True,target="parallel")
def func_diff_calc_numba_v3(X,refY,Y,lower,upper):
fac=1000
if X >= lower and X < upper:
return (Y-refY)**2*fac
else:
return 0.0
@njit
def func_diff_calc_numba(X,refY,Y,lower,upper):
fac=1000
arr=np.zeros(len(X))
for i in range(len(X)):
if X[i] >=lower and X[i] <upper:
arr[i]=(Y[i]-refY[i])**2*fac
else:
arr[i] = 0
return arr
np.random.seed(69)
X=np.arange(10000)
refY = np.random.rand(10000)
Y = np.random.rand(10000)
lower=1
upper=10000
print("func_diff_calc_numba: {:.5f}".format(timeit.timeit(stmt="func_diff_calc_numba(X,refY,Y,lower,upper)", number=10000, globals=globals())))
print("func_diff_calc_numba_v2: {:.5f}".format(timeit.timeit(stmt="func_diff_calc_numba_v2(X,refY,Y,lower,upper)", number=10000, globals=globals())))
print("func_diff_calc_numba_v3: {:.5f}".format(timeit.timeit(stmt="func_diff_calc_numba_v3(X,refY,Y,lower,upper)", number=10000, globals=globals())))

v2和v3的加速明显不同:

func_diff_calc_numba: 0.58257
func_diff_calc_numba_v2: 0.49573
func_diff_calc_numba_v3: 1.07519

如果我将迭代次数从10,000更改为100,000，则:

func_diff_calc_numba: 1.67251
func_diff_calc_numba_v2: 4.85828
func_diff_calc_numba_v3: 11.63361

我原以为vectorize和guvectorize在加速方面几乎是相似的，但是njit和guvectorize在时间上几乎是相等的，而vectorize分别比guvectorize和njit慢2倍和10倍。是我的实现有什么问题还是别的什么?

Answer 1

任务(函数+输入)可能太小/太简单，无法有效地并行化，从而导致这样做的开销增加了总运行时间。如果您将两者都编译为默认的cpu目标，我想差异会消失吗?

因为你的输入是一维的，在给定的ufunc签名下，guvectorize不会并行化任何东西，因为只有一个任务。

可以通过将签名设置为"(),(),(),(),()->()"来完成类似的并行比较，基本上告诉它也(像vectorize一样)应用函数元素。结果应该又很接近了。但是你会发现，在这种情况下，并行化的开销使这两种情况都变得更糟。

对我来说时序是:

1. 同时使用target="parallel"和"(m),(m),(m),(),()->(m)":

   numba_guvec    : 0.26364
   numba_vec      : 3.26960
   

2. 同时使用target="cpu"和"(m),(m),(m),(),()->(m)":

   numba_guvec    : 0.21886
   numba_vec      : 0.26198
   

3. 同时使用target="parallel"和"(),(),(),(),()->()":

   numba_guvec    : 3.05748
   numba_vec      : 3.15587
   

如果将@njit(parallel=True)与numba.prange进行比较，您可能会发现类似的行为。

最后，并行化一些东西只是涉及一些额外的工作，这只有在足够大(慢)的任务时才值得。