Numba：手动循环比使用 numpy 数组的 a += c * b 更快？

Question

我想用 numpy 使用 numba 做一个 'daxpy'（将第二个向量的标量倍数添加到一个向量中，并将结果分配给第一个向量）。在进行以下测试时，我注意到自己编写循环比 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 中的 'manual looping' 非常快，与 numpy 操作使用的 c-loop 基本相同。

在第一个例子中有两个操作，一个临时数组（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

Numba：手动循环比使用 numpy 数组的 a += c * b 更快？

Numba: Manual looping faster than a += c * b with numpy arrays?

python

numba