Use wrapped versions of jit and guvectorize with custom defaults

Author: timothymillar

sgkit/accelerate.py

@@ -0,0 +1,31 @@
+import os
+from typing import Callable
+
+from numba import guvectorize, jit
+
+_DISABLE_CACHE = os.environ.get("SGKIT_DISABLE_NUMBA_CACHE", "0")
+
+try:
+    CACHE_NUMBA = {"0": True, "1": False}[_DISABLE_CACHE]
+except KeyError as e:  # pragma: no cover
+    raise KeyError(
+        "Environment variable 'SGKIT_DISABLE_NUMBA_CACHE' must be '0' or '1'"
+    ) from e
+
+
+DEFAULT_NUMBA_ARGS = {
+    "nopython": True,
+    "cache": CACHE_NUMBA,
+}
+
+
+def numba_jit(*args, **kwargs) -> Callable:  # pragma: no cover
+    kwargs_ = DEFAULT_NUMBA_ARGS.copy()
+    kwargs_.update(kwargs)
+    return jit(*args, **kwargs_)
+
+
+def numba_guvectorize(*args, **kwargs) -> Callable:  # pragma: no cover
+    kwargs_ = DEFAULT_NUMBA_ARGS.copy()
+    kwargs_.update(kwargs)
+    return guvectorize(*args, **kwargs_)

sgkit/distance/metrics.py

@@ -9,8 +9,9 @@
 from typing import Any
 
 import numpy as np
-from numba import cuda, guvectorize, types
+from numba import cuda, types
 
+from sgkit.accelerate import numba_guvectorize
 from sgkit.typing import ArrayLike
 
 # The number of parameters for the map step of the respective distance metric
@@ -20,15 +21,13 @@
 }
 
 
