Numba:手动循环比带有numpy数组的+ = c*b更快？

Question

我想做一个'daxpy'(将向量添加到第二个向量的标量倍数并将结果赋给第一个)并numpy使用numba.做了以下测试,我注意到自己编写循环要快得多a += c * b.

我没想到这个.这种行为的原因是什么？

import numpy as np
from numba import jit

x = np.random.random(int(1e6))
o = np.random.random(int(1e6))
c = 3.4

@jit(nopython=True)
def test1(a, b, c):
    a += c * b
    return a

@jit(nopython=True)
def test2(a, b, c):
    for i in range(len(a)):
        a[i] += c * b[i]
    return a

%timeit -n100 -r10 test1(x, o, c)
>>> 100 loops, best of 10: 2.48 ms per loop
%timeit -n100 -r10 test2(x, o, c)
>>> 100 loops, best of 10: 1.2 ms per loop

Answer 1

要记住的一件事是“手动循环”numba非常快，本质上与 numpy 操作使用的 c 循环相同。

在第一个示例中，有两个操作，c * b分配/计算临时数组 ( )，然后将该临时数组添加到a。在第二个示例中，两个计算都在同一个循环中进行，没有中间结果。

理论上，numba可以融合循环并优化 #1 以执行与 #2 相同的操作，但它似乎并没有这样做。如果您只是想优化 numpy 操作，numexpr也可能值得一看，因为它正是为此而设计的 - 尽管可能不会比显式融合循环做得更好。

In [17]: import numexpr as ne

In [18]: %timeit -r10 test2(x, o, c)
1000 loops, best of 10: 1.36 ms per loop

In [19]: %timeit ne.evaluate('x + o * c', out=x)
1000 loops, best of 3: 1.43 ms per loop