我试图弄清楚如何使用numba
生成矢量化数组操作的numpy
样式的ufunc。我注意到我的性能非常慢,因此我尝试通过按numba FAQ在我的代码中调用以下内容来进行调试:
import llvmlite.binding as llvm
llvm.set_option('', '--debug-only=loop-vectorize')
显然,由于内存冲突,我的循环没有被向量化。
同样从FAQ页面,我看到这种情况发生在'当内存访问模式不平凡时'。我不清楚这是什么意思,但是我尝试向量化的代码对我来说似乎微不足道:
@guvectorize(['void(f4[:,:], b1[:,:], f8, f4, f4[:,:])'],
'(n,m), (n,m), (), () -> (n,m)', cache=True)
def enforce_cutoff(img, mask, max, nodata, out):
for i in range(img.shape[0]):
for j in range(img.shape[1]):
if mask[i,j]:
out[i,j] = nodata
else:
if img[i,j]<max:
out[i,j] = img[i,j]
else:
out[i,j] = max-0.1
关于为何无法将其向量化以及如何解决它的任何线索,我们将不胜感激。我对numba相当陌生,对LLVM也不熟悉,所以我对此不太了解。
LLVM的完整输出在这里:
LV: Checking a loop in "_ZN7AtmCorr18enforce_cutoff$241E5ArrayIfLi2E1A7mutable7alignedE5ArrayIbLi2E1A7mutable7alignedEdf5ArrayIfLi2E1A7mutable7alignedE" from enforce_cutoff
LV: Loop hints: force=? width=0 unroll=0
LV: Found a loop: B40.us
LV: Found an induction variable.
LV: Found an induction variable.
LV: Can't vectorize due to memory conflicts
LV: Not vectorizing: Cannot prove legality.
LV: Checking a loop in "__gufunc__._ZN7AtmCorr18enforce_cutoff$241E5ArrayIfLi2E1A7mutable7alignedE5ArrayIbLi2E1A7mutable7alignedEdf5ArrayIfLi2E1A7mutable7alignedE" from <numba.npyufunc.wrappers._GufuncWrapper object at 0x0000020A848A6438>
LV: Loop hints: force=? width=0 unroll=0
LV: Not vectorizing: Cannot prove legality.
LV: Checking a loop in "_ZN7AtmCorr18enforce_cutoff$241E5ArrayIfLi2E1A7mutable7alignedE5ArrayIbLi2E1A7mutable7alignedEdf5ArrayIfLi2E1A7mutable7alignedE" from <numba.npyufunc.wrappers._GufuncWrapper object at 0x0000020A848A6438>
LV: Loop hints: force=? width=0 unroll=0
LV: Found a loop: B40.us
LV: Found an induction variable.
LV: Found an induction variable.
LV: Found an induction variable.
LV: Found an induction variable.
LV: Did not find one integer induction var.
LV: Can't vectorize due to memory conflicts
LV: Not vectorizing: Cannot prove legality.
LV: Checking a loop in "_ZN7AtmCorr18enforce_cutoff$241E5ArrayIfLi2E1A7mutable7alignedE5ArrayIbLi2E1A7mutable7alignedEdf5ArrayIfLi2E1A7mutable7alignedE" from <numba.npyufunc.wrappers._GufuncWrapper object at 0x0000020A848A6438>
LV: Loop hints: force=? width=0 unroll=0
LV: Found a loop: B20.us.us
LV: Found an induction variable.
LV: Can't vectorize due to memory conflicts
LV: Not vectorizing: Cannot prove legality.
一种可能的解决方法是确保数组是C连续的。如果它们不连续,将被复制。
示例
import numba as nb
import numpy as np
@nb.njit(cache=True,parallel=True)
def enforce_cutoff_2(img, mask, max, nodata, out):
#create a contigous copy if array isn't c-contiguous
img=np.ascontiguousarray(img)
mask=np.ascontiguousarray(mask)
for i in nb.prange(img.shape[0]):
for j in range(img.shape[1]):
if mask[i,j]:
out[i,j] = nodata
else:
if img[i,j]<max:
out[i,j] = img[i,j]
else:
out[i,j] = max-0.1
Timings
#contiguous arrays
img=np.random.rand(1000,1000).astype(np.float32)
mask=np.random.rand(1000,1000)>0.5
max=0.5
nodata=1.
out=np.empty((img.shape[0],img.shape[1]),dtype=np.float32)
%timeit enforce_cutoff_2(img, mask, max, nodata, out)
#single-thread
#678 µs ± 3.72 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#parallel
#143 µs ± 1.87 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
#non contigous arrays
img=np.random.rand(2000,2000).astype(np.float32)
mask=np.random.rand(2000,2000)>0.5
img=img[0:-1:2,0:-1:2]
mask=mask[0:-1:2,0:-1:2]
max=0.5
nodata=1.
out=np.empty((img.shape[0],img.shape[1]),dtype=np.float32)
%timeit enforce_cutoff_2(img, mask, max, nodata, out)
#single threaded
#with contiguous copy
#1.78 ms ± 9.58 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#without contiguous copy
#5.76 ms ± 20.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
#parallel
#with contiguous copy
#1.42 ms ± 7.03 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
##without contiguous copy
#1.08 ms ± 75.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)