Search code examples
pythonpython-3.xperformancenumba

When numba is effective?


I know numba creates some overheads and in some situations (non-intensive computation) it become slower that pure python. But what I don't know is where to draw the line. Is it possible to use order of algorithm complexity to figure out where?

for example for adding two arrays (~O(n)) shorter that 5 in this code pure python is faster:

def sum_1(a,b):
    result = 0.0
    for i,j in zip(a,b):
            result += (i+j)
    return result

@numba.jit('float64[:](float64[:],float64[:])')
def sum_2(a,b):
    result = 0.0
    for i,j in zip(a,b):
            result += (i+j)
    return result

# try 100
a = np.linspace(1.0,2.0,5)
b = np.linspace(1.0,2.0,5)
print("pure python: ")
%timeit -o sum_1(a,b)
print("\n\n\n\npython + numba: ")
%timeit -o sum_2(a,b)

UPDADE: what I am looking for is a similar guideline like here:

"A general guideline is to choose different targets for different data sizes and algorithms. The “cpu” target works well for small data sizes (approx. less than 1KB) and low compute intensity algorithms. It has the least amount of overhead. The “parallel” target works well for medium data sizes (approx. less than 1MB). Threading adds a small delay. The “cuda” target works well for big data sizes (approx. greater than 1MB) and high compute intensity algorithms. Transfering memory to and from the GPU adds significant overhead."


Solution

  • It's hard to draw the line when numba becomes effective. However there are a few indicators when it might not be effective:

    • If you cannot use jit with nopython=True - whenever you cannot compile it in nopython mode you either try to compile too much or it won't be significantly faster.

    • If you don't use arrays - When you deal with lists or other types that you pass to the numba function (except from other numba functions), numba needs to copy these which incurs a significant overhead.

    • If there is already a NumPy or SciPy function that does it - even if numba can be significantly faster for short arrays it will almost always be as fast for longer arrays (also you might easily neglect some common edge cases that these would handle).

    There's also another reason why you might not want to use numba in cases where it's just "a bit" faster than other solutions: Numba functions have to be compiled, either ahead-of-time or when first called, in some situations the compilation will be much slower than your gain, even if you call it hundreds of times. Also the compilation times add up: numba is slow to import and compiling the numba functions also adds some overhead. It doesn't make sense to shave off a few milliseconds if the import overhead increased by 1-10 seconds.

    Also numba is complicated to install (without conda at least) so if you want to share your code then you have a really "heavy dependency".


    Your example is lacking a comparison with NumPy methods and a highly optimized version of pure Python. I added some more comparison functions and did a benchmark (using my library simple_benchmark):

    import numpy as np
    import numba as nb
    from itertools import chain
    
    def python_loop(a,b):
        result = 0.0
        for i,j in zip(a,b):
            result += (i+j)
        return result
    
    @nb.njit
    def numba_loop(a,b):
        result = 0.0
        for i,j in zip(a,b):
                result += (i+j)
        return result
    
    def numpy_methods(a, b):
        return a.sum() + b.sum()
    
    def python_sum(a, b):
        return sum(chain(a.tolist(), b.tolist()))
    
    from simple_benchmark import benchmark, MultiArgument
    
    arguments = {
        2**i: MultiArgument([np.zeros(2**i), np.zeros(2**i)])
        for i in range(2, 17)
    }
    b = benchmark([python_loop, numba_loop, numpy_methods, python_sum], arguments, warmups=[numba_loop])
    
    %matplotlib notebook
    b.plot()
    

    enter image description here

    Yes, the numba function is fastest for small arrays, however the NumPy solution will be slightly faster for longer arrays. The Python solutions are slower but the "faster" alternative is already significantly faster than your original proposed solution.

    In this case I would simply use the NumPy solution because it's short, readable and fast, except when you're dealing with lots of short arrays and call the function a lot of times - then the numba solution would be significantly better.