TopolearnParallel loops can multiply speed by the number of cores with a one-word change, if the loop is safe.
@njit(parallel=True) enables two things. First, Numba automatically parallelises supported array expressions and reductions. Second, you can replace range with numba.prange to mark a loop whose iterations run in parallel threads. The GIL is not involved: the compiled code runs natively on multiple cores.
Rules for a correct prange loop:
- Iterations must be independent. Writing to
out[i]from iterationiis safe. Writing to a shared element from several iterations is a race condition and gives wrong results. - Reductions into a scalar are recognised:
total += f(a[i])inside aprangeloop is turned into a per-thread partial sum combined at the end. This works for+=,*=and similar operators on scalars, not for arbitrary updates of array elements. - Only the outermost
prangein a nest is parallelised; inner ones run serially.
Control and inspect the threading:
numba.set_num_threads(n)andget_num_threads()change the thread count at run time, and theNUMBA_NUM_THREADSenvironment variable sets the maximum.f.parallel_diagnostics(level=1)reports which loops were parallelised and fused.- Threading layers (TBB, OpenMP, workqueue) can be selected with
NUMBA_THREADING_LAYER.
Parallelism has overhead: for small arrays or cheap loop bodies the single-threaded version is faster. Floating-point reductions may also give slightly different results from the serial sum, because the order of additions changes.
Example
import numpy as np
from numba import njit, prange, get_num_threads
@njit(parallel=True)
def pairwise_min_dist(p):
n = p.shape[0]
out = np.empty(n)
for i in prange(n): # rows processed in parallel
best = np.inf
for j in range(n):
if i != j:
best = min(best, ((p[i] - p[j]) ** 2).sum())
out[i] = np.sqrt(best) # each i writes only out[i]: safe
return out
@njit(parallel=True)
def total(a):
s = 0.0
for i in prange(a.size):
s += a[i] # recognised reduction
return s
pts = np.random.default_rng(0).random((2000, 3))
print(pairwise_min_dist(pts).shape, get_num_threads()) # (2000,) 8 (your core count)
print(total(np.ones(1_000_000))) # 1000000.0Remember
- parallel=True plus prange runs independent iterations on several cores
- Scalar reductions with += are handled; shared array writes are races
- Overhead and changed summation order are the trade-offs
Watch out for
- Updating a shared histogram bin from a prange loop, which silently loses counts.