diff --git a/suitesparse_graphblas/__init__.py b/suitesparse_graphblas/__init__.py index c61a6b9..0408426 100644 --- a/suitesparse_graphblas/__init__.py +++ b/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"): +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,12 +76,25 @@ def initialize(*, blocking=False, memory_manager="numpy"): 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 : {'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 '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 (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": @@ -93,6 +106,20 @@ def initialize(*, blocking=False, memory_manager="numpy"): # See: https://github.com/GraphBLAS/python-suitesparse-graphblas/issues/40 for attr in dir(lib): getattr(lib, attr) + jit._get_suitesparse_original_configs() + 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() def libget(name): @@ -302,3 +329,5 @@ def __repr__(self): burble = burble() + +from . import jit # noqa: E402 isort:skip diff --git a/suitesparse_graphblas/jit.py b/suitesparse_graphblas/jit.py new file mode 100644 index 0000000..31d5eb1 --- /dev/null +++ b/suitesparse_graphblas/jit.py @@ -0,0 +1,156 @@ +import sysconfig + +from suitesparse_graphblas import _error_code_lookup, ffi, lib + +_SUITESPARSE_ORIG_CONFIG = {} + + +def _get_suitesparse_original_configs(): + """Save the initial configuration for the SuiteSparse:GraphBLAS JIT.""" + if _SUITESPARSE_ORIG_CONFIG: + raise RuntimeError("suitesparse defaults already gotten!") + val_ptr = ffi.new("int32_t*") + # GxB_JIT_C_CONTROL + info = lib.GxB_Global_Option_get_INT32(lib.GxB_JIT_C_CONTROL, val_ptr) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info]("Failed to get config for GxB_JIT_C_CONTROL.") + _SUITESPARSE_ORIG_CONFIG["C_CONTROL"] = val_ptr[0] + # GxB_JIT_USE_CMAKE + info = lib.GxB_Global_Option_get_INT32(lib.GxB_JIT_USE_CMAKE, val_ptr) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info]("Failed to get config for GxB_JIT_USE_CMAKE.") + _SUITESPARSE_ORIG_CONFIG["USE_CMAKE"] = bool(val_ptr[0]) + + val_ptr = ffi.new("char**") + get_config = lib.GxB_Global_Option_get_CHAR + for config_name in ( + "C_COMPILER_NAME", + "C_COMPILER_FLAGS", + "C_LINKER_FLAGS", + "C_LIBRARIES", + "C_CMAKE_LIBS", + "C_PREFACE", + "ERROR_LOG", + "CACHE_PATH", + ): + key_obj = getattr(lib, f"GxB_JIT_{config_name}") + info = get_config(key_obj, val_ptr) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info](f"Failed to get config for GxB_JIT_{config_name}.") + _SUITESPARSE_ORIG_CONFIG[config_name] = ffi.string(val_ptr[0]).decode() + + +def _set_defaults_common(**override): + set_config = lib.GxB_Global_Option_set_INT32 + # GxB_JIT_C_CONTROL (turn on) + info = set_config(lib.GxB_JIT_C_CONTROL, ffi.cast("int32_t", lib.GxB_JIT_ON)) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info]("Failed to set config for GxB_JIT_C_CONTROL") + # GxB_JIT_USE_CMAKE + info = set_config( + lib.GxB_JIT_USE_CMAKE, ffi.cast("int32_t", _SUITESPARSE_ORIG_CONFIG["USE_CMAKE"]) + ) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info]("Failed to set config for GxB_JIT_USE_CMAKE") + + set_config = lib.GxB_Global_Option_set_CHAR + for config_name in ( + "C_COMPILER_NAME", + "C_COMPILER_FLAGS", + "C_LINKER_FLAGS", + "C_LIBRARIES", + "C_CMAKE_LIBS", + "C_PREFACE", + "ERROR_LOG", + "CACHE_PATH", + ): + key_obj = getattr(lib, f"GxB_JIT_{config_name}") + if config_name in override: + val = override[config_name] + else: + val = _SUITESPARSE_ORIG_CONFIG[config_name] + val_obj = ffi.new("char[]", val.encode()) + info = set_config(key_obj, val_obj) + if info != lib.GrB_SUCCESS: + raise _error_code_lookup[info](f"Failed to set config for GxB_JIT_{config_name}") + + +def set_suitesparse_defaults(): + """Enable the JIT and set all JIT configs to the SuiteSparse:GraphBLAS defaults. + + Using the SuiteSparse:GraphBLAS JIT requires a C compiler. + + This sets values for: + - GxB_JIT_C_CONTROL + - GxB_JIT_USE_CMAKE + - GxB_JIT_C_COMPILER_NAME + - GxB_JIT_C_COMPILER_FLAGS + - GxB_JIT_C_LINKER_FLAGS + - GxB_JIT_C_LIBRARIES + - GxB_JIT_C_CMAKE_LIBS + - GxB_JIT_C_PREFACE + - GxB_JIT_ERROR_LOG + - GxB_JIT_CACHE_PATH + + See Also + -------- + set_python_defaults : set compiler config from Python's ``sysconfig`` module + """ + _set_defaults_common() + + +def set_python_defaults(): + """Enable the JIT and set configs with compiler configs from ``sysconfig`` module. + + Using the SuiteSparse:GraphBLAS JIT requires a C compiler. + + This uses ``sysconfig`` to set the values for: + - GxB_JIT_C_COMPILER_NAME + - GxB_JIT_C_COMPILER_FLAGS + - GxB_JIT_C_LIBRARIES + + and uses the SuiteSparse:GraphBLAS defaults for: + - GxB_JIT_USE_CMAKE + - GxB_JIT_C_LINKER_FLAGS + - GxB_JIT_C_CMAKE_LIBS + - GxB_JIT_C_PREFACE + - GxB_JIT_ERROR_LOG + - GxB_JIT_CACHE_PATH + + See Also + -------- + 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=cc, + C_COMPILER_FLAGS=f"{cflags} -I{include}", + C_LIBRARIES=libs, + ) + + +def disable(): + """Completely disable the SuiteSparse:GraphBLAS JIT. + + This sets GxB_JIT_C_CONTROL to GxB_JIT_OFF. + """ + set_config = lib.GxB_Global_Option_set_INT32 + 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")