Parallelize convolve_2d numpy kernel (#3615)

[brendancol]

Fixes #3615.

Problem

_convolve_2d_numpy iterates its two outer loops with numba.prange but was decorated @jit(nopython=True, nogil=True) with no parallel=True. Numba only parallelizes prange when parallel=True is set, so as shipped the loop ran serial on a single core. This is the main CPU convolution path, used by the numpy backend and per chunk by the dask+numpy backend.

Fix

• Add parallel=True to the decorator.
• Serialize the kernel launch behind a module-level threading.Lock. Dask calls the kernel per chunk under its threaded scheduler, and numba's default workqueue threading layer is not threadsafe across host threads (SIGABRT on macOS). This is the same hazard and fix already applied to terrain.py (Streaming reproject thread pool aborts the process when numba parallel kernels run concurrently #3141). 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.

The cupy and dask+cupy paths use a separate CUDA kernel and are untouched.

Measurement

20-core host, 2000x2000 float64 raster, 15x15 kernel, end-to-end through convolve_2d:

• before: ~356 ms
• after: ~53 ms

Results are identical (verified np.allclose(..., equal_nan=True)).

Verification

• xrspatial/tests/test_convolution.py: 6 passed.
• numpy vs dask+numpy parity across all four boundary modes (nan, nearest, reflect, wrap): match.
• cupy vs numpy parity: match.
• Concurrent dask threaded scheduler (8 workers): runs clean, no crash or deadlock.

Benchmark

Adds benchmarks/benchmarks/convolution.py with Convolve2d (numpy/cupy/dask across sizes and kernel sizes) and CircleKernel, so the serial regression cannot silently return.

Performance context

• OOM verdict: SAFE. Peak dask memory scales with chunk size plus overlap depth, not the full array (~20 tasks per chunk).
• Bottleneck: compute-bound.
• Affected backends: numpy (direct speedup), dask+numpy (per-chunk speedup, serialized behind the lock).

[brendancol]

PR Review: Parallelize convolve_2d numpy kernel (#3615)

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

• benchmarks/benchmarks/convolution.py:20-21 — the dask case of time_convolve_2d never runs the computation. convolve_2d(self.agg.data, self.kernel) on a dask-backed input returns a lazy array, so ASV times graph construction, not the kernel. Since this PR's whole point is dask throughput, the dask numbers won't move. Most dask benchmarks in this directory (watershed.py, twi.py, flood.py, interpolate.py, flow_length.py) call result.data.compute() (or .compute() on the returned array) in the dask branch to force execution. Consider doing the same here.

Nits (optional improvements)

• benchmarks/benchmarks/convolution.py:20-21 — the cupy case doesn't synchronize the device, so the timing may end before the kernel finishes. This matches how the other cupy benchmarks in the repo are written, so it's consistent, just worth knowing the cupy number is a launch time rather than an execution time.

What looks good

• The parallel=True change is guarded against the #3141 SIGABRT hazard exactly the way the terrain/reproject kernels were: _PARALLEL_KERNEL_LOCK serializes every kernel launch, and the dask path (the only concurrent caller, via map_overlap under the threaded scheduler) goes through _convolve_2d_numpy_locked. I checked every call site — there is no path that reaches _convolve_2d_numpy without holding the lock. The comment on the lock spells out the reasoning for the next reader.
• No data race inside the kernel: the outer prange writes disjoint out[i, j], the num accumulator is loop-private, and out[:] = np.nan runs before the parallel region, so the boundary cells stay NaN. Numerical results are unchanged — only the outer loop is parallelized, inner accumulation order is untouched.
• A benchmark now exists for convolve_2d and circle_kernel, which the module was missing.

Checklist

• Algorithm matches reference/paper (unchanged; only loop parallelism added)
• All implemented backends produce consistent results (numpy path unchanged numerically)
• NaN handling is correct (edge cells stay NaN)
• Dask chunk boundaries handled correctly (locked wrapper, depth unchanged)
• No premature materialization or unnecessary copies
• [~] Benchmark exists (added; dask case should force compute — see Suggestions)
• No macOS concurrent-kernel SIGABRT hazard (#3141) — serialized behind a lock
• Comments present and accurate

[brendancol]

PR Review (follow-up): Parallelize convolve_2d numpy kernel (#3615)

Follow-up on commit ce8b097c, which addresses the one Suggestion from the first review.

Resolved

• The dask case of time_convolve_2d now forces the compute (benchmarks/benchmarks/convolution.py:22-23), so asv times the kernel rather than graph construction. The guard is hasattr(result, 'compute') on the returned array — convolve_2d returns the bare backend array, so numpy (ndarray) and cupy (cupy.ndarray) skip it and only dask computes. I ran all three backend paths (this box has CUDA): numpy, cupy, and dask each construct and time without error, and the dask path materializes.

Merge with main

• Merged origin/main (16 commits, including #3613's isort/unused-import cleanup on convolution.py). The auto-merge kept both the style fixes and this branch's threading import + parallel=True kernel. flake8 is clean on convolution.py and the benchmark, and test_convolution.py passes (6/6).

Still open

Nothing blocking. The cupy-doesn't-synchronize nit from the first review stands and remains consistent with the other cupy benchmarks in the repo, so I'd leave it.