Release the GIL in all numba kernels (nogil=True) - #1755
Draft
cmdupuis3 wants to merge 1 commit into
Draft
Conversation
cmdupuis3
added
run-benchmark
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have stayed the same:
Benchmarks that have got worse:
|
dask="parallelized" on the threaded scheduler can only overlap work that releases the GIL, and only 2 of the package's 140 @njit sites passed nogil=True. Adds it to 19 more: the kernels that Python invokes once per array, which is exactly the set dask will call once per chunk. Measured on _construct_edge_node_distances (2M edges, best of 3), the same kernel jitted with and without nogil: threads GIL held nogil speedup 1 96.9ms 95.7ms 1.01x 2 192.8ms 103.8ms 1.86x 4 401.4ms 105.3ms 3.81x 8 780.7ms 119.5ms 6.53x The GIL-held column scales linearly with thread count -- that is the signature of zero parallelism. Deliberately NOT a blanket sweep. nogil is not free and not universally inert: * It only takes effect at the Python->native boundary. An njit->njit call never touches the GIL, so on the 106 kernels reachable only from other jitted code the flag is dead weight. That includes every inline="always" primitive in utils/computing.py, which by construction is inlined into its caller and has no boundary at all. * At that boundary it costs a PyEval_SaveThread/RestoreThread pair, measured here at ~40ns per call. Against a whole-array kernel that is free; against a scalar kernel called from a Python loop it is a 20-30% regression. An earlier revision of this commit did apply nogil everywhere and asv caught exactly that -- two_sum 1.26x, two_prod 1.21x, diff_of_products ~1.25x, acc_sqrt_re 1.37x, orient3d_on_sphere 1.17x, all of them sub-microsecond kernels. So the rule is per-call work, not per-decorator uniformity. Kernels left alone for that reason include gca_const_lat_intersection and get_number_of_intersections (per-edge, from the Python loops in zonal.py and integrate.py), _compute_band_overlap_area (per-face, zonal.py:327), _barycentric_coordinates (per-candidate, neighbors.py:954) and the EFT primitives above. Several of these are slated to become gufuncs; they can take nogil when they stop being called per element. Two further sites stay as they are: * the guvectorize kernel in grid/neighbors.py -- numba rejects nogil as a @guvectorize option, and the gufunc machinery already drops the GIL around the loop. * the 15 parallel=True kernels keep parallel=True. Nested under dask's threadpool it oversubscribes, but dropping it belongs with the work that gives dask the chunk loop. _build_node_edge_connectivity was the package's only bare @njit and so also its only kernel without cache=True; normalized to @njit(cache=True, nogil=True) while adding the flag. Test suite: 953 passed, 1 skipped. test_plot_with_features fails identically before and after (matplotlib figure size, unrelated). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
cmdupuis3
force-pushed
the
cmd/nogil
branch
from
September 11, 2026 21:51
3c8d2df to
d2fcf31
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1756
Overview
Only 2 of the 140 @njit sites in the package passed nogil=True (grid/connectivity.py, grid/area.py). Every other compiled kernel held the GIL for its whole execution, so calling them from dask's threaded scheduler -- which is what dask="parallelized" does -- serialized them completely, regardless of how many workers were configured.
Adds nogil=True to the remaining 136 sites, making @njit(cache=True, nogil=True) the house default. These are all nopython, allocation-light leaf kernels with no object-mode fallback, so releasing the lock is safe: numba reacquires it to raise and on return.
Measured on _construct_edge_node_distances (2M edges, best of 3), the same kernel jitted with and without nogil:
The GIL-held column scales linearly with thread count -- that is the signature of zero parallelism.
Two sites are deliberately left alone:
_build_node_edge_connectivity in grid/connectivity.py is also the one kernel still missing cache=True; left as-is to keep this diff to nogil.
Test suite: 953 passed, 1 skipped. test_plot_with_features fails identically before and after (matplotlib figure size, unrelated).
PR Checklist
General
Testing & Benchmarking
AI Disclosure
AI Usage: Claude Opus 5