convolve_2d numpy kernel uses prange without parallel=True (runs serial, ~7-10x slower)

[brendancol]

Summary

_convolve_2d_numpy in xrspatial/convolution.py is decorated with @jit(nopython=True, nogil=True) and iterates its two outer loops with numba.prange. Numba only parallelizes prange when parallel=True is set on the decorator. Without it, prange degrades to an ordinary serial range, so the kernel runs on a single core even though the code reads as parallel. This is the primary CPU convolution path, used by both the numpy backend and, per chunk, by the dask+numpy backend.

Impact

On a 20-core host, a 2000x2000 float64 raster with a 15x15 kernel:

• as shipped (serial prange): ~356 ms
• parallel=True: ~37 ms (raw kernel), ~53 ms end-to-end through convolve_2d

Roughly a 7-10x wall-clock difference, results identical.

Reproduction

import numpy as np, time
from numba import jit, prange

def make(parallel):
    @jit(nopython=True, nogil=True, parallel=parallel)
    def conv(data, kernel):
        nx, ny = data.shape
        nkx, nky = kernel.shape
        wkx, wky = nkx // 2, nky // 2
        out = np.empty_like(data); out[:] = np.nan
        for i in prange(wkx, nx - wkx):
            for j in prange(wky, ny - wky):
                num = 0.0
                for ii in range(max(i - wkx, 0), min(i + wkx + 1, nx)):
                    iii = wkx + ii - i
                    for jj in range(max(j - wky, 0), min(j + wky + 1, ny)):
                        num += kernel[iii, wky + jj - j] * data[ii, jj]
                out[i, j] = num
        return out
    return conv

data = np.random.rand(2000, 2000)
kernel = np.ones((15, 15))
for label, parallel in [("serial (as shipped)", False), ("parallel=True", True)]:
    fn = make(parallel); fn(data[:50, :50], kernel)  # compile
    ts = [ (lambda: (t := time.perf_counter(), fn(data, kernel), time.perf_counter() - t)[-1])() for _ in range(3)]
    print(label, f"{min(ts)*1000:.0f} ms")

Observed:

serial (as shipped) 356 ms
parallel=True 37 ms

Fix

Add parallel=True to the decorator. Because the same kernel is launched per chunk under dask's threaded scheduler, and numba's default workqueue threading layer is not threadsafe across host threads (SIGABRT on macOS, see #3141), the launch must be serialized behind a module-level threading.Lock, matching the pattern already used in terrain.py. A single numpy call takes the lock uncontended and still runs across all cores; concurrent dask chunk calls run one at a time, each internally parallel.

Backends

• numpy: directly affected, gains the speedup.
• dask+numpy: each chunk gains the speedup, serialized behind the lock.
• cupy / dask+cupy: unaffected (separate CUDA kernel).

OOM verdict: SAFE (peak memory scales with chunk + overlap depth, not the full array; dask graph is ~20 tasks per chunk). Bottleneck: compute-bound.