Numba: Manual looping faster than a += c * b with numpy arrays?
<p>I would like to do a 'daxpy' (add to a vector the scalar multiple of a second vector and assign the result to the first) with <code>numpy</code> using <code>numba</code>. Doing the following test, I noticed that writing the loop myself was much faster than doing <code>a += c * b</code>.</p>

<p>I was not expecting this. What is the reason for this behavior?</p>

<pre><code>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)
&gt;&gt;&gt; 100 loops, best of 10: 2.48 ms per loop
%timeit -n100 -r10 test2(x, o, c)
&gt;&gt;&gt; 100 loops, best of 10: 1.2 ms per loop
</code></pre>

<p>One thing to keep in mind is 'manual looping' in <code>numba</code> is very fast, essentially the same as the c-loop used by numpy operations.</p>

<p>In the first example there are two operations, a temporary array (<code>c * b</code>) is allocated / calculated, then that temporary array is added to <code>a</code>.  In the second example, both calculations are happening in the same loop with no intermediate result.</p>

<p>In theory, <code>numba</code> could fuse loops and optimize #1 to do the same as #2, but it doesn't seem to be doing it.  If you just want to optimize numpy ops, <code>numexpr</code> may also be worth a look as it was designed for exactly that - though probably won't do any better than the explicit fused loop.</p>

<pre><code>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
</code></pre>

python numba