Switching between parallel / serial mode

Numba Community Support
4

The solution I ended up using is the following, which allows the user to switch between parallel and serial mode, by passing a boolean argument to the wrapper-function, which then selects either the parallel or serial jitted function. Here is a working example:

import numpy as np
from numba import njit, prange

def foo(x, parallel=True):
    # Do something ...

    if parallel:
        # Parallel execution.
        y = bar_parallel(x)
    else:
        # Serial execution.
        y = bar_serial(x)

    # Do something ...

    return y

# Jit parallel version.
@njit(parallel=True)
def bar_parallel(x):
    n = len(x)
    y = np.zeros(n)
    
    # Parallel loop.
    for i in prange(n):
        
        # Inner-loop.
        for j in range(i, i + 100):
            # Some "heavy" computation.
            y[i] += np.cos(j + x[i])

    return y

# Jit serial version.
# This takes the underlying Python fuction from the function bar
# and wraps it in Numba Jit again.
bar_serial = njit(bar_parallel.py_func, parallel=False)

# Test array.
x = np.arange(10000)

# Run the underlying function `bar` in parallel mode.
foo(x=x, parallel=True)

# Run the underlying function `bar` in serial mode.
foo(x=x, parallel=False)

This is a simple and elegant solution that works well.

2 Likes