feat: compile njit functions with nogil=True

src/fast_array_utils/numba/__init__.py

@@ -49,14 +49,19 @@ def _is_on_unsafe_thread() -> bool:
 
 
 @overload
-def njit[**P, R](fn: Callable[P, R], /) -> Callable[P, R]: ...
+def njit[**P, R](fn: Callable[P, R], /, *, nogil: bool = True) -> Callable[P, R]: ...
 @overload
-def njit[**P, R]() -> Callable[[Callable[P, R]], Callable[P, R]]: ...
-def njit[**P, R](fn: Callable[P, R] | None = None, /) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+def njit[**P, R](*, nogil: bool = True) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+def njit[**P, R](fn: Callable[P, R] | None = None, /, *, nogil: bool = True) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
     """Jit-compile a function using numba.
 
     On call, this function dispatches to a parallel or serial numba function,
     depending on the current threading environment.
+
+    Parameters
+    ----------
+    nogil
+        Release the GIL while the compiled function runs (see :func:`numba.jit`).
     """
     # See https://github.com/numbagg/numbagg/pull/201/files#r1409374809
 
@@ -69,7 +74,9 @@ def decorator(f: Callable[P, R], /) -> Callable[P, R]:
 
         # use distinct names so numba doesn’t reuse the wrong version’s cache
         fns: dict[bool, Callable[P, R]] = {
-            parallel: numba.njit(_copy_function(f, __qualname__=f"{f.__qualname__}-{'parallel' if parallel else 'serial'}"), cache=True, parallel=parallel)
+            parallel: numba.njit(
+                _copy_function(f, __qualname__=f"{f.__qualname__}-{'parallel' if parallel else 'serial'}"), cache=True, parallel=parallel, nogil=nogil
+            )
             for parallel in (True, False)
         }
 

tests/test_numba.py

@@ -83,9 +83,10 @@ def _set_runtime(
         monkeypatch.setattr(probe, "LAYERS", layers)
 
 
-def _install_fake_njit(monkeypatch: pytest.MonkeyPatch, calls: list[bool]) -> None:
-    def fake_njit(_fn: object, /, *, cache: bool, parallel: bool) -> Callable[..., bool]:
+def _install_fake_njit(monkeypatch: pytest.MonkeyPatch, calls: list[bool], *, expected_nogil: bool = True) -> None:
+    def fake_njit(_fn: object, /, *, cache: bool, parallel: bool, nogil: bool) -> Callable[..., bool]:
         assert cache is True
+        assert nogil is expected_nogil
 
         def compiled(*_args: object, **_kwargs: object) -> bool:
             calls.append(parallel)
@@ -287,3 +288,8 @@ def test_serial_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
         result = wrapped(values)
 
     assert result == pytest.approx(np.sum(values))
+
+
+def test_njit_nogil(monkeypatch: pytest.MonkeyPatch) -> None:
+    _install_fake_njit(monkeypatch, [], expected_nogil=False)
+    fa_numba.njit(nogil=False)(_return_true)

tests/test_stats.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: MPL-2.0
 from __future__ import annotations
 
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 from importlib.metadata import version
 from importlib.util import find_spec
 from pathlib import Path
@@ -409,3 +411,16 @@ def test_stats_benchmark(
     @benchmark
     def call() -> None:
         func(arr, axis=axis)
+
+
+@pytest.mark.benchmark
+@pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Matrix | Flags.Dask | Flags.Disk | Flags.Gpu)
+def test_mean_var_threaded_benchmark(benchmark: BenchmarkFixture, array_type: ArrayType[CpuArray, None]) -> None:
+    """Calls from worker threads run serially, so they only overlap if numba releases the GIL."""
+    mats = [array_type.random((20_000, 5_000), density=0.1, dtype=np.float64) for _ in range(4)]
+    stats.mean_var(mats[0], axis=0)  # warmup: numba compile
+
+    @benchmark
+    def call() -> None:
+        with ThreadPoolExecutor(len(mats)) as pool:
+            list(pool.map(partial(stats.mean_var, axis=0), mats))