Add "default" and "load" options

suitesparse_graphblas/__init__.py

@@ -60,7 +60,7 @@ def supports_complex():
     return hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64")
 
 
-def initialize(*, blocking=False, memory_manager="numpy", jit_config="python"):
+def initialize(*, blocking=False, memory_manager="numpy", jit_config="default"):
     """Initialize GraphBLAS via GrB_init or GxB_init.
 
     This must be called before any other GraphBLAS functions are called.
@@ -76,25 +76,25 @@ def initialize(*, blocking=False, memory_manager="numpy", jit_config="python"):
         allocators, which makes it safe to perform zero-copy to and from numpy,
         and allows Python to track memory usage via tracemalloc (if enabled).
         'c' uses the default allocators.  Default is 'numpy'.
-    jit_config : {'python', 'suitesparse', 'disable'}, optional
+    jit_config : {'default', 'python', 'suitesparse', 'load', 'disable'}, optional
         Choose how to configure the SuiteSparse:GraphBLAS JIT, or disable the JIT.
         'python' uses config from the builtin ``sysconfig`` module, which is what
         Python uses to enable compiling with C. 'suitesparse' uses config from when
         SuiteSparse:GraphBLAS was compiled. 'disable' completely turns off the JIT.
+        'load' allows pre-compiled JIT kernels to be loaded, but does not compile.
+        'default' tries to use 'python', but uses 'load' if no compiler was found.
         If you are using a distributed binary of SuiteSparse:GraphBLAS, then you
         probably want to use 'python', otherwise you may want to use 'suitesparse'
         if SuiteSparse:GraphBLAS was compiled locally. Using the
-        SuiteSparse:GraphBLAS JIT requires a C compiler. Default is 'python'.
+        SuiteSparse:GraphBLAS JIT requires a C compiler. Default is 'default'.
 
     The global variable `suitesparse_graphblas.is_initialized` indicates whether
     GraphBLAS has been initialized.
     """
     if is_initialized():
         raise RuntimeError("GraphBLAS is already initialized!  Unable to initialize again.")
-    if jit_config not in {"python", "suitesparse", "disable"}:
-        raise ValueError(
-            f"jit_config must be 'python', 'suitesparse', or 'disable'; got: {jit_config!r}"
-        )
+    if jit_config not in (allowed := {"default", "python", "suitesparse", "load", "disable"}):
+        raise ValueError(f"jit_config must be one of {allowed}; got: {jit_config!r}")
     blocking = lib.GrB_BLOCKING if blocking else lib.GrB_NONBLOCKING
     memory_manager = memory_manager.lower()
     if memory_manager == "numpy":
@@ -107,10 +107,17 @@ def initialize(*, blocking=False, memory_manager="numpy", jit_config="python"):
     for attr in dir(lib):
         getattr(lib, attr)
     jit._get_suitesparse_original_configs()
-    if jit_config == "python":
+    if jit_config == "default":
+        try:
+            jit.set_python_defaults()
+        except RuntimeError:
+            jit.load()
+    elif jit_config == "python":
         jit.set_python_defaults()
     elif jit_config == "suitesparse":
         jit.set_suitesparse_defaults()
+    elif jit_config == "load":
+        jit.load()
     else:
         jit.disable()
 

suitesparse_graphblas/jit.py

@@ -121,10 +121,16 @@ def set_python_defaults():
     --------
     set_suitesparse_defaults : reset JIT config with SuiteSparse:GraphBLAS defaults.
     """
+    cc = sysconfig.get_config_var("CC")
+    cflags = sysconfig.get_config_var("CFLAGS")
+    include = sysconfig.get_path("include")
+    libs = sysconfig.get_config_var("LIBS")
+    if cc is None or cflags is None or include is None or libs is None:
+        raise RuntimeError("C compiler configuration not found by `sysconfig` module")
     _set_defaults_common(
-        C_COMPILER_NAME=sysconfig.get_config_var("CC"),
-        C_COMPILER_FLAGS=sysconfig.get_config_var("CFLAGS") + f" -I{sysconfig.get_path('include')}",
-        C_LIBRARIES=sysconfig.get_config_var("LIBS"),
+        C_COMPILER_NAME=cc,
+        C_COMPILER_FLAGS=f"{cflags} -I{include}",
+        C_LIBRARIES=libs,
     )
 
 
@@ -137,3 +143,14 @@ def disable():
     info = set_config(lib.GxB_JIT_C_CONTROL, ffi.cast("int32_t", lib.GxB_JIT_OFF))
     if info != lib.GrB_SUCCESS:
         raise _error_code_lookup[info]("Failed to set config for GxB_JIT_C_CONTROL")
+
+
+def load():
+    """Allow the SuiteSparse:GraphBLAS JIT to load JIT kernels, but don't compile.
+
+    This sets GxB_JIT_C_CONTROL to GxB_JIT_LOAD.
+    """
+    set_config = lib.GxB_Global_Option_set_INT32
+    info = set_config(lib.GxB_JIT_C_CONTROL, ffi.cast("int32_t", lib.GxB_JIT_LOAD))
+    if info != lib.GrB_SUCCESS:
+        raise _error_code_lookup[info]("Failed to set config for GxB_JIT_C_CONTROL")