Numba is not enhancing the performance

Question

I am testing numba performance on some function that takes a numpy array, and compare:

import numpy as np
from numba import jit, vectorize, float64
import time
from numba.core.errors import NumbaWarning
import warnings

warnings.simplefilter('ignore', category=NumbaWarning)

@jit(nopython=True, boundscheck=False) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a):     # Function is compiled to machine code when called the first time
    trace = 0.0
    for i in range(a.shape[0]):   # Numba likes loops
        trace += np.tanh(a[i, i]) # Numba likes NumPy functions
    return a + trace              # Numba likes NumPy broadcasting
   
class Main(object):
    def __init__(self) -> None:
        super().__init__()
        self.mat     = np.arange(100000000, dtype=np.float64).reshape(10000, 10000)

    def my_run(self):
        st = time.time()
        trace = 0.0
        for i in range(self.mat.shape[0]):   
            trace += np.tanh(self.mat[i, i]) 
        res = self.mat + trace
        print('Python Diration: ', time.time() - st)
        return res                           
    
    def jit_run(self):
        st = time.time()
        res = go_fast(self.mat)
        print('Jit Diration: ', time.time() - st)
        return res
        
obj = Main()
x1 = obj.my_run()
x2 = obj.jit_run()

The output is:

Python Diration:  0.2164750099182129
Jit Diration:  0.5367801189422607

How can I obtain an enhance version of this example ?

I cannot reproduce the problem on my machine when the compilation time of Numba is excluded (ie. ignoring the first run of the JIT function): both takes about 0.1 s. — Jérôme Richard, Dec 22 '21 at 23:10
That's the answer. Testing on my machine, I get `Python duration: 0.23` then `Jit duration: 0.79 0.20 0.20 0.20 0.20 ...`. — Guimoute, Dec 22 '21 at 23:10
Yes, I am timing the first run. Is there any way to initialize the jit function without running it ? since I will be running the function once in normal work — zezo, Dec 22 '21 at 23:22

Jérôme Richard · Accepted Answer · 2021-12-22T23:52:55.203

1

The slower execution time of the Numba implementation is due to the compilation time since Numba compile the function at the time it is used (only the first time unless the type of the argument change). It does that because it cannot know the type of the arguments before the function is called. Hopefully, you can specify the argument type to Numba so it can compile the function directly (when the decorator function is executed). Here is the resulting code:

@njit('float64[:,:](float64[:,:])')
def go_fast(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

Note that njit is a shortcut for jit+nopython=True and that boundscheck is already set to False by default (see the doc).

On my machine this result in the same execution time for both Numpy and Numba. Indeed, the execution time is not bounded by the computation of the tanh function. It is bounded by the expression a + trace (for both Numba and Numpy). The same execution time is expected since both implement this the same way: they create a temporary new array to perform the addition. Creating a new temporary array is expensive because of page faults and the use of the RAM (a is fully read from the RAM and the temporary array is fully stored in RAM). If you want a faster computation, then you need to perform the operation in-place (this prevent page faults and expensive cache-line write allocations on x86 platforms).

edited Dec 22 '21 at 23:52

answered Dec 22 '21 at 23:43

Jérôme Richard

41,678
6
29
59

Thanks for the explanation. I have three questions, 1) how you read this expression 'float64[:,:](float64[:,:])' (what does it mean), 2) If I have a second input (of type numba.type.pyobject), how can I add it to this expression, and 3) how can I add the type of the output to this expression (suppose it is an array of same type float64) ? – zezo Dec 23 '21 at 00:18
1) the syntax is `ReturnType(Arg1_Type, Arg2_Type, ...)`. The type of a 1D array is `ItemType[:]`. For a 2D array it is `ItemType[:, :]` and so on. Here the type of the items is a 64-bit float (hence `float64`). See the [doc](https://numba.readthedocs.io/en/stable/reference/types.html?highlight=types#basic-types) for more information. 2) you cannot use CPython objects in Numba nopython jitted code because not using slow CPython dynamic types is what makes Numba fast (it uses static types and no GC). Otherwise you can use `pyobject` with `nopython=False` (but I this is not very useful). – Jérôme Richard Dec 23 '21 at 00:31
For the question 3, I think the answer 1 already gives the answer: the first parameter before the parenthesis is the return type. – Jérôme Richard Dec 23 '21 at 00:32
I have used the syntax as following for Arg1(numpy array) and Arg2(pyobject): @jit(('float64[:,:](float64[:,:])', 'pyobject'), nopython=False).....-> gives me this error [TypeError: invalid type in signature: expected a type instance, got 'float64[:,:](float64[:,:])'] – zezo Dec 23 '21 at 00:47
I am not sure to understand what you exactly want but I think you need that: `@njit('float64[:,:](float64[:,:], pyobject)')`. – Jérôme Richard Dec 23 '21 at 00:52
Got it, Thanks a lot – zezo Dec 23 '21 at 01:04

Numba is not enhancing the performance

1 Answers1

Linked