In [1]:
from __future__ import print_function
import numba
import numpy as np
import math
import llvm
import ctypes
print("numba version: %s \nNumPy version: %s\nllvm version: %s" % (numba.__version__,np.__version__, llvm.__version__))
NumPy provides a compact, typed container for homogenous arrays of data. This is ideal to store data homogeneous data in Python with little overhead. NumPy also provides a set of functions that allows manipulation of that data, as well as operating over it. There is a rich ecosystem around Numpy that results in fast manipulation of Numpy arrays, as long as this manipulation is done using pre-baked operations (that are typically vectorized). This operations are usually provided by extension modules and written in C, using the Numpy C API.
numba allows generating native code from Python functions just by adding decorators. This code is wrapped and directly callable from within Python.
There are many cases where you want to apply code to your NumPy data, and need that code to execute fast. You may get lucky and have the functions you want already written in the extensive NumPy ecosystem. Otherwise you will end with some code that is not that fast, but that you can improve execution time by writing code the "NumPy way". If it is not fast enough, you can write an extension module using the Numpy C API. Writing an extension module will take quite a bit of time, and forces you to a slow compile-install-test cycle.
Wouldn't it be great if you could just write code in Python that describes your function and execute it at speed similar to that of what you could achieve with the extension module, all without leaving the Python interpreter? numba allows that.
Numba is NumPy aware. This means:
NumPy arrays are understood by numba. By using the numba.typeof we can see that numba not only knows about the arrays themshelves, but also about its shape and underlying dtypes:
In [2]:
array = np.arange(2000, dtype=np.float_)
numba.typeof(array)
Out[2]:
In [3]:
numba.typeof(array.reshape((2,10,100)))
Out[3]:
From the point of view of numba, there are three factors that identify the array type:
The base type (dtype)
The number of dimensions (len(shape)). Note that for numba the arity of each dimension is not considered part of the type, only the dimension count.
The arrangement of the array. 'C' for C-like, 'F' for FORTRAN-like, 'A' for generic strided array.
It is easy to illustrate how the arity of an array is not part of the dtype in numba with the following samples:
In [4]:
numba.typeof(array.reshape((2,10,100))) == numba.typeof(array.reshape((4,10,50)))
Out[4]:
In [5]:
numba.typeof(array.reshape((2,10,100))) == numba.typeof(array.reshape((40,50)))
Out[5]:
In numba you can build the type specification by basing it on the base type for the array.
So if numba.float32 specifies a single precision floating point number:
In [6]:
numba.float32
Out[6]:
numba.float32[:] specifies an single dimensional array of single precision floating point numbers:
In [7]:
numba.float32[:]
Out[7]:
Adding dimensions is just a matter of tweaking the slice description passed:
In [8]:
numba.float32[:,:]
Out[8]:
In [9]:
numba.float32[:,:,:,:]
Out[9]:
In [10]:
numba.float32[:,::1]
Out[10]:
In [11]:
numba.float32[:,:,::1]
Out[11]:
column-major arrays (F-type) have elements in the first dimension packed together:
In [12]:
numba.float32[::1,:]
Out[12]:
In [13]:
numba.float32[::1,:,:]
Out[13]:
The use of any other dimension as consecutive is handled as a strided array:
In [14]:
numba.float32[:,::1,:]
Out[14]:
Note that the array arrangement does change the type, although numba will easily coerce a C or FORTRAN array into a strided one:
In [15]:
numba.float32[::1,:] == numba.float32[:,::1]
Out[15]:
In [16]:
numba.float32[::1,:] == numba.float32[:,:]
Out[16]:
In all cases, NumPy arrays are passed to numba functions by reference. This means that any change performed on the argument in the function will modify the contents of the original matrix. This behavior maps the usual NumPy semantics. So the array values passed as arguments to a numba functions can be considered as input/output arguments.
Indexing and slicing of NumPy arrays are handled natively by numba. This means that it is possible to index and slice a Numpy array in numba compiled code without relying on the Python runtime. In practice this means that numba code running on NumPy arrays will execute with a level of efficiency close to that of C.
Let's make a simple function that uses indexing. For example a really naive implementation of a sum:
In [17]:
def sum_all(A):
"""Naive sum of elements of an array... assumes one dimensional array of floats"""
acc = 0.0
for i in xrange(A.shape[0]):
acc += A[i]
return acc
In [18]:
sample_array = np.arange(10000.0)
The pure Python approach of this naive function is quite underwhelming speed-wise:
In [19]:
%timeit sum_all(sample_array)
If we relied on NumPy it would be much faster:
In [20]:
%timeit np.sum(sample_array)
But with numba the speed of that naive code is quite good:
In [21]:
sum_all_jit = numba.jit('float64(float64[:])')(sum_all)
In [22]:
%timeit sum_all_jit(sample_array)
This is in part possible because of the native support for indexing in numba. The function can be compiled in a nopython context, that makes it quite fast:
In [23]:
sum_all_jit = numba.jit('float64(float64[:])', nopython=True)(sum_all)
In NumPy there are universal functions(ufuncs) and generalized universal functions (gufuncs).
ufuncs are quite established and allows mapping of scalar operations over NumPy arrays. The resulting vectorized operation follow Numpy's broadcasting rules.
gufuncs are a generalization of ufuncs that allow vectorization of kernels that work over the inner dimensions of the arrays. In this context a ufunc would just be a gufunc where all the operands of its kernels have 0 dimensions (i.e. are scalars).
ufuncs and gufuncs are typically built using Numpy's C API. Numba offers the possibility to create ufuncs and gufuncs within the Python interpreter, using Python functions to describe the kernels. To access this functionality numba provides the vectorize decorator and the GUVectorize class.
vectorize is the decorator to be used to build ufuncs. Note that as of this writing, it is not in the numba namespace, but in numba.vectorize.
In [24]:
print(numba.vectorize.__doc__)
For example, let's write a sample ufunc that performs a lineal interpolation between A and B. The 'kernel' will look like this:
In [25]:
def lerp(A,B,factor):
"""interpolates A and B by factor"""
return (1-factor)*A + factor*B
In [26]:
lerp(0.0, 10.0, 0.3)
Out[26]:
Now let's do a ufunc for the floating point types. I will be using vectorize as a function, but remember that you could just add the decorator in the definition of the kernel itself.
In [27]:
lerp_ufunc = numba.vectorize(['float32(float32, float32, float32)', 'float64(float64, float64, float64)'])(lerp)
Now we can run our lerp with all of NumPy's niceties, like broadcasting of one operand (in this case the factor).
In [28]:
A = np.arange(0.0, 100000.0, 2.0)
B = np.arange(100000.0, 0.0, -2.0)
F = np.array([0.5] * 50000)
In [29]:
lerp_ufunc(A,B,0.5)
Out[29]:
It is also quite fast:
In [30]:
%timeit lerp_ufunc(A, B, 0.5)
In [31]:
%timeit lerp(A, B, 0.5)
Note that in this case the same original function can be used to generate the ufunc and to execute the equivalent NumPy vectorized version. When executing there will be differences in how the expression is evaluated.
When using NumPy the expression is evaluated one operation at a time, over the entire vector. Numba generated code will evaluate the full expression in one go, for each element. The numba approach approach avoids having temporal intermmediate arrays built, as well as avoiding revisiting operands that are being used more than once in a expression. This is useful with big arrays of data where there will be savings in process memory usage as well as better cache usage.
In [32]:
def sample_poly(x):
return x - x*x*x/6.0 + x*x*x*x*x/120.0
In [33]:
S = np.arange(0, np.pi, np.pi/36000000)
In [34]:
sample_poly_ufunc = numba.vectorize(['float32(float32)', 'float64(float64)'])(sample_poly)
In [35]:
%timeit sample_poly(S)
In [36]:
%timeit sample_poly_ufunc(S)
It is also worth noting that numba's vectorize provides similar convenience to that of NumPy's vectorize, but with performance similar to an ufunc.
For example, let's take the example in NumPy's vectorize documentation:
In [37]:
def myfunc(a, b):
"Return a-b if a>b, otherwise return a+b"
if a > b:
return a - b
else:
return a + b
In [38]:
myfunc_input = np.arange(100000.0)
In [39]:
numpy_vec_myfunc = np.vectorize(myfunc)
%timeit numpy_vec_myfunc(myfunc_input, 50000)
In [40]:
numba_vec_myfunc = numba.vectorize(['float64(float64, float64)'])(myfunc)
%timeit numba_vec_myfunc(myfunc_input, 50000)
In the same way the vectorize allows building NumPy's ufuncs from inside the Python interpreter just by writing the expression that forms the kernel; guvectorize allows building Numpy's gufuncs without the need of writing a C extension module.
In [41]:
print(numba.guvectorize.__doc__)
Generalized universal functions require a dimension signature for the kernel they implement. We call this the NumPy generalized-ufunc signature. Do not confuse this dimension signature with the type signature that numba requires.
The dimension signature describe the dimensions of the operands, as well as constraints to the values of those dimensions so that the function can work. For example, a matrix multiply gufunc will have a dimension signature like '(m,n), (n,p) -> (m,p)'. This means:
First operand has two dimensions (m,n).
Second operand has two dimensions (n,p).
Result has two dimensions (m,p).
The names of the dimensions are symbolic, and dimensions having the same name must match in arity (number of elements). So in our matrix multiply example the following constraints have to be met:
As you can see, the arity of the dimensions of the result can be infered from the source operands:
Result will have as many rows as rows has the first operand. Both are 'm'.
Result will have as many columns as columns has the second operand. Both are 'p'.
You can find more information about Numpy generalized-ufunc signature in NumPy's documentation.
When building a gufunc you start by writing the kernel function. You have to bear in mind which is the dimension signature and write the code to handle a single element. The function will take both, input arguments and results, as parameters. The result will be the last argument of the function. There shouldn't be any return value to the function, as the result should be placed directly in the last argument. The result of modifying an argument other than the result argument is undefined.
In [42]:
def matmulcore(A, B, C):
m, n = A.shape
n, p = B.shape
for i in range(m):
for j in range(p):
C[i, j] = 0
for k in range(n):
C[i, j] += A[i, k] * B[k, j]
Note how the m, n and p are extracted from the input arguments. The extraction of n is done twice to reinforce the notion that both are the same. That extraction is not really needed, as you could directly index inside the shape when defining the range.
To build a generalized-ufunc from the function is just a matter of using the guvectorize decorator. The interface to guvectorize is akin that of vectorize, but also requires the NumPy generalized-ufunc signature.
In [43]:
gu_matmul = numba.guvectorize(['float32[:,:], float32[:,:], float32[:,:]', 'float64[:,:], float64[:,:], float64[:,:]' ],
'(m,n),(n,p)->(n,p)')(matmulcore)
The result is a gufunc, that can be used as any othe gufunc in NumPy. Broadcasting and type promotion rules are those on NumPy.
In [44]:
matrix_ct = 10000
gu_test_A = np.arange(matrix_ct * 2 * 4, dtype=np.float32).reshape(matrix_ct, 2, 4)
gu_test_B = np.arange(matrix_ct * 4 * 5, dtype=np.float32).reshape(matrix_ct, 4, 5)
In [45]:
%timeit gu_matmul(gu_test_A, gu_test_B)
Some recap on the difference between vectorize and guvectorize:
vectorize generates ufuncs, guvectorize generates generalized-ufuncs
In both, type signatures for the arguments and return type are given in a list, but in vectorize function signatures are used to specify them, while on guvectorize a list of types is used instead, the resulting type being specified last.
When using guvectorize the NumPy generalized-ufunc signature needs to be supplied. This signature must be coherent with the type signatures.
Remember that with guvectorize the result is passed as the last argument to the kernel, while in vectorize the result is returned by the kernel.
There are some points to take into account when dealing with NumPy arrays inside numba compiled functions:
In numba generated code no range checking is performed when indexing. No range checking is performed as to allow generating code that performs better. So you need to be careful about the code as any indexing that goes out of range can cause a bad-access or a memory overwrite, potentially crashing the interpreter process.
In [46]:
arr = np.arange(16.0).reshape((2,8))
print(arr)
print(arr.strides)
As indexing in Python is 0-based, the following line will cause an exception error, as arr.shape[1] is 8, and the range for the column number is (0..7):
In [47]:
arr[0, arr.shape[1]] = 42.0
However, as numba doesn't have range checks, it will index anyways. As the index is out of bounds, and the array is in C order, the value will overflow into the next row. In this case, in the place reserved for element (1, 0).
In [48]:
@numba.jit("void(f8[:,:])")
def bad_access(array):
array[0, array.shape[1]] = 42.0
bad_access(arr)
print(arr)
In this sample case we where lucky, as the out-of-bounds access fell into the allocated range. Unchecked indexing can potentially cause illegal accesses and crash the process running the Python interpreter. However, it allows for code generation that produces faster code.