4. Integration with Numba
PyCOMPSs can also be used with Numba. Numba (http://numba.pydata.org/) is an Open Source JIT compiler for Python which provides a set of decorators and functionalities to translate Python functios to optimized machine code.
4.1. Basic usage
PyCOMPSs’ tasks can be decorated with Numba’s
@jit
/@njit
decorator (with the appropiate
parameters) just below the @task decorator in order to apply
Numba to that task.
from pycompss.api.task import task # Import @task decorator
from numba import jit
@task(returns=1)
@jit()
def numba_func(a, b):
...
The task will be optimized by Numba within the worker node, enabling COMPSs to use the most efficient implementation of the task (and exploiting the compilation cache – any task that has already been compiled does not need to be recompiled in subsequent invocations).
4.2. Advanced usage
PyCOMPSs can be also used in conjuntion with the Numba’s
@vectorize
, @guvectorize
, @stencil
and @cfunc
.
But since these decorators do not preserve the original argument specification
of the original function, their usage is done through the numba parameter
withih the @task
decorator.
The numba parameter accepts:
- Boolean:
True: Applies jit to the function.
- Dictionary{k, v}:
Applies jit with the dictionary parameters to the function (allows to specify specific jit parameters (e.g.*nopython=True*)).
- String:
“jit”: Applies jit to the function.
“njit”: Applies jit with nopython=True to the function.
“generated_jit”: Applies generated_jit to the function.
“vectorize”: Applies vectorize to the function. Needs some extra flags in the @task decorator:
numba_signature: String with the vectorize signature.
“guvectorize”: Applies guvectorize to the function. Needs some extra flags in the @task decorator:
numba_signature: String with the guvectorize signature.
numba_declaration: String with the guvectorize declaration.
“stencil”: Applies stencil to the function.
“cfunc”: Applies cfunc to the function. Needs some extra flags in the @task decorator:
numba_signature: String with the cfunc signature.
Moreover, the @task decorator also allows to define specific flags for the jit, njit, generated_jit, vectorize, guvectorize and cfunc functionalities with the numba_flags hint. This hint is used to declare a dictionary with the flags expected to use with these numba functionalities. The default flag included by PyCOMPSs is the cache=True in order to exploit the function caching of Numba across tasks.
For example, to apply Numba jit to a task:
from pycompss.api.task import task
@task(numba='jit') # Aternatively: @task(numba=True)
def jit_func(a, b):
...
And if the developer wants to use specific flags with jit (e.g. parallel=True), the numba_flags must be defined with a dictionary where the key is the numba flag name, and the value, the numba flag value to use):
from pycompss.api.task import task
@task(numba='jit', numba_flags={'parallel':True})
def jit_func(a, b):
...
Other Numba’s functionalities require the specification of the function signature and declaration. In the next example a task that will use the vectorize with three parameters and a specific flag to target the cpu is shown:
from pycompss.api.task import task
@task(returns=1,
numba='vectorize',
numba_signature=['float32(float32, float32, float32)'],
numba_flags={'target':'cpu'})
def vectorize_task(a, b, c):
return a * b * c
Details about numba and the specification of the signature, declaration and flags can be found in the Numba’s webpage (http://numba.pydata.org/).