-@guvectorize(  # type: ignore
+@numba_guvectorize(  # type: ignore
     [
         "void(float32[:], float32[:], float32[:], float32[:])",
         "void(float64[:], float64[:], float64[:], float64[:])",
         "void(int8[:], int8[:], int8[:], float64[:])",
     ],
     "(n),(n),(p)->(p)",
-    nopython=True,
-    cache=True,
 )
 def euclidean_map_cpu(
     x: ArrayLike, y: ArrayLike, _: ArrayLike, out: ArrayLike
@@ -78,15 +77,13 @@ def euclidean_reduce_cpu(v: ArrayLike) -> ArrayLike:  # pragma: no cover
     return out
 
 
-@guvectorize(  # type: ignore
+@numba_guvectorize(  # type: ignore
     [
         "void(float32[:], float32[:], float32[:], float32[:])",
         "void(float64[:], float64[:], float64[:], float64[:])",
         "void(int8[:], int8[:], int8[:], float64[:])",
     ],
     "(n),(n),(p)->(p)",
-    nopython=True,
-    cache=True,
 )
 def correlation_map_cpu(
     x: ArrayLike, y: ArrayLike, _: ArrayLike, out: ArrayLike
@@ -141,14 +138,12 @@ def correlation_map_cpu(
     )
 
 
-@guvectorize(  # type: ignore
+@numba_guvectorize(  # type: ignore
     [
         "void(float32[:, :], float32[:])",
         "void(float64[:, :], float64[:])",
     ],
     "(p, m)->()",
-    nopython=True,
-    cache=True,
 )
 def correlation_reduce_cpu(v: ArrayLike, out: ArrayLike) -> None:  # pragma: no cover
     """Corresponding "reduce" function for pearson correlation

sgkit/stats/aggregation.py

@@ -3,11 +3,11 @@
 import dask.array as da
 import numpy as np
 import xarray as xr
-from numba import guvectorize
 from typing_extensions import Literal
 from xarray import Dataset
 
 from sgkit import variables
+from sgkit.accelerate import numba_guvectorize
 from sgkit.stats.utils import cohort_sum
 from sgkit.typing import ArrayLike
 from sgkit.utils import (
@@ -19,16 +19,14 @@
 Dimension = Literal["samples", "variants"]
 
 
-@guvectorize(  # type: ignore
+@numba_guvectorize(  # type: ignore
     [
         "void(int8[:], uint8[:], uint8[:])",
         "void(int16[:], uint8[:], uint8[:])",
         "void(int32[:], uint8[:], uint8[:])",
         "void(int64[:], uint8[:], uint8[:])",
     ],
     "(k),(n)->(n)",
-    nopython=True,
-    cache=True,
 )
 def count_alleles(
     g: ArrayLike, _: ArrayLike, out: ArrayLike

sgkit/stats/conversion.py

@@ -1,21 +1,19 @@
 import dask.array as da
 import numpy as np
-from numba import guvectorize
 from xarray import Dataset
 
 from sgkit import variables
+from sgkit.accelerate import numba_guvectorize
 from sgkit.typing import ArrayLike
 from sgkit.utils import conditional_merge_datasets, create_dataset
 
 
-@guvectorize(  # type: ignore
+@numba_guvectorize(  # type: ignore
     [
         "void(float64[:], uint8[:], float64, int8[:])",
         "void(float32[:], uint8[:], float64, int8[:])",
     ],
     "(p),(k),()->(k)",
-    nopython=True,
-    cache=True,
 )
 def _convert_probability_to_call(
     gp: ArrayLike, _: ArrayLike, threshold: float, out: ArrayLike

sgkit/stats/hwe.py

@@ -2,10 +2,10 @@
 
 import dask.array as da
 import numpy as np
-from numba import njit
 from xarray import Dataset
 
 from sgkit import variables
+from sgkit.accelerate import numba_jit
 from sgkit.stats.aggregation import count_genotypes
 from sgkit.typing import NDArray
 from sgkit.utils import conditional_merge_datasets, create_dataset
@@ -100,7 +100,9 @@ def hardy_weinberg_p_value(obs_hets: int, obs_hom1: int, obs_hom2: int) -> float
 
 
 # Benchmarks show ~25% improvement w/ fastmath on large (~10M) counts
-hardy_weinberg_p_value_jit = njit(hardy_weinberg_p_value, fastmath=True, nogil=True)
+hardy_weinberg_p_value_jit = numba_jit(
+    hardy_weinberg_p_value, fastmath=True, nogil=True
+)
 
 
 def hardy_weinberg_p_value_vec(
@@ -118,7 +120,7 @@ def hardy_weinberg_p_value_vec(
     return p
 
 
-hardy_weinberg_p_value_vec_jit = njit(
+hardy_weinberg_p_value_vec_jit = numba_jit(
     hardy_weinberg_p_value_vec, fastmath=True, nogil=True
 )
 

sgkit/stats/ld.py

@@ -7,15 +7,15 @@
 import numpy as np
 import pandas as pd
 from dask.dataframe import DataFrame
-from numba import njit
 from xarray import Dataset
 
 from sgkit import variables
+from sgkit.accelerate import numba_jit
 from sgkit.typing import ArrayLike, DType
 from sgkit.window import _get_chunked_windows, _sizes_to_start_offsets, has_windows
 
 
-@njit(nogil=True, fastmath=False, cache=True)  # type: ignore
+@numba_jit(nogil=True, fastmath=False)  # type: ignore
 def rogers_huff_r_between(gn0: ArrayLike, gn1: ArrayLike) -> float:  # pragma: no cover
     """Rogers Huff *r*.
 
@@ -67,7 +67,7 @@ def rogers_huff_r_between(gn0: ArrayLike, gn1: ArrayLike) -> float:  # pragma: n
     return r
 
 
-@njit(nogil=True, fastmath=True, cache=True)  # type: ignore
+@numba_jit(nogil=True, fastmath=True)  # type: ignore
 def rogers_huff_r2_between(gn0: ArrayLike, gn1: ArrayLike) -> float:  # pragma: no cover
     return rogers_huff_r_between(gn0, gn1) ** 2  # type: ignore
 
@@ -202,7 +202,7 @@ def to_ld_df(x: ArrayLike, chunk_index: int) -> DataFrame:
     )
 
 
-@njit(nogil=True, cache=True)  # type: ignore
+@numba_jit(nogil=True)  # type: ignore
 def _ld_matrix_jit(
     x: ArrayLike,
     chunk_window_starts: ArrayLike,
@@ -302,7 +302,7 @@ def _ld_matrix(
     return df
 
 
-@njit(nogil=True, cache=True)  # type: ignore
+@numba_jit(nogil=True)  # type: ignore
 def _maximal_independent_set_jit(
     idi: ArrayLike, idj: ArrayLike, cmp: ArrayLike
 ) -> List[int]:  # pragma: no cover