Automatic Parallelisation

Overview

  • Time: 10 min

  1. Learn how to automatically parallelize code in Numba.

  2. Learn how to parallelise loops in Numba.

Setting the parallel option for jit() enables Numba to automatically parallelize and optimize parts of a function, though it currently only works on CPUs. Instead of parallelizing each operation individually, which can lead to inefficiencies, Numba identifies and fuses adjacent parallelizable operations into kernels that run in parallel, improving performance.

 1@jit(nopython=True, parallel=True)
 2def reduction_with_parallel(n):
 3    shp = (13, 17)
 4    result1 = 2 * np.ones(shp, np.int_)
 5    tmp = 2 * np.ones_like(result1)
 6
 7    for i in prange(n):
 8        result1 *= tmp
 9
10    return result1

When parallel=True, Numba supports explicit parallel loops using prange instead of range. This allows you to specify that a loop can be parallelized, but you must ensure there are no dependencies between iteration, except for those allowed in reductions.

 1from numba import njit, prange
 2
 3
 4@njit(parallel=True)
 5def prange_test(A):
 6    s = 0
 7    # Without "parallel=True" in the jit-decorator
 8    # the prange statement is equivalent to range
 9    for i in prange(A.shape[0]):
10        s += A[i]
11    return s

Without parallel=True in the jit decorator, the prange statement behaves the same as range.

Key Points

  1. @jit(nopython=True, parallel=True) automatically parallelise functions.

  2. Numba supports explicit parallel by replacing range with prange.

  3. This functionality only works for CPU.