Numba: “weapon of mass optimization” | by Alex Diaz | Towards Data Science

Numba is a Python compiler, specifically for numerical functions and allows you to accelerate your applications with high performance functions written directly in Python.
Numba generates machine code optimized from pure Python code using LLVM. With a couple of simple changes, our Python code (function-oriented) can be optimized “just-in-time” to obtain a performance similar to that of C, C++, without having to change the language.
You can find the entire code in my GitHub! :)
What is Numba?First steps: Compile for the CPUNumba for GPUConclusion
What is Numba?
First steps: Compile for the CPU
Numba for GPU
Conclusion
Numba is a compiler that allows you to accelerate Python code (numerical functions) for both CPU and GPU:
Function Compiler: Numba compiles Python functions, not whole applications or parts of it. Basically, Numba is another Python module to improve the performance of our functions.Just-in-time: (Dynamic translation) Numba translates the bytecode (intermediate code more abstract than the machine code) to machine code immediately before its execution to improve the execution speed.Numerically-focused: Numba is focused on numerical data, such as int, float, complex. For now, there are limitations to use it with string data.
Function Compiler: Numba compiles Python functions, not whole applications or parts of it. Basically, Numba is another Python module to improve the performance of our functions.
Just-in-time: (Dynamic translation) Numba translates the bytecode (intermediate code more abstract than the machine code) to machine code immediately before its execution to improve the execution speed.
Numerically-focused: Numba is focused on numerical data, such as int, float, complex. For now, there are limitations to use it with string data.
Numba is not the only way to program in CUDA, it is usually programmed in C/C ++ directly for it. But Numba allows you to program directly in Python and optimize it for both CPU and GPU with few changes in our code. In relation to Python, there are other alternatives such as pyCUDA, here is a comparison between them:
CUDA C/C++:
It is the most common and flexible way to program in CUDAAccelerate applications in C, C ++.
It is the most common and flexible way to program in CUDA
Accelerate applications in C, C ++.
pyCUDA
It is the most efficient CUDA form for PythonIt requires programming C in our Python code and, in general, many code modifications.
It is the most efficient CUDA form for Python
It requires programming C in our Python code and, in general, many code modifications.
Numba
Less efficient than pyCUDAIt allows you to write your code in Python and optimize it with few modificationsIt also optimizes the Python code for CPU
Less efficient than pyCUDA
It allows you to write your code in Python and optimize it with few modifications
It also optimizes the Python code for CPU
The objectives of this talk are the following:
Use Numba to compile functions on the CPU
Understand how Numba works
Accelerate Numpy ufuncs in GPU
Write Kernels using Numba (Next tutorial)
Numba, apart from being able to speed up the functions in the GPU, can be used to optimize functions in the CPU. To do this, Python decorators (function modifiers) are used.
First of all, we are going to start to evaluate the function hypot to try how Numba works. We need to use the decorator @jit in our function.
>>> # Numba function>>> hypot(3.0, 4.0)5.0>>> # Python function>>> hypot.py_func(3.0, 4.0)5.0
The result is the same in Numba as in the Python function, since Numba saves the original implementation of the function in .py_func.
Naturally, it is important to measure the performance of our code, check if Numba really works well and observe the difference between Python implementation and Numba implementation. Furthermore, the library math already includes the hypot function, we can evaluate it also.
>>> # Python function>>> %timeit hypot.py_func(3.0, 4.0)The slowest run took 17.62 times longer than the fastest. This could mean that an intermediate result is being cached. 1000000 loops, best of 3: 260 ns per loop>>> # Numba function>>> %timeit hypot(3.0, 4.0)The slowest run took 33.89 times longer than the fastest. This could mean that an intermediate result is being cached. 1000000 loops, best of 3: 216 ns per loop>>> # math function>>> %timeit math.hypot(3.0, 4.0)The slowest run took 105.55 times longer than the fastest. This could mean that an intermediate result is being cached. 10000000 loops, best of 3: 133 ns per loop
The math.hypot function is even faster than Numba!! This is because Numba introduces a certain overhead to each function call, which is larger than the Python function call overhead, very fast functions (like the previous one) will be affected by this.
(However, if you call a Numba function from another, there is very little overhead, sometimes even zero if the compiler integrates the function into the other. In short, check if the functions are really accelerating with Numba).
numba.pydata.org
When we initialize the hypot function:
IR Intermediate Representations
Bytecode Analysis Intermediate code more abstract than machine code
LLVM Low Level Virtual Machine, infrastructure to develop compilers
NVVM It is an IR compiler based on LLVM, it is designed to represent GPU kernels
Each line of python is preceded by several lines of Numba IR code. The most useful is to see the type annotations that show us how Numba treats the variables, for example, in “pyobject” it is indicating that Numba does not know the np.sin function and that he should call it from Python. We can inspect the process for the hypot using .inspect_types().
>>> @jit>>> def foo_np(x):>>>     return np.sin(x)>>> foo_np(2)>>> foo_np.inspect_types()foo_np (int64,)----------------------------# File: <ipython-input-18-02574ac7ba04> # --- LINE 1 ---  # label 0  @jit  # --- LINE 2 ---   def foo_np(x):    # --- LINE 3 ---    #   x = arg(0, name=x)  :: int64   #   $0.1 = global(np: <module 'numpy' from '/usr/local/lib/python3.6/dist-packages/numpy/__init__.py'>)  :: Module(<module 'numpy' from '/usr/local/lib/python3.6/dist-packages/numpy/__init__.py'>)   #   $0.2 = getattr(value=$0.1, attr=sin)  :: Function(<ufunc 'sin'>)   #   del $0.1   #   $0.4 = call $0.2(x, func=$0.2, args=[Var(x, <ipython-input-18-02574ac7ba04> (3))], kws=(), vararg=None)  :: (int64,) -> float64   #   del x   #   del $0.2   #   $0.5 = cast(value=$0.4)  :: float64   #   del $0.4   #   return $0.5    return np.sin(x)   ==============================================
We are going to measure the performance of creating fractals using the Mandelbrot Set, and we will see how Numba helps us to improve our performance.
1 loop, best of 3: 4.62 s per loop<matplotlib.image.AxesImage at 0x7f986ce23780>
It takes about 4.62 seconds to generate a fractal using Mandelbrot set, now we are going to improve the performance using Numba, we just have to add the @jit decorator.
The slowest run took 4.17 times longer than the fastest. This could mean that an intermediate result is being cached. 1 loop, best of 3: 52.4 ms per loop
We can observe how we have achieved to reduce the time to build fractals from 4.62 seconds to 52.4 ms... and this has been done only adding a decorator!!
We have said that Numba only works for numeric functions, although Numba compiles and runs any Python code, there are some types of data that it cannot compile yet (such as dictionaries) and it also makes no sense to compile them.
>>> @jit>>> def dictionary(dict_test):>>>    return dict_test['house']dictionary({'house': 2, 'car': 35})2
But it has not failed !! We have said that Numba does not compile dictionaries... The point here is that Numba creates 2 functions, one of them in Python and another in Numba. So here we are seeing the python solution, we can verify this by doing nopython = True.
jit(nopython = True) is equivalent to njit
>>> @jit(nopython = True)>>> def dictionary(dict_test):>>>    return dict_test['house']dictionary({'house': 2, 'car': 35})-----------------------------------------TypingError                Traceback (most recent call last)<ipython-input-31-14d1c8683c01> in <module>()      3   return dict_test['house']      4 ----> 5 dictionary({'house': 245, 'car': 350})2 frames/usr/local/lib/python3.6/dist-packages/numba/six.py in reraise(tp, value, tb)    656             value = tp()    657         if value.__traceback__ is not tb:--> 658             raise value.with_traceback(tb)    659         raise value    660TypingError: Failed in nopython mode pipeline (step: nopython frontend)Internal error at <numba.typeinfer.ArgConstraint object at 0x7f986c1bed68>:--%<-----------------------------------------
There are two ways to program in GPU using Numba:
1. ufuncs/gufuncs__
2. CUDA Python Kernels (Next tutorial)
One of the main design features of the GPU is the ability to handle data in parallel, so the universal functions of numpy (ufunc) are an ideal candidate to implement them in GPU programming.
Note: ufunc are functions that perform the same operation on each element of a numpy array. For instance:
As we have said before, the ufunc functions are an ideal candidate to use them with GPU because of their parallelism. Therefore, Numba has the ability to create compiled ufunc functions without using C. To do this, we must use the decorator @vectorize.
Let’s start with an example using @vectorize to compile and optimize a CPU ufunc.
array([ 24, 343,  15,   9])
Instead of using the CPU to compile and execute the previous function, we are going to use CUDA in the GPU and for this we must use a “target attribute”. We will indicate what type each variable is (arguments and return value).
return_value_type(argument1_value_type, argument2_value_type, ...)
To do this, we will use the previous function that expects 2 int64 values and returns another int64 value. We will designate target = 'cuda' to be able to execute it in the GPU.
array([ 24, 343,  15,   9])
We can check the speed of running it on CPU or GPU:
>>> %timeit np.add(a, b)   # Numpy en CPUThe slowest run took 38.66 times longer than the fastest. This could mean that an intermediate result is being cached. 1000000 loops, best of 3: 511 ns per loop>>> %timeit add_ufunc_gpu(a, b)   # Numpy en GPUThe slowest run took 4.01 times longer than the fastest. This could mean that an intermediate result is being cached. 1000 loops, best of 3: 755 μs per loop
GPU is slower than the CPU!!! Quiet, this has an explanation ... But first let’s see what happens when we call that function...
When we execute this function, Numba produces:
Compile a CUDA kernel to execute the ufunc function in parallel over all the elements of the input arrayAssign the inputs and outputs to the GPU memoryCopy the input to the GPURun the CUDA KernelCopy the results back from the GPU to the CPUReturn the results as a numpy array
Compile a CUDA kernel to execute the ufunc function in parallel over all the elements of the input array
Assign the inputs and outputs to the GPU memory
Copy the input to the GPU
Run the CUDA Kernel
Copy the results back from the GPU to the CPU
Return the results as a numpy array
As compared to the implementation in C, Numba allows you to perform these types of tasks in a more concise way.
Why is the GPU slower than the CPU?
Our inputs are too small: The GPU achieves better performance using parallelism operating on thousands of values ​​at once. Our inputs are 4 or 64 dimensions, we need larger arrays to keep the GPU occupied.
Very simple calculation: Sending the calculation to the GPU requires a lot of “effort” compared to calling the CPU function. If our function does not require excessive mathematical computation (this is usually called * arithmetic intensity *), then the GPU could take longer than in the CPU.
Numba copies the data to the GPU.
The variable types of our inputs are larger than necessary: ​​Our example uses int64 and we probably do not need them. Practically, in CPU, 32 and 64-bits have the same computing speed, but in GPU, 64-bits have a slight increase in computational speed (it can reach up to 24x slower respectively to 32-bits). So it is important to keep this in mind when executing our function in the GPU.
With this in mind, we will try to apply what has been learned in the previous points to see if it is really faster to operate on the GPU than the CPU. We are going to calculate a density function, which is a slightly more complex operation with larger arrays.
Let’s calculate the value of a Gaussian density function in x given the mean and sigma:
>>> %timeit norm_pdf.pdf(x, loc=mean, scale=sigma)  # CPU function10 loops, best of 3: 60.8 ms per loop>>> %timeit gaussian_dens_gpu(x, mean, sigma) # GPU function100 loops, best of 3: 6.88 ms per loop
We can even define our function to be executed in the CPU using Numba.
>>> %timeit gaussian_dens_cpu(x, mean, sigma) # CPU10 loops, best of 3: 23.6 ms per loop
It is even faster than the function written in Python, but slower than the one executed in the GPU.
Unfortunately, there are several functions that do not fall within the scope of the definition of ufunc, therefore, to execute functions in the GPU that do not meet that requirement we use cuda.jit. We can use "device functions" that run on our GPU.
Note: “devices functions” are functions that can only be called from a kernel or from another “device” function.
>>> %timeit polar_distance(rho1, theta1, rho2, theta2)The slowest run took 23.16 times longer than the fastest. This could mean that an intermediate result is being cached. 1 loop, best of 3: 10.2 ms per loop
To sum up, Numba is a Python compiler, specifically for numerical functions and allows you to accelerate your applications with high performance functions written directly in Python.
It is a stable tool, which allows you to optimize code oriented to operate with arrays. Thanks to its ease of use (just a decorator!!) provides us with a very powerful tool to improve the performance of our code.
Suggestions and reviews are welcome. Follow me and thank you for reading! :)