diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index 7c62a0e..fd5c1dd 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/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) } diff --git a/tests/test_numba.py b/tests/test_numba.py index 6a04bad..b71ce2f 100644 --- a/tests/test_numba.py +++ b/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) diff --git a/tests/test_stats.py b/tests/test_stats.py index 0860334..4d5afc6 100644 --- a/tests/test_stats.py +++ b/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))