From 424dd76bb858ce6ddbbc88eef9621341d6289a8c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 17 Dec 2023 08:23:56 -0600 Subject: [PATCH 1/3] Enable the SuiteSparse:GraphBLAS JIT --- suitesparse_graphblas/__init__.py | 23 ++++- suitesparse_graphblas/jit.py | 135 ++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 suitesparse_graphblas/jit.py diff --git a/suitesparse_graphblas/__init__.py b/suitesparse_graphblas/__init__.py index c61a6b9..aea18a9 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="python"): """Initialize GraphBLAS via GrB_init or GxB_init. This must be called before any other GraphBLAS functions are called. @@ -76,12 +76,24 @@ 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 : {'python', 'suitesparse', '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. + 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. Default is 'python'. 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}" + ) blocking = lib.GrB_BLOCKING if blocking else lib.GrB_NONBLOCKING memory_manager = memory_manager.lower() if memory_manager == "numpy": @@ -93,6 +105,13 @@ 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 == "python": + jit.set_python_defaults() + elif jit_config == "suitesparse": + jit.set_suitesparse_defaults() + else: + jit.disable() def libget(name): @@ -302,3 +321,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..c72ae0c --- /dev/null +++ b/suitesparse_graphblas/jit.py @@ -0,0 +1,135 @@ +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. + + 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. + + 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. + """ + _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"), + ) + + +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") From 7397eff87fdc43af6b7a9c49ff121499f90225ef Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 17 Dec 2023 08:26:14 -0600 Subject: [PATCH 2/3] Add note that it requires a compiler. --- suitesparse_graphblas/__init__.py | 3 ++- suitesparse_graphblas/jit.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/suitesparse_graphblas/__init__.py b/suitesparse_graphblas/__init__.py index aea18a9..f6ad23a 100644 --- a/suitesparse_graphblas/__init__.py +++ b/suitesparse_graphblas/__init__.py @@ -83,7 +83,8 @@ def initialize(*, blocking=False, memory_manager="numpy", jit_config="python"): SuiteSparse:GraphBLAS was compiled. 'disable' completely turns off the JIT. 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. Default is 'python'. + if SuiteSparse:GraphBLAS was compiled locally. Using the + SuiteSparse:GraphBLAS JIT requires a C compiler. Default is 'python'. The global variable `suitesparse_graphblas.is_initialized` indicates whether GraphBLAS has been initialized. diff --git a/suitesparse_graphblas/jit.py b/suitesparse_graphblas/jit.py index c72ae0c..9a8f080 100644 --- a/suitesparse_graphblas/jit.py +++ b/suitesparse_graphblas/jit.py @@ -78,6 +78,8 @@ def _set_defaults_common(**override): 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 @@ -100,6 +102,8 @@ def set_suitesparse_defaults(): 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 From 372fcc34524bcbff4497de399c199f0b9e56f9c7 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 17 Dec 2023 09:01:44 -0600 Subject: [PATCH 3/3] Add "default" and "load" options --- suitesparse_graphblas/__init__.py | 23 +++++++++++++++-------- suitesparse_graphblas/jit.py | 23 ++++++++++++++++++++--- 2 files changed, 35 insertions(+), 11 deletions(-) diff --git a/suitesparse_graphblas/__init__.py b/suitesparse_graphblas/__init__.py index f6ad23a..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", 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() diff --git a/suitesparse_graphblas/jit.py b/suitesparse_graphblas/jit.py index 9a8f080..31d5eb1 100644 --- a/suitesparse_graphblas/jit.py +++ b/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")