diff --git a/docs/user_guide/init.rst b/docs/user_guide/init.rst index ffb6a3463..141e1f250 100644 --- a/docs/user_guide/init.rst +++ b/docs/user_guide/init.rst @@ -57,3 +57,74 @@ raised. However, if the arguments do not match, an error will be raised. There is no mechanism to change the initialization arguments after graphblas has been initialized the first time. + +.. _suitesparse_jit: + +The SuiteSparse:GraphBLAS JIT +----------------------------- + +SuiteSparse:GraphBLAS ships precompiled kernels for common operations. For the rest it +can compile a kernel with a C compiler (its JIT), or else run a generic kernel that is +often several times slower (``Matrix.reduce_rowwise`` with ``lor`` on ``BOOL`` or +``plus`` on ``UINT64`` ran 2.7 to 8 times slower on 10 million entries). The JIT control, +``gb.ss.config["jit_c_control"]``, is ``"run"`` by default: run only kernels already +loaded in the process, and compile nothing. UDT operators, and types and operators +registered from C source, raise this default to ``"on"`` when they need a kernel (see +:doc:`udt`), but leave a control you set as it is. To opt in for everything, set the +control to ``"on"`` (compile, load and run): + +.. code-block:: bash + + export GRAPHBLAS_JIT_C_CONTROL=on + +.. code-block:: python + + import graphblas as gb + + gb.config.set(jit_c_control="on") # before graphblas initializes + gb.ss.config["jit_c_control"] = "on" # or later, for the rest of the process + +``graphblas.config["jit_c_control"]`` is read once, when graphblas initializes. It can be +``"on"``, ``"load"`` (load and run cached kernels, never compile), ``"run"``, ``"pause"`` +(run no JIT kernels) or ``"off"`` (run none, and unload them). Set it with the +``GRAPHBLAS_JIT_C_CONTROL`` environment variable, with ``gb.config.set`` before graphblas +initializes, or in a YAML file in ``~/.config/graphblas/``. Left unset, it is the +default ``"run"``, which UDT work raises. With nothing configured, a ``"run"`` written +later to ``gb.ss.config`` is that same default, so to keep UDT work from compiling, set +``"run"`` in one of those three ways, or write ``"load"`` or ``"off"`` later. With +``"on"``, python-graphblas checks the C compiler by compiling a tiny type (once per cache +directory and compiler), and falls back to ``"run"``, with no error, where the JIT cannot +work: SuiteSparse built without the JIT (as in the PyPI wheels), SuiteSparse 7, no working +C compiler (in an activated conda environment, ``conda install c-compiler``), or a cache +directory it cannot write to. On Windows, SuiteSparse built with MSVC (as conda-forge +builds it) compiles JIT kernels only through cmake, and python-graphblas does not turn the +JIT on there yet, so a configured ``"on"`` becomes ``"run"``. + +With the JIT on: + +- **Do not share one cache directory between processes running at the same time** on + SuiteSparse 10.1.1 or later, which no longer locks the cache. They race when they + compile the same new kernel, and calls fail with ``JitError``, leave their outputs + invalid or cleared, return wrong values with no error, or hang + (`GraphBLAS issue #449 `_). + With multiprocessing, dask, pytest-xdist and the like, set + ``GRAPHBLAS_JIT_C_CONTROL=run`` in the environment the workers start in, so that none + of them compiles: UDT work raises only the default ``"run"``, not one you set. That + costs UDTs their compiled kernels: UDT operators call their Numba functions instead + (typically 2 to 3 times slower for elementwise operations), and registering types or + operators from C source fails. To keep the kernels, give each process its own cache + instead, by setting the ``GRAPHBLAS_CACHE_PATH`` environment variable before graphblas + initializes. The default will change once SuiteSparse publishes compiled kernels + atomically. +- The first use of each new kernel compiles it, typically in 0.2 to 1 second. Kernels are + cached on disk, in ``gb.ss.config["jit_cache_path"]`` (by default under + ``~/.SuiteSparse`` on Linux and macOS), so later runs load them instead. +- The cache is keyed only on the SuiteSparse:GraphBLAS version, so two builds of one + version (a conda build and a source build, say) can load each other's kernels, and get + wrong results or errors if their internals differ. Give each build its own cache path. + +If you hit a bug, freeze or crash that may involve the JIT, set +``GRAPHBLAS_JIT_C_CONTROL=run`` (which UDT work leaves alone) and delete the cache +directory, then try again. Please +`report it `_ either way. The +``suitesparse-vanilla`` backend leaves the control as the SuiteSparse build sets it. diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 2a0b86555..44af6a1f0 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -442,25 +442,32 @@ but never compiles a ``.dylib`` or ``.so``, and silently falls back to the cfunc path. The 2-3x JIT speedup is silently lost. python-graphblas repairs the compiler path when ``graphblas.ss`` is -imported, which the first access to ``gb.ss`` does. If -``jit_c_compiler_name`` doesn't exist on disk, it is replaced with one from -``$CONDA_PREFIX/bin/`` (or from ``sysconfig`` for pure-pip installs). The +imported, which the first access to ``gb.ss`` does (and at init, when the JIT +control is configured to ``'on'``). If ``jit_c_compiler_name`` doesn't exist on +disk, it is replaced with one from ``$CONDA_PREFIX/bin/`` (or from +``sysconfig`` for pure-pip installs). The import changes nothing else, so reading ``gb.ss.about`` does not change what later operations compute. -``jit_c_control`` is raised from the SS default ``'run'`` (run cached kernels -only; no compile, no load from disk) to ``'on'`` (compile, load, and run) by -each of these, since each asks SuiteSparse to compile C source: +``jit_c_control`` starts at ``graphblas.config["jit_c_control"]``, ``'run'`` by +default (run cached kernels only; no compile, no load from disk); see +:ref:`suitesparse_jit`. That default ``'run'`` is raised to ``'on'`` (compile, +load, and run) by each of these, since each asks SuiteSparse to compile C source: - a UDT operator that gets C source, such as the first ``binary.plus[udt]``; - ``gb.dtypes.ss.register_new``, which registers a type from a C typedef; - ``gb.binary.ss.register_new`` and its ``unary``, ``indexunary``, ``select``, and ``indexbinary`` siblings, which register an operator from a C definition. -A setting other than ``'run'`` is kept, including ``'off'``, ``'pause'``, and -``'load'`` (load and run cached kernels, never compile). SuiteSparse itself -sets ``'load'`` after a compile fails, so keeping it means a failing compile -is tried once, not again at every later request. +A control you set is kept: ``'off'``, ``'pause'``, ``'load'`` (load and run +cached kernels, never compile), and ``'run'`` too when it was set before init +(``GRAPHBLAS_JIT_C_CONTROL``, ``gb.config.set``, or a YAML config file). With +nothing configured, a ``'run'`` set later in ``gb.ss.config`` is the default +again. Under a ``'run'`` you set, UDT operators use their Numba functions, and +the registrations above fail with an error saying how to opt in: +``gb.ss.config["jit_c_control"] = "on"``, or ``None`` for the default. +SuiteSparse itself sets ``'load'`` after a compile fails, so keeping it means a +failing compile is tried once, not again at every later request. Call the helper manually to re-fix or verify:: diff --git a/graphblas/__init__.py b/graphblas/__init__.py index 86759f570..2de45dda9 100644 --- a/graphblas/__init__.py +++ b/graphblas/__init__.py @@ -137,6 +137,9 @@ def _init(backend_arg, blocking, automatic=False): # Already initialized with these parameters; nothing more to do return + if backend_arg == "suitesparse": + # Checked before anything is initialized, so a bad value changes nothing + jit_c_control = _jit_c_control_config() backend = backend_arg if backend in {"suitesparse", "suitesparse-vanilla"}: try: @@ -197,6 +200,35 @@ class Lib: core.ffi = ffi core.lib = lib core.NULL = ffi.NULL + if backend == "suitesparse": + from .core.ss.jit_config import _set_jit_c_control_at_init + + _set_jit_c_control_at_init(jit_c_control) + + +_JIT_C_CONTROLS = ("on", "load", "run", "pause", "off") + + +def _jit_c_control_config(): + """``config["jit_c_control"]``, the SuiteSparse JIT control set at init (see graphblas.yaml). + + None when it is not set, which is the default: ``'run'``, raised by UDT work. + """ + val = config.get("jit_c_control") + if val is None: + return None + # YAML 1.1 reads an unquoted on/off as a bool, so a config file saying + # ``jit_c_control: off`` arrives here as False. + if val is True or val is False: + return "on" if val else "off" + if isinstance(val, str) and val.lower() in _JIT_C_CONTROLS: + return val.lower() + raise ValueError( + "Bad value for graphblas.config['jit_c_control'] (or the GRAPHBLAS_JIT_C_CONTROL " + f"environment variable): {val!r}. Must be one of: " + + ", ".join(repr(x) for x in _JIT_C_CONTROLS) + + ", or unset" + ) # Ideally this is in operator.py, but lives here to avoid circular references diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index e17889d2a..491a3e0ac 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -409,6 +409,7 @@ def _pair_dtype(op, dtype, dtype2): _has_jit_set, _maybe_warn_jit_skipped, compile_udt_binary_wrapper, + set_jit_c_mixed_on_op, set_jit_c_on_op, ) @@ -689,13 +690,13 @@ def _compile_udt(self, dtype, dtype2): op = _finalize_udt_op( self, dtype, dtype2, ret_type, binary_wrapper, wrapper_sig, TypedUserBinaryOp ) - # Set JIT C definition for auto-generated built-in UDT ops. Only - # same-type pairs apply: mixed UDT+scalar ops don't have a - # meaningful single-type C definition. The arithmetic kernel also - # declares ``z`` with the operands' type, so it fits only an op whose - # result is that type too, which ``truediv`` on integers is not. - # ``_has_jit_set`` gates access to ``lib.GrB_BinaryOp_set_String``, - # which is absent on SS < 9. + # Set JIT C definition for auto-generated built-in UDT ops. A same-type + # pair whose result is that type too gets the one-type kernel below; a + # mixed pair (a scalar with a UDT, or two UDTs of one layout with other + # element types) gets a kernel with a C type for each of z, x and y, + # when it has a C form (``set_jit_c_mixed_on_op``); eq/ne on a mixed + # pair keep the cfunc. ``_has_jit_set`` gates access to + # ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 9. if is_lifted and self.name in _BUILTIN_UDT_BINARY_OPS: is_jittable = ret_type == dtype else: @@ -723,6 +724,20 @@ def _compile_udt(self, dtype, dtype2): is_eq=(self.name == "eq"), ) _maybe_warn_jit_skipped(op._jit_c_info, self.name, dtype.name) + elif _has_numba and _has_jit_set and is_lifted and self.name in _BUILTIN_UDT_BINARY_OPS: + # (no warning when there is none: the cfunc is this op's documented path) + from ..udt import _udt_op_types + + _, input_types = _udt_op_types(self.name, dtype, dtype2) + op._jit_c_info = set_jit_c_mixed_on_op( + op.gb_obj, + self.name, + _BUILTIN_UDT_BINARY_OPS[self.name], + dtype, + dtype2, + ret_type, + input_types, + ) return op @classmethod diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index 81d78f73f..41fba757a 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -17,6 +17,7 @@ """ import ast +import hashlib import itertools import linecache from functools import lru_cache, reduce @@ -1419,6 +1420,102 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): return c_name, f"void {c_name} ({params}) {{ {assigns} }}" +def _c_type_name(dtype): + """The C name SuiteSparse knows ``dtype`` by: a UDT's pinned JIT name, or the C type of + a built-in scalar dtype. ``None`` when it has none. + """ + if dtype.np_type.names is None and dtype.np_type.subdtype is None: + return NP_TO_C_TYPES.get(dtype.np_type) + typedef_info = _udt_c_typedef(dtype.jit_c_name or dtype.name, dtype.np_type) + return None if typedef_info is None else typedef_info[0] + + +def _c_operand(side, path, stored, computed): + """C source reading operand ``side`` (``*x``, or the element at ``path``) in the type the + op computes in. A conversion in C (integer to floating, a real to complex, a narrower + type to a wider one) gives what numpy's ``astype`` gives, which the cfunc uses. + """ + expr = f"(*{side})" if path is None else f"{side}->{path}" + if stored == computed: + return expr + return f"(({NP_TO_C_TYPES[computed]}) ({expr}))" + + +def _make_jit_c_mixed_definition(op_name, py_op, dtype, dtype2, ret_type, input_types): + """Generate a JIT C function definition for a lifted op whose operand types differ, or + whose result type differs from them: a scalar with an array or record UDT (either + order), two array UDTs of one shape with different element types, two records with the + same fields and different leaf types. + + The signature has a C type for each of z, x and y, so the C name carries all three. A + scalar operand is read as ``(*x)`` for every element or leaf. Returns ``(c_name, + c_defn)``, or ``None`` when one of the types has no C name, when an array operand + broadcasts (another shape than the result's), when a record has an array leaf, or when + an element is not computed in the result's element type (integer ``truediv`` computes + in its operands' type and returns float); those keep the Numba cfunc. + + ``input_types`` is what :func:`udt._udt_op_types` returns with ``ret_type``: each + element is computed as the built-in op on the operands' element types computes it, + after converting both operands to that type, as the cfunc does. + """ + names = [_c_type_name(t) for t in (ret_type, dtype, dtype2)] + if None in names: + return None + z_name, x_name, y_name = names + np_x, np_y, np_z = dtype.np_type, dtype2.np_type, ret_type.np_type + + def assign(target, x_path, y_path, x_stored, y_stored, types, out): + in_x, in_y = types + if not in_x == in_y == out: + return None + lhs = _c_operand("x", x_path, x_stored, in_x) + rhs = _c_operand("y", y_path, y_stored, in_y) + return _c_assign_binary(py_op, target, lhs, rhs, out) + + if np_z.names is not None: + leaves_z = list(_iter_record_leaves(np_z)) + leaves_x = list(_iter_record_leaves(np_x)) if np_x.names is not None else None + leaves_y = list(_iter_record_leaves(np_y)) if np_y.names is not None else None + stmts = [] + for i, (_py, c, leaf) in enumerate(leaves_z): + # A record with an array leaf has no C type name, so it returned above + if leaf.shape: # pragma: no cover (safety) + return None + x_path, x_stored = (None, np_x) if leaves_x is None else (c, leaves_x[i][2]) + y_path, y_stored = (None, np_y) if leaves_y is None else (c, leaves_y[i][2]) + stmt = assign(f"z->{c}", x_path, y_path, x_stored, y_stored, input_types[i], leaf) + if stmt is None: + return None + stmts.append(stmt) + body = " ".join(stmts) + else: + base_z, shape_z = np_z.subdtype + operands = [] + for np_t in (np_x, np_y): + if np_t.subdtype is None: + operands.append((None, np_t)) + elif np_t.subdtype[1] != shape_z: + return None # (broadcasting reads through an index map) + else: + operands.append(("v[i]", np_t.subdtype[0])) + (x_path, x_stored), (y_path, y_stored) = operands + stmt = assign("z->v[i]", x_path, y_path, x_stored, y_stored, input_types[0], base_z) + if stmt is None: + return None + body = f"for (int64_t i = 0 ; i < {reduce(mul, shape_z)} ; i++) {{ {stmt} }}" + # The length of x's name keeps two pairs from sharing a C name ("A" and "B_C", "A_B" and + # "C"), and keeps it from any one-type name (``{op_name}_{type}``, whose type name cannot + # start with a digit). SuiteSparse finds a compiled kernel by its op's name, so two ops + # with one name would run each other's kernel. + c_name = f"{op_name}_{len(x_name)}_{x_name}_{y_name}" + if len(c_name) >= lib.GxB_MAX_NAME_LEN: + # SuiteSparse keeps at most 127 characters (see ``_set_jit_c_strings``). A length of 0 + # cannot come from the readable form, so this name is still the pair's own. + c_name = f"{op_name}_0_{hashlib.sha256(c_name.encode()).hexdigest()[:32]}" + params = f"{z_name} *z, const {x_name} *x, const {y_name} *y" + return c_name, f"void {c_name} ({params}) {{ {body} }}" + + def _make_jit_c_comparison_definition(op_name, dtype, *, is_eq): """Generate a JIT C function definition for ``binary.eq[udt]`` / ``binary.ne[udt]``. @@ -1499,8 +1596,15 @@ def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): from ..ss.jit_config import _enable_jit_for_udt _enable_jit_for_udt() - set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + # SuiteSparse refuses a name of GxB_MAX_NAME_LEN (128) characters or more, then takes the + # definition under the empty name: its kernels call no function and leave their outputs + # uninitialized. So the definition goes only with the name, and without it the op keeps + # its cfunc. + info = set_string_func(gb_obj, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME) + if info != lib.GrB_SUCCESS: + return False set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) + return True def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq): @@ -1516,7 +1620,23 @@ def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq result = _make_jit_c_comparison_definition(op_name, dtype, is_eq=is_eq) if result is None: return None - _set_jit_c_strings(gb_obj, *result, set_string_func) + if not _set_jit_c_strings(gb_obj, *result, set_string_func): + return None + return result + + +def set_jit_c_mixed_on_op(gb_obj, op_name, py_op, dtype, dtype2, ret_type, input_types): + """Generate and set the JIT C name and definition of a lifted binary op on two types + (:func:`_make_jit_c_mixed_definition`). Returns ``(c_name, c_defn)``, or ``None`` when + the op has no C form (it keeps the Numba cfunc). Callers check ``_has_jit_set`` first, + as for :func:`set_jit_c_on_op`: without it, SuiteSparse has no + ``GrB_BinaryOp_set_String``. + """ + result = _make_jit_c_mixed_definition(op_name, py_op, dtype, dtype2, ret_type, input_types) + if result is None: + return None + if not _set_jit_c_strings(gb_obj, *result, lib.GrB_BinaryOp_set_String): + return None return result @@ -1532,5 +1652,6 @@ def set_jit_c_on_op(gb_obj, op_name, py_op, dtype, set_string_func, arity=2): result = _make_jit_c_definition(op_name, py_op, dtype, arity) if result is None: return None - _set_jit_c_strings(gb_obj, *result, set_string_func) + if not _set_jit_c_strings(gb_obj, *result, set_string_func): + return None return result diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 677b979d8..62d78a44a 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -18,16 +18,27 @@ ``jit_compiler_is_usable()`` is a non-invasive check that returns True iff the configured compiler exists on disk. Useful for emitting a one-time warning at import or at first JIT use. + +``jit_c_control`` itself is set when the ``suitesparse`` backend initializes, +from ``graphblas.config["jit_c_control"]`` (``'run'`` unless configured +otherwise; see :func:`_set_jit_c_control_at_init`). UDT work raises that +default ``'run'`` to ``'on'`` when it needs C compiled, but keeps a ``'run'`` +the user set (:func:`_enable_jit_for_udt`). +The trade-offs of the JIT, and what to do when it misbehaves, are in the user +guide's initialization page (docs/user_guide/init.rst). """ import contextlib +import hashlib import os import pathlib import platform import re import subprocess -from .. import lib # noqa: F401 (sets up cffi) +# ``ffi`` and ``lib`` are imported in the functions that use them: importing them +# here initializes the backend when this module is imported first, and init +# imports this module for ``_set_jit_c_control_at_init``. # Mapping from the values that show up in baked-in ``-arch`` flags to what # ``platform.machine()`` returns on the running host. @@ -53,7 +64,10 @@ def jit_compiler_is_usable(): Cheap (no compile attempt); use before suggesting ``fix_jit_config()``. """ - cfg = _ss_config() + return _compiler_exists(_ss_config()) + + +def _compiler_exists(cfg): cc = cfg.get("jit_c_compiler_name", "") return bool(cc) and pathlib.Path(cc).exists() @@ -73,8 +87,8 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): compiler from ``sysconfig.get_config_var("CC")``. Set to ``False`` to restrict the repair to a conda environment only. probe : bool, default True - After fixing the config, try to JIT-register a trivial UDT to verify - the compiler actually works. A failed compile makes SuiteSparse lower + After fixing the config, JIT-compile a trivial type to verify the + compiler actually works. A failed compile makes SuiteSparse lower ``jit_c_control`` to ``'load'``, and from SuiteSparse 9.4 on it also raises ``JitError`` that one time; the probe takes that failure, so later operations use the generic kernels without raising. @@ -94,8 +108,8 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): rv = _repair_jit_compiler(cfg, use_sysconfig=use_sysconfig) if rv is None: return None - # Enabling compilation is what the caller asked for. The import-time - # repair deliberately stops short of this; see ``_auto_fix_jit_at_import``. + # Enabling compilation is what the caller asked for, whatever + # ``graphblas.config["jit_c_control"]`` set at init. cfg["jit_c_control"] = "on" # An explicit user-driven fix is a clean opportunity to re-arm # ``NoJITWarning``: if the repair worked, the next UDT auto-lift that @@ -127,14 +141,19 @@ def _repair_jit_compiler(cfg, *, use_sysconfig=True): def _fix_from_conda(cfg, conda_prefix): - """Conda-aware fix; swap the compiler path with one from $CONDA_PREFIX/bin/.""" + """Conda-aware fix; swap the compiler path with one from the environment's binaries: + ``$CONDA_PREFIX/bin/``, or ``$CONDA_PREFIX/Library/bin/`` with ``.exe`` names (Windows). + """ jit_cc = cfg["jit_c_compiler_name"] if not pathlib.Path(jit_cc).exists(): - cc_basename = pathlib.Path(jit_cc).name - bin_dir = pathlib.Path(conda_prefix) / "bin" - for candidate in [cc_basename, "cc", "clang", "gcc"]: - local_cc = bin_dir / candidate - if local_cc.exists(): + # (a Windows path's basename, also when read on another system) + cc_basename = re.split(r"[\\/]", jit_cc)[-1] + prefix = pathlib.Path(conda_prefix) + names = [cc_basename, "cc", "clang", "gcc"] + names += [n + ".exe" for n in names if not n.lower().endswith(".exe")] + for bin_dir in [prefix / "bin", prefix / "Library" / "bin"]: + local_cc = next((bin_dir / n for n in names if (bin_dir / n).exists()), None) + if local_cc is not None: cfg["jit_c_compiler_name"] = str(local_cc) break else: @@ -195,39 +214,158 @@ def _strip_mismatched_arch(flags): ) +# The settings a compile depends on. The probe's type is named after them; see +# ``_probe_jit``. +_TOOLCHAIN_KEYS = ( + "jit_c_compiler_name", + "jit_c_compiler_flags", + "jit_c_linker_flags", + "jit_c_libraries", +) + + def _probe_jit(cfg): - """Probe a trivial JIT compile to verify the config works. - - ``jit_c_control`` is left wherever SuiteSparse puts it, and that is not - predictable from here. A library built without the JIT clamps every write - down to ``'run'`` (``GB_jitifyer_set_control``), so ``'on'`` never takes on - such a build no matter how good the compiler is; separately, a compile - failure drops it to ``'load'``, a load failure to ``'run'``, and a failed - hash insert to ``'pause'``. The contract here is only "did this work?", so - the caller gets a bool. + """Compile a trivial type with the JIT, or load it from the cache, to check the config works. + + ``jit_c_control`` must be ``'on'``, and is left wherever SuiteSparse puts + it, which is not predictable from here. A library built without the JIT + clamps every write down to ``'run'`` (``GB_jitifyer_set_control``), so + ``'on'`` never takes on such a build no matter how good the compiler is; + separately, a compile failure drops it to ``'load'``, a load failure to + ``'run'``, and a failed hash insert to ``'pause'``. The contract here is + only "did this work?", so the caller gets a bool. + + It goes through the C API, so nothing is registered on the Python side, + and it can run while the backend initializes. Once a kernel is in the + cache, a later probe loads it (about 2 ms) instead of compiling (about + 0.3 s). The type is named after the toolchain settings, so a compiler or + flags this cache has not seen compile a probe of their own, rather than + load one that another toolchain left in a shared cache. """ - from ... import dtypes as _dtypes - - # The probe dtype is installed at ``dtypes.ss._jit_probe`` on first - # success; reuse it on a second probe so a repeat call doesn't appear to - # fail. Without the hasattr short-circuit, ``register_new`` would raise - # ``ValueError("name unavailable")`` on the second call. - probe_name = "_jit_probe" - if hasattr(_dtypes.ss, probe_name): - return True - global _probing_jit - _probing_jit = True - try: - _dtypes.ss.register_new(probe_name, "typedef struct { int _probe ; } _jit_probe ;") - except Exception: - # ``register_new`` can raise ``JitError`` (bad path / flags / arch / - # SDK), ``RuntimeError`` (SS<8 has no JIT), or one of its input - # validation ``ValueError``s. The probe's contract is "did this - # work?", so absorb every failure mode here. + from .. import ffi, lib + + toolchain = "\n".join(cfg.get(key, "") for key in _TOOLCHAIN_KEYS) + name = "_jit_probe_" + hashlib.sha256(toolchain.encode()).hexdigest()[:16] + defn = f"typedef struct {{ int _probe ; }} {name} ;" + type_ptr = ffi.new("GrB_Type*") + # Size 0 asks SuiteSparse to get the size from a compiled kernel. + if lib.GxB_Type_new(type_ptr, 0, name.encode(), defn.encode()) != lib.GrB_SUCCESS: return False - finally: - _probing_jit = False - return True + lib.GrB_Type_free(type_ptr) + return cfg["jit_c_control"] == "on" + + +class _CConfig: + """The ``gb.ss.config`` keys :func:`_set_jit_c_control_at_init` needs, through the C API. + + That runs while the backend initializes, possibly in the middle of + importing a core module, and importing ``graphblas.ss`` (for its + ``config``) would import the core modules. + """ + + _strings = { + "jit_c_compiler_name": "GxB_JIT_C_COMPILER_NAME", + "jit_c_compiler_flags": "GxB_JIT_C_COMPILER_FLAGS", + "jit_c_linker_flags": "GxB_JIT_C_LINKER_FLAGS", + "jit_c_libraries": "GxB_JIT_C_LIBRARIES", + "jit_cache_path": "GxB_JIT_CACHE_PATH", + } + + def __init__(self): + from .. import lib + + self._controls = { + "off": lib.GxB_JIT_OFF, + "pause": lib.GxB_JIT_PAUSE, + "run": lib.GxB_JIT_RUN, + "load": lib.GxB_JIT_LOAD, + "on": lib.GxB_JIT_ON, + } + + def __getitem__(self, key): + from .. import ffi, lib + + if key == "jit_c_control": + val = ffi.new("int32_t*") + info = lib.GxB_Global_Option_get_INT32(lib.GxB_JIT_C_CONTROL, val) + else: + val = ffi.new("char**") + info = lib.GxB_Global_Option_get_CHAR(getattr(lib, self._strings[key]), val) + if info != lib.GrB_SUCCESS: # pragma: no cover (safety) + raise RuntimeError(f"Could not get SuiteSparse:GraphBLAS option {key!r}") + if key == "jit_c_control": + return next(k for k, v in self._controls.items() if v == val[0]) + return ffi.string(val[0]).decode() + + def __setitem__(self, key, value): + from .. import lib + + if key == "jit_c_control": + info = lib.GxB_Global_Option_set_INT32(lib.GxB_JIT_C_CONTROL, self._controls[value]) + else: + option = getattr(lib, self._strings[key]) + info = lib.GxB_Global_Option_set_CHAR(option, value.encode()) + if info != lib.GrB_SUCCESS: # pragma: no cover (safety) + raise RuntimeError(f"Could not set SuiteSparse:GraphBLAS option {key!r}") + + def get(self, key, default=None): + return self[key] if key == "jit_c_control" or key in self._strings else default + + +def _cache_is_writable(cfg): + """Whether SuiteSparse can write new kernels to its cache. + + A probe kernel already in the cache loads from a read-only one, and the + first compile after it would fail (with ``JitError`` from SuiteSparse 9.4 + on), so the probe alone does not answer this. + """ + path = cfg.get("jit_cache_path") + return bool(path) and all( + os.access(pathlib.Path(path, sub), os.W_OK | os.X_OK) for sub in ("c", "lib", "tmp") + ) + + +def _set_jit_c_control_at_init(control): + """Set ``jit_c_control`` as the backend initializes, from ``graphblas.config``. + + ``control`` is ``graphblas.config["jit_c_control"]``, already checked, or + None when it is not set: then the control is the default ``'run'``, which + UDT work raises (:func:`_enable_jit_for_udt`). A ``'run'`` that is set is + the user's choice, and UDT work keeps it (``_user_set_run``). + Every value but ``'on'`` is set as given. ``'on'`` is set only where it + works, which :func:`_probe_jit` checks by compiling a trivial type (or + loading it from the cache, after the first time). Where it does not, the + control is set to ``'run'`` (where the conda-forge and PyPI builds of + SuiteSparse start it), with no error or warning: a library built without + the JIT, which clamps every write to ``'run'``; no compiler on disk, even + after :func:`_repair_jit_compiler`; a cache directory it cannot write to; + or a compiler that fails the probe. + The probe is load-bearing: from SuiteSparse 9.4 on a failed compile raises + ``JitError`` (once; SuiteSparse then drops the control to ``'load'``), and + without the probe that would be the first user operation that needs a new + kernel, such as a slice, which SuiteSparse precompiles no kernel for. + """ + from . import _IS_SSGB7 + + if _IS_SSGB7: + return # no JIT before SuiteSparse:GraphBLAS 8 + global _jit_enabled_for_udt, _configured_jit_c_control, _user_set_run + cfg = _CConfig() + _configured_jit_c_control = control + _user_set_run = control == "run" + if control != "on": + cfg["jit_c_control"] = "run" if control is None else control + return + if not _compiler_exists(cfg): + _repair_jit_compiler(cfg) + _jit_enabled_for_udt = False + if _compiler_exists(cfg) and _cache_is_writable(cfg): + cfg["jit_c_control"] = "on" + _jit_enabled_for_udt = cfg["jit_c_control"] == "on" and _probe_jit(cfg) + if not _jit_enabled_for_udt: + # Not the 'load' a failed probe leaves: 'load' still loads kernels from + # the on-disk cache, which other builds of this version may have written. + cfg["jit_c_control"] = "run" def _auto_fix_jit_at_import(): @@ -237,12 +375,14 @@ def _auto_fix_jit_at_import(): breaks ``import graphblas.ss``. The body sticks to dict ops and the string rewriting in :func:`_repair_jit_compiler`. - Deliberately leaves ``jit_c_control`` at whatever SuiteSparse set. An - attribute access such as ``gb.ss.about["library_version"]`` imports this - submodule, and that must not change what any later operation computes or - which kernels SuiteSparse is willing to load from its on-disk cache. - Compilation is enabled later, by :func:`_enable_jit_for_udt`, when a UDT - actually needs a kernel built, or a type or op is registered from C source. + Deliberately leaves ``jit_c_control`` alone. An attribute access such as + ``gb.ss.about["library_version"]`` imports this submodule, and that must + not change what any later operation computes or which kernels SuiteSparse + is willing to load from its on-disk cache. The control is set when the + backend initializes (:func:`_set_jit_c_control_at_init`), which also + repairs the compiler when it sets ``'on'``. Otherwise compilation is + enabled later, by :func:`_enable_jit_for_udt`, when a UDT actually needs a + kernel built, or a type or op is registered from C source. """ cfg = _ss_config() if "jit_c_control" not in cfg: @@ -251,57 +391,61 @@ def _auto_fix_jit_at_import(): _repair_jit_compiler(cfg) -# Tri-state: ``None`` until the first request to compile C source while the -# control allows compiling, then the answer to "can this process -# JIT-compile?" until ``fix_jit_config`` re-arms it. +# Tri-state: ``None`` until the compiler is probed (at init when that sets +# ``'on'``, or at the first request to compile C source while the control +# allows compiling), then the answer to "can this process JIT-compile?" until +# ``fix_jit_config`` re-arms it. _jit_enabled_for_udt = None -# True only while ``_probe_jit`` is registering its own UDT. That registration -# goes through ``dtypes.ss.register_new``, which asks to enable the JIT, which -# would probe again; the second probe would register the same name a second -# time and the first one would then die on cffi's "multiple declarations". -_probing_jit = False +# ``graphblas.config["jit_c_control"]`` as init found it, None when it was not set. +_configured_jit_c_control = None +# Whether a ``'run'`` in effect is one the user set, which ``_enable_jit_for_udt`` +# keeps: a ``'run'`` configured before init, or one written to ``gb.ss.config`` +# (whose setter keeps this current) while something else is configured. With +# nothing configured, a ``'run'`` written there is the default, as before it. +_user_set_run = False def _enable_jit_for_udt(): """Enable JIT compilation the first time C source needs compiling. - SuiteSparse defaults ``jit_c_control`` to ``'run'``, which runs kernels - already loaded but neither compiles nor loads any. UDT auto-lift wants - ``'on'``; without it every UDT op falls back to the Numba - function-pointer path (typically 2-3x slower for elementwise ops). An op - or type registered from C source has no fallback: under ``'run'`` its - creation fails outright. + ``jit_c_control`` defaults to ``'run'`` (``graphblas.config``, set at + init), which runs kernels already loaded but neither compiles nor loads + any. UDT auto-lift wants ``'on'``; without it every UDT op falls back to + the Numba function-pointer path (typically 2-3x slower for elementwise + ops). An op or type registered from C source has no fallback: under + ``'run'`` its creation fails outright. This is where that bump belongs, rather than at import: registering a type or an op from C source, or arming an op with C source, is an act that plainly involves compiling C, so enabling the compiler is not a surprise. Reading ``gb.ss.about`` is not, so it leaves the setting alone. - Only the default ``'run'`` is raised. An explicit ``'off'``, ``'pause'``, - or ``'load'`` is honored, and ``'load'`` is also where SuiteSparse leaves - the control after a compile fails, so raising it would retry that compile - at every later UDT op (raising ``JitError`` each time on SuiteSparse 9.4 - and later). Returns True iff a kernel armed now will be compiled. Whether - this process can compile at all is probed once and cached. + Only the default ``'run'`` is raised. A ``'run'`` the user set (through + ``graphblas.config["jit_c_control"]``, or ``gb.ss.config``; see + ``_user_set_run``) is honored, often to keep processes that run at the + same time from compiling into one cache (GraphBLAS #449). So are + ``'off'``, ``'pause'``, and ``'load'``, and ``'load'`` is also where + SuiteSparse leaves the control after a compile fails, so raising it would + retry that compile at every later UDT op (raising ``JitError`` each time on + SuiteSparse 9.4 and later). Returns True iff a kernel armed now will be + compiled. Whether this process can compile at all is probed once and + cached. """ global _jit_enabled_for_udt - if _probing_jit: - # Re-entered from the probe's own registration. The probe is the thing - # deciding this answer, so say yes and let it finish rather than - # starting a second one inside it. - return True cfg = _ss_config() control = cfg.get("jit_c_control") - if control in (None, "off", "pause", "load"): - # ``None`` is a library with no JIT at all. The others are honored as - # set and not cached, so a saved ``'run'`` restored afterwards is - # still raised. + if control in (None, "off", "pause", "load") or (control == "run" and _user_set_run): + # ``None`` is a library with no JIT at all. The others are honored and + # not cached, so they decide nothing once the control is back at the + # default. return False if _jit_enabled_for_udt is None: if not jit_compiler_is_usable(): _repair_jit_compiler(cfg) - if _jit_enabled_for_udt := jit_compiler_is_usable(): + # A probe kernel already in a read-only cache loads, and the first new + # kernel would then fail to compile; see ``_cache_is_writable``. + if _jit_enabled_for_udt := jit_compiler_is_usable() and _cache_is_writable(cfg): cfg["jit_c_control"] = "on" # A library built without the JIT clamps that write back down to # ``'run'``. Otherwise the probe is load-bearing: a failed compile @@ -366,6 +510,15 @@ def _compiling_c_source(name): "itself after a compile fails. Set gb.ss.config['jit_c_control'] = " "'on' to compile it." ) + elif after == before == "run" and _user_set_run: + reason = ( + "jit_c_control was set to 'run' (by graphblas.config['jit_c_control'], " + "the GRAPHBLAS_JIT_C_CONTROL environment variable, or gb.ss.config), " + "which neither compiles nor loads, and python-graphblas keeps a 'run' " + "that was set. To compile it, set gb.ss.config['jit_c_control'] = 'on', " + "or None for the default 'run', which python-graphblas raises to 'on' " + "for C source." + ) elif after == before == "run": reason = ( "jit_c_control is still 'run', which neither compiles nor loads: " diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index 45aa47c01..93a6748ef 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -1,2 +1,8 @@ autocompute: True mapnumpy: True +# SuiteSparse:GraphBLAS C JIT control, set when graphblas initializes the "suitesparse" +# backend: "on", "load", "run", "pause" or "off". Environment variable: GRAPHBLAS_JIT_C_CONTROL. +# Unset (null) is "run", raised to "on" when UDT work needs a kernel compiled; a value set +# through this key is kept. +# See "The SuiteSparse:GraphBLAS JIT" in the user guide's initialization page. +jit_c_control: null diff --git a/graphblas/ss/__init__.py b/graphblas/ss/__init__.py index 33a342373..87bc3e2cd 100644 --- a/graphblas/ss/__init__.py +++ b/graphblas/ss/__init__.py @@ -14,11 +14,11 @@ # Repair the JIT compiler path on import. Conda-built psg bakes in a # compiler path that almost never exists at runtime, so JIT silently falls # back to the cfunc path (2-3x slowdown). This only rewrites the compiler - # name and flags; ``jit_c_control`` is left where SuiteSparse set it, so - # importing this submodule cannot change what a later operation computes. - # Compilation is enabled on demand by ``_enable_jit_for_udt``. Skipped on - # vanilla: the repair reads ``gb.ss.config``, which needs GxB callables - # vanilla strips. + # name and flags; ``jit_c_control`` is left where the backend's init set + # it (``_set_jit_c_control_at_init``), so importing this submodule cannot + # change what a later operation computes. Compilation is enabled on demand + # by ``_enable_jit_for_udt``. Skipped on vanilla: the repair reads + # ``gb.ss.config``, which needs GxB callables vanilla strips. if backend == "suitesparse": _auto_fix() del _auto_fix diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index f036b3bb7..40c7695f0 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -143,9 +143,17 @@ class GlobalConfig(BaseConfig): "run" : run JIT kernels if already loaded, but don't load or compile "load" : able to load and run JIT kernels; may not compile "on" : full JIT: able to compile, load, and run - Only available for SuiteSparse:GraphBLAS >=8 + Only available for SuiteSparse:GraphBLAS >=8. + python-graphblas sets it at init from ``graphblas.config["jit_c_control"]``, + or to "run" when that is not set, and UDT work raises that default "run" to + "on" when it needs a kernel compiled. A "run" configured there is kept, and + so is one set here over another configured value; with nothing configured, + "run" set here is the default again, as None is in any case. The user + guide's initialization page covers the trade-offs, and when "on" is safe. jit_use_cmake : bool Whether to use cmake to compile the JIT kernels. + Always True for SuiteSparse built with MSVC (on Windows), where + python-graphblas does not turn the JIT on yet. Only available for SuiteSparse:GraphBLAS >=8 jit_c_compiler_name : str C compiler for JIT kernels. @@ -254,6 +262,25 @@ class GlobalConfig(BaseConfig): "on": lib.GxB_JIT_ON, } + def __setitem__(self, key, val): + key = key.lower() + if key != "jit_c_control": + super().__setitem__(key, val) + return + from ..core.ss import jit_config + + # UDT work keeps a 'run' the user set. One written here counts when it + # changes the configured control (graphblas.config, read at init) or repeats + # a configured 'run'. With nothing configured it is the default 'run' again, + # which UDT work raises, so saving and restoring the default changes nothing. + # None is that default whatever is configured. + super().__setitem__(key, "run" if val is None else val) + if isinstance(val, str): + val = val.lower() + elif val is not None: + val = self._reverse_enumerations[key].get(val, val) + jit_config._user_set_run = val == "run" and jit_config._configured_jit_c_control is not None + class About(Mapping): _modes = { diff --git a/graphblas/tests/test_core.py b/graphblas/tests/test_core.py index a15da6f54..bfad8279d 100644 --- a/graphblas/tests/test_core.py +++ b/graphblas/tests/test_core.py @@ -41,6 +41,27 @@ def test_bad_init(): gb.init(**params) +def test_jit_c_control_config(): + """``config["jit_c_control"]``, which init sets SuiteSparse's JIT control to.""" + import yaml + + path = pathlib.Path(gb.__file__).parent / "graphblas.yaml" + # Unset by default, which init takes as a 'run' that UDT work may raise + assert yaml.safe_load(path.read_text())["jit_c_control"] is None + # YAML 1.1 reads an unquoted on/off in a user's config file as a bool + cases = [("on", "on"), ("LOAD", "load"), ("run", "run"), ("pause", "pause"), ("off", "off")] + cases += [(True, "on"), (False, "off"), (None, None)] + for val, expected in cases: + with gb.config.set(jit_c_control=val): + assert gb._jit_c_control_config() == expected + for val in ["of", "", 1]: + with ( + gb.config.set(jit_c_control=val), + pytest.raises(ValueError, match="GRAPHBLAS_JIT_C_CONTROL"), + ): + gb._jit_c_control_config() + + def test_bad_libget(): with pytest.raises(AttributeError, match="GrB_bad_name"): gb.core.base.libget("GrB_bad_name") diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 5fb3aad63..8ca033625 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -498,8 +498,8 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): with pytest.warns(NoJITWarning, match="without JIT") as record: plus_udt = gb.binary.plus[udt] # The cause has to name the UDT, on every host. Nothing armed a kernel for - # this dtype, so ``jit_c_control`` is still whatever SuiteSparse set, and - # blaming it would send the reader to a setting that cannot fix anything. + # this dtype, so ``jit_c_control`` is still whatever init set, and blaming + # it would send the reader to a setting that cannot fix anything. assert "not expressible as a C struct" in str(record[0].message) assert plus_udt.jit_c_source is None # no JIT # ``eq`` has a separate JIT codegen path from arithmetic (BOOL output, leaf @@ -521,8 +521,8 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): def test_no_jit_warning_blames_the_udt_not_the_host_jit_state(): """The cause named must not move with the host's JIT state. - On a runner whose JIT probe fails, SuiteSparse leaves ``jit_c_control`` - at ``'run'`` and ``_enable_jit_for_udt`` caches its ``False``. Codegen + On a runner whose JIT probe fails, init leaves ``jit_c_control`` at + ``'run'`` and caches the probe's ``False``. Codegen for this UDT read neither, so naming either as the cause points the reader at a setting that cannot fix it. CI hit exactly that: on ubuntu-latest the message blamed ``jit_c_control``. @@ -592,13 +592,14 @@ def test_udt_layouts_under_one_name_get_their_own_jit_c_names(): with pytest.raises(ValueError, match="unavailable"): dtypes.ss.register_new("_SameJitName", "typedef struct { int8_t q ; } _SameJitName ;") # The rest compiles kernels for both layouts, which needs a C compiler the JIT can use. + # Asked under 'on': a 'run' the session set (GRAPHBLAS_JIT_C_CONTROL) would answer no. from graphblas.core.ss.jit_config import _enable_jit_for_udt - if not _enable_jit_for_udt(): - pytest.skip("the JIT cannot compile here (no usable C compiler)") jit_c_control = gb.ss.config["jit_c_control"] gb.ss.config["jit_c_control"] = "on" try: + if not _enable_jit_for_udt(): + pytest.skip("the JIT cannot compile here (no usable C compiler)") for udt in [small, big]: shape = udt.np_type.shape one = np.full(shape, 1, np.int16) diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index b1469b7a0..4e553f51b 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -246,9 +246,14 @@ def test_openmp_enabled(): ) -def test_global_config(): +def test_global_config(monkeypatch): d = {} config = gb.ss.config + if "jit_c_control" in config: # SuiteSparse:GraphBLAS 8+ + from graphblas.core.ss import jit_config + + # Writing jit_c_control also sets whether UDT work keeps a "run"; undo that after + monkeypatch.setattr(jit_config, "_user_set_run", jit_config._user_set_run) for k in config: d[k] = config[k] assert d == config @@ -271,6 +276,13 @@ def test_global_config(): for k in config: if k in config._defaults: config[k] = None + elif k == "jit_c_control": + # None is the default "run" whatever graphblas.config has, and not a "run" + # the user chose, so UDT work may still raise it to "on" + config[k] = None + assert config[k] == "run" + assert jit_config._user_set_run is False + config[k] = d[k] else: with pytest.raises(ValueError, match="Unable to set default value for"): config[k] = None diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 9e083a2a1..3721eee1f 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -47,12 +47,22 @@ # can call it themselves. Tests use it through the public surface. _fix_jit_config = gb.ss.fix_jit_config if not _IS_SSGB7 else (lambda: None) -# ``_auto_fix_jit_at_import`` repairs the compiler path but leaves -# ``jit_c_control`` alone, so compilation is enabled on demand by -# ``_enable_jit_for_udt``. Ask it once here, before the autouse ``_setup_jit`` -# fixture mutates the live config: it answers whether this environment can -# JIT-compile at all, and caches that answer for the rest of the process. -_JIT_WORKS_AT_IMPORT = not _IS_SSGB7 and backend == "suitesparse" and _enable_jit_for_udt() + +def _jit_works_here(): + """Whether this environment can JIT-compile at all, whatever control the session has. + + The backend's init set ``jit_c_control`` from ``gb.config`` (``'run'`` by + default) and, under ``'on'``, probed the compiler. ``_enable_jit_for_udt`` + raises the default ``'run'`` and probes, but answers no under ``'off'``, + ``'pause'``, ``'load'`` or a ``'run'`` the user set without probing, so ask + it under ``'on'``, and put the control back. It caches the answer for the + rest of the process. Asked once here, before the autouse ``_setup_jit`` + fixture mutates the live config. + """ + if _IS_SSGB7 or backend != "suitesparse": + return False + with _jit_mode("on"): + return _enable_jit_for_udt() @pytest.fixture(scope="module", autouse=True) @@ -114,19 +124,28 @@ def _jit_mode(mode): gb.ss.config["jit_c_control"] = prev -def _run_in_fresh_process(script): +_JIT_WORKS_AT_IMPORT = _jit_works_here() + + +def _run_in_fresh_process(script, jit_c_control=None, **extra_env): """Run ``script`` in a new interpreter on this tree; return its ``key: value`` lines. The script must print ``package: `` so the child can be shown - to have imported this tree before anything it reports is trusted. + to have imported this tree before anything it reports is trusted. The + child gets ``GRAPHBLAS_JIT_C_CONTROL=jit_c_control``, or none at all, and + ``extra_env``. """ repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root), **extra_env) + env.pop("GRAPHBLAS_JIT_C_CONTROL", None) + if jit_c_control is not None: + env["GRAPHBLAS_JIT_C_CONTROL"] = jit_c_control result = subprocess.run( [sys.executable, "-c", script], capture_output=True, text=True, check=False, - env=dict(os.environ, PYTHONPATH=str(repo_root)), + env=env, ) report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" assert result.returncode == 0, report @@ -159,17 +178,115 @@ def control(): @pytest.mark.skipif("_IS_SSGB7") -def test_ss_import_leaves_jit_c_control_alone(): +@pytest.mark.parametrize("jit_c_control", [None, "run", "on", "pause"]) +def test_ss_import_leaves_jit_c_control_alone(jit_c_control): """Reading from ``gb.ss`` must not change the numerical machinery. ``gb.ss`` is reached by attribute access, so importing it has to be inert. It used to raise ``jit_c_control`` to ``'on'``, which decides whether SuiteSparse loads kernels from its on-disk cache, so a line as harmless as ``print(gb.ss.about["library_version"])`` changed how every later - operation in the process was computed. + operation in the process was computed. The control is set by the + backend's init, from ``gb.config`` (here ``GRAPHBLAS_JIT_C_CONTROL``), + and ``'on'`` becomes ``'run'`` where the JIT cannot compile. Not set, it + is ``'run'``. """ - lines, report = _run_in_fresh_process(_SS_IMPORT_PROBE) + lines, report = _run_in_fresh_process(_SS_IMPORT_PROBE, jit_c_control) assert lines["after"] == lines["before"], report + if jit_c_control != "on": + expected = getattr(gb.core.lib, f"GxB_JIT_{(jit_c_control or 'run').upper()}") + assert lines["before"] == str(expected), report + elif _JIT_WORKS_AT_IMPORT: + assert lines["before"] == str(gb.core.lib.GxB_JIT_ON), report + else: + # psg's wheels (built without the JIT), or no working C compiler + assert lines["before"] == str(gb.core.lib.GxB_JIT_RUN), report + + +_BAD_JIT_C_CONTROL_PROBE = """ +import graphblas as gb +import suitesparse_graphblas as ssgb + +print("package: " + gb.__file__) +try: + gb.Vector +except ValueError as exc: + print("error: " + str(exc)) +print("initialized: " + str(ssgb.is_initialized())) +""" + + +def test_bad_jit_c_control_config_fails_init_before_it_starts(): + """A bad ``GRAPHBLAS_JIT_C_CONTROL`` says so, and leaves SuiteSparse uninitialized.""" + lines, report = _run_in_fresh_process(_BAD_JIT_C_CONTROL_PROBE, "of") + assert "GRAPHBLAS_JIT_C_CONTROL" in lines.get("error", ""), report + assert "'of'" in lines["error"], report + assert lines["initialized"] == "False", report + + +@pytest.mark.skipif("_IS_SSGB7") +def test_set_jit_c_control_at_init(monkeypatch): + """Init sets each configured value as given, and ``'on'`` only where the JIT compiles. + + Where it cannot, it sets ``'run'`` with no error: a library built without + the JIT, no compiler on disk, a cache it cannot write to, or a compiler + that fails the probe. Nothing configured (None) sets ``'run'`` too, and + only that ``'run'`` is left for UDT work to raise. Called directly here, + as init calls it. + """ + from graphblas.core.ss import jit_config + from graphblas.core.ss.jit_config import _set_jit_c_control_at_init + + cfg = gb.ss.config + prev_control = cfg["jit_c_control"] + prev_cc = cfg["jit_c_compiler_name"] + for name in ["_jit_enabled_for_udt", "_configured_jit_c_control", "_user_set_run"]: + monkeypatch.setattr(jit_config, name, getattr(jit_config, name)) + try: + for control in ["load", "run", "pause", "off"]: + _set_jit_c_control_at_init(control) + # A library built without the JIT clamps 'load' down to 'run'. + expected = {control, "run"} if control == "load" else {control} + assert cfg["jit_c_control"] in expected + assert jit_config._configured_jit_c_control == control + assert jit_config._user_set_run is (control == "run") + _set_jit_c_control_at_init(None) + assert cfg["jit_c_control"] == "run" + assert jit_config._configured_jit_c_control is None + assert jit_config._user_set_run is False + _set_jit_c_control_at_init("on") + assert jit_config._configured_jit_c_control == "on" + assert cfg["jit_c_control"] == ("on" if jit_config._jit_enabled_for_udt else "run") + if _JIT_WORKS_AT_IMPORT: + assert cfg["jit_c_control"] == "on" + # A cache it cannot write to: the cached probe would load, and the first + # new kernel would fail to compile. + with monkeypatch.context() as m: + m.setattr(jit_config, "_cache_is_writable", lambda cfg: False) + _set_jit_c_control_at_init("on") + assert cfg["jit_c_control"] == "run" + # A "compiler" that exists but cannot compile anything. The probe's type + # is named after the toolchain, so the probe kernel this process has + # already loaded does not answer for this one. SuiteSparse built with + # MSVC always compiles through cmake, whose Visual Studio generator + # chooses the compiler itself, so this setting cannot fail the probe. + if not cfg["jit_use_cmake"]: + cfg["jit_c_compiler_name"] = str(Path(__file__).parent) + _set_jit_c_control_at_init("on") + assert cfg["jit_c_control"] == "run" + assert jit_config._jit_enabled_for_udt is False + # No compiler on disk and nothing to repair it from: no probe at all. + cfg["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + probes = [] + monkeypatch.setattr(jit_config, "_repair_jit_compiler", lambda cfg: None) + monkeypatch.setattr(jit_config, "_probe_jit", probes.append) + _set_jit_c_control_at_init("on") + assert cfg["jit_c_control"] == "run" + assert jit_config._jit_enabled_for_udt is False + assert probes == [] + finally: + cfg["jit_c_compiler_name"] = prev_cc + cfg["jit_c_control"] = prev_control # Run in a subprocess: this is about the very first call in a process, and by @@ -187,11 +304,11 @@ def test_ss_import_leaves_jit_c_control_alone(): def test_fix_jit_config_first_call_agrees_with_the_second(): """``fix_jit_config()`` must report on the environment, not on itself. - Its probe registers a UDT, and registering a UDT asks to enable the JIT, - which would start a second probe inside the first. The inner probe would - register the probe type and the outer one would then die redeclaring it, - so a working environment would answer ``False`` the first time and - ``True`` every time after. + Its probe once registered a UDT, and registering a UDT asks to enable the + JIT, which started a second probe inside the first. The inner probe + registered the probe type and the outer one then died redeclaring it, so a + working environment answered ``False`` the first time and ``True`` every + time after. The probe now goes through the C API and registers nothing. """ lines, report = _run_in_fresh_process(_FIX_JIT_TWICE_PROBE) assert lines["first"] == lines["second"], report @@ -199,8 +316,8 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): # Run in a subprocess: this process has already probed the JIT and cached the # answer, and a fresh one has to probe at its first registration. Each -# registration starts from ``'run'``, so one registrar raising the control -# cannot cover for another that does not. +# registration starts from the default ``'run'`` (None), so one registrar +# raising the control cannot cover for another that does not. _SS_REGISTER_NEW_PROBE = """ import graphblas as gb from graphblas import binary, indexbinary, indexunary, select, unary @@ -238,7 +355,7 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): "FP64", "FP64", "FP64", "FP64", ) for kind, register in registrars.items(): - gb.ss.config["jit_c_control"] = "run" + gb.ss.config["jit_c_control"] = None try: register() except Exception as exc: @@ -250,7 +367,7 @@ def test_fix_jit_config_first_call_agrees_with_the_second(): @pytest.mark.skipif("_IS_SSGB7") def test_ss_register_new_enables_the_jit(): - """Registering an op from C source raises SuiteSparse's default ``'run'``. + """Registering an op from C source raises the default ``'run'``. SuiteSparse compiles the C definition, or loads it from its cache, when the op is created. ``'run'`` does neither, so without the raise the @@ -298,9 +415,9 @@ def test_ss_register_new_says_why_the_jit_cannot_build_it(monkeypatch): unary.ss.register_new(op_name, cdef, "FP64", "FP64") with pytest.raises(InvalidValue, match=match): dtypes.ss.register_new(type_name, tdef) - # A process that cannot compile leaves SuiteSparse's default in place. + # A process that cannot compile leaves the default 'run' in place. monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", False) - gb.ss.config["jit_c_control"] = "run" + gb.ss.config["jit_c_control"] = None with pytest.raises(NullPointer, match="jit_c_control is still 'run'"): unary.ss.register_new(op_name, cdef, "FP64", "FP64") # A probe that fails in this very call leaves 'load' before the op is created. @@ -314,13 +431,141 @@ def failed_probe(cfg): monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", None) monkeypatch.setattr(jit_config, "jit_compiler_is_usable", lambda: True) monkeypatch.setattr(jit_config, "_probe_jit", failed_probe) - gb.ss.config["jit_c_control"] = "run" + gb.ss.config["jit_c_control"] = None with pytest.raises(NullPointer, match="the C compiler does not work here"): unary.ss.register_new(op_name, cdef, "FP64", "FP64") finally: gb.ss.config["jit_c_control"] = prev +@pytest.mark.skipif("_IS_SSGB7") +def test_ss_register_new_under_a_chosen_run_says_how_to_opt_in(monkeypatch): + """C source cannot be built under a ``'run'`` the user set, and the error says so. + + UDT work raises only the default ``'run'``. One that was set, often to keep + processes from compiling into one cache at the same time, is kept, so the + registration fails, and the message has to say why and what to set instead. + Here the user configured ``'on'`` before init and writes ``'run'`` later. + """ + from graphblas.core.ss import jit_config + from graphblas.exceptions import InvalidValue, NullPointer + + monkeypatch.setattr(jit_config, "_configured_jit_c_control", "on") + monkeypatch.setattr(jit_config, "_user_set_run", jit_config._user_set_run) + suffix = uuid.uuid4().hex[:8] + op_name = f"chosen_neg_{suffix}" + type_name = f"chosen_t_{suffix}" + prev = gb.ss.config["jit_c_control"] + try: + gb.ss.config["jit_c_control"] = "run" + assert jit_config._user_set_run is True + match = "jit_c_control was set to 'run' .*GRAPHBLAS_JIT_C_CONTROL.* = 'on', or None" + with pytest.raises(NullPointer, match=match): + unary.ss.register_new( + op_name, + f"void {op_name} (double *z, double *x) {{ (*z) = -(*x) ; }}", + "FP64", + "FP64", + ) + with pytest.raises(InvalidValue, match=match): + dtypes.ss.register_new(type_name, f"typedef struct {{ int64_t c_a ; }} {type_name} ;") + assert gb.ss.config["jit_c_control"] == "run" + finally: + gb.ss.config["jit_c_control"] = prev + + +# Run in a subprocess: the forms that set the control are read once, at init. +# ``form`` is defined by a line put in front of this. +_UDT_WORK_PROBE = """ +import uuid + +import numpy as np + +import graphblas as gb + +if form == "gb.config.set": + gb.config.set(jit_c_control="run") +elif form == "gb.ss.config": + gb.ss.config["jit_c_control"] = "run" +print("package: " + gb.__file__) +print("before: " + gb.ss.config["jit_c_control"]) +name = "uw_" + uuid.uuid4().hex[:8] +try: + gb.dtypes.ss.register_new(name, "typedef struct { int64_t uw_a ; } " + name + " ;") +except Exception as exc: + print("register_new: " + type(exc).__name__ + ": " + str(exc)) +else: + print("register_new: ok") +print("after register_new: " + gb.ss.config["jit_c_control"]) +try: + import numba +except ImportError: + numba = None +if numba is not None: + udt = gb.dtypes.register_anonymous( + np.dtype([("uw_x", np.int64), ("uw_y", np.int64)], align=True), "UdtWorkProbe" + ) + v = gb.Vector.from_coo([0, 1], np.array([(1, 2), (3, 4)], dtype=udt.np_type), dtype=udt) + w = v.ewise_add(v, gb.binary.plus).new() + print("udt op: " + str(w[1].new().value.tolist())) + print("after udt op: " + gb.ss.config["jit_c_control"]) +""" + + +@pytest.mark.skipif("_IS_SSGB7") +@pytest.mark.parametrize( + ("form", "env_control", "kept"), + [ + pytest.param("default", None, False, id="default"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "run", True, id="env-run"), + pytest.param("gb.config.set", None, True, id="config-set-run"), + pytest.param("YAML file", None, True, id="yaml-run"), + # With nothing configured, a 'run' written later is the default again. + pytest.param("gb.ss.config", None, False, id="ss-config-run"), + # Written over a configured 'on', it changes the control, so it counts. + pytest.param("gb.ss.config", "on", True, id="env-on-then-ss-config-run"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "off", True, id="env-off"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "pause", True, id="env-pause"), + pytest.param("GRAPHBLAS_JIT_C_CONTROL", "load", True, id="env-load"), + ], +) +def test_udt_work_raises_only_the_default_run(form, env_control, kept, tmp_path): + """UDT work raises the default ``'run'`` to ``'on'``, and keeps a control the user set. + + A ``'run'`` set before init (``GRAPHBLAS_JIT_C_CONTROL``, ``gb.config.set`` + or a YAML config file) keeps processes from compiling, at the cost of UDT + ops running their Numba functions and C source failing to register. One + written to ``gb.ss.config`` later counts only when it changes the configured + control. ``'off'``, ``'pause'`` and ``'load'`` were kept before the default + was told apart. + """ + if not kept and not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; nothing raises to 'on'") + extra_env = {} + if form == "YAML file": + (tmp_path / "jit.yaml").write_text("jit_c_control: run\n") + extra_env["GRAPHBLAS_CONFIG"] = str(tmp_path) + script = f"form = {form!r}\n" + _UDT_WORK_PROBE + lines, report = _run_in_fresh_process(script, env_control, **extra_env) + control = env_control if form == "GRAPHBLAS_JIT_C_CONTROL" else "run" + if not kept: + assert lines["before"] == "run", report + assert lines["register_new"] == "ok", report + assert lines["after register_new"] == "on", report + expected = "on" + else: + # A library built without the JIT clamps 'load' down to 'run'. + expected = lines["before"] + assert expected in {control, "run"}, report + assert lines["register_new"].startswith("InvalidValue: Cannot create"), report + if control == "run": + assert "was set to 'run'" in lines["register_new"], report + assert lines["after register_new"] == expected, report + if numba is not None: + assert lines["udt op"] == "(6, 8)", report + assert lines["after udt op"] == expected, report + + @pytest.mark.skipif("_IS_SSGB7") def test_ss_register_new_says_the_compile_failed(): """C source that does not compile says so, and that ``jit_c_control`` is now ``'load'``.""" @@ -377,12 +622,18 @@ def test_enable_jit_for_udt_re_raises_a_lowered_control(monkeypatch): fake = {} monkeypatch.setattr(jit_config, "_ss_config", lambda: fake) monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", True) + monkeypatch.setattr(jit_config, "_user_set_run", False) # An explicit opt-out is a choice, not the default to raise, and ``'load'`` is # also where SuiteSparse parks the control after a compile fails. for start, expected in (("run", "on"), ("load", "load"), ("off", "off"), ("pause", "pause")): fake["jit_c_control"] = start assert _enable_jit_for_udt() is (expected == "on") assert fake["jit_c_control"] == expected + # A 'run' the user set is an opt-out too. + jit_config._user_set_run = True + fake["jit_c_control"] = "run" + assert _enable_jit_for_udt() is False + assert fake["jit_c_control"] == "run" @pytest.mark.skipif("_IS_SSGB7") @@ -390,9 +641,10 @@ def test_enable_jit_for_udt_does_not_cache_an_opt_out(monkeypatch): """An ``'off'`` at the first UDT op must not decide the rest of the process. It says nothing about whether this process can compile. Caching a "no" - for it would leave a saved ``'run'``, restored afterwards, unraised for + for it would leave the default ``'run'``, restored afterwards, unraised for good, so every later UDT op would run without the JIT and nothing would - say so. The same goes for ``'pause'`` and ``'load'``. + say so. The same goes for ``'pause'``, ``'load'``, and a ``'run'`` the user + chose. """ from graphblas.core.ss import jit_config @@ -401,19 +653,36 @@ def test_enable_jit_for_udt_does_not_cache_an_opt_out(monkeypatch): monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", None) monkeypatch.setattr(jit_config, "jit_compiler_is_usable", lambda: True) monkeypatch.setattr(jit_config, "_probe_jit", lambda cfg: True) - for start in ("off", "pause", "load"): + monkeypatch.setattr(jit_config, "_cache_is_writable", lambda cfg: True) + for start in ("off", "pause", "load", "run"): + # The last is a 'run' the user set. + monkeypatch.setattr(jit_config, "_user_set_run", start == "run") fake["jit_c_control"] = start assert _enable_jit_for_udt() is False assert fake["jit_c_control"] == start assert jit_config._jit_enabled_for_udt is None + monkeypatch.setattr(jit_config, "_user_set_run", False) + # A cache it cannot write to: the default 'run' stays, and the answer is cached. + fake["jit_c_control"] = "run" + with monkeypatch.context() as m: + m.setattr(jit_config, "_cache_is_writable", lambda cfg: False) + assert _enable_jit_for_udt() is False + assert fake["jit_c_control"] == "run" + assert jit_config._jit_enabled_for_udt is False + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", None) fake["jit_c_control"] = "run" assert _enable_jit_for_udt() is True assert fake["jit_c_control"] == "on" @pytest.mark.skipif("_IS_SSGB7") -def test_enable_jit_for_udt_re_raises_the_live_control(): - """The same raise, through SuiteSparse's own config rather than a stand-in.""" +def test_enable_jit_for_udt_re_raises_the_live_control(monkeypatch): + """The same raise, through SuiteSparse's own config rather than a stand-in. + + A ``'run'`` written to ``gb.ss.config`` is the default, and raised, while + nothing is configured; written over a configured ``'on'``, it is the user's, + and kept. None is the default either way. + """ if not _JIT_WORKS_AT_IMPORT: pytest.skip("JIT not compiled in or not working here; the control need not hold 'on'") @@ -423,9 +692,23 @@ def test_enable_jit_for_udt_re_raises_the_live_control(): prev_enabled = jit_config._jit_enabled_for_udt try: jit_config._jit_enabled_for_udt = True + monkeypatch.setattr(jit_config, "_configured_jit_c_control", None) + monkeypatch.setattr(jit_config, "_user_set_run", jit_config._user_set_run) gb.ss.config["jit_c_control"] = "run" assert _enable_jit_for_udt() is True assert gb.ss.config["jit_c_control"] == "on" + monkeypatch.setattr(jit_config, "_configured_jit_c_control", "on") + gb.ss.config["jit_c_control"] = "RUN" + assert jit_config._user_set_run is True + assert _enable_jit_for_udt() is False + assert gb.ss.config["jit_c_control"] == "run" + gb.ss.config["jit_c_control"] = None + assert gb.ss.config["jit_c_control"] == "run" + assert jit_config._user_set_run is False + assert _enable_jit_for_udt() is True + assert gb.ss.config["jit_c_control"] == "on" + gb.ss.config["jit_c_control"] = gb.core.lib.GxB_JIT_RUN + assert jit_config._user_set_run is True finally: gb.ss.config["jit_c_control"] = prev_control jit_config._jit_enabled_for_udt = prev_enabled @@ -465,7 +748,7 @@ def test_failed_compile_is_not_retried_at_the_next_arming(): prev_enabled = jit_config._jit_enabled_for_udt try: # Start from the default, so the first arming shows the raise. - gb.ss.config["jit_c_control"] = "run" + gb.ss.config["jit_c_control"] = None # A compiler that does not exist fails under every toolchain, where a # flag it does not know is only a warning to MSVC. The cache is pinned # to yes (true here) so the arming does not repair the path first. @@ -507,8 +790,8 @@ def test_public_fix_jit_config_repairs_a_broken_compiler(): assert result is True assert gb.ss.jit_compiler_is_usable() # Calling fix_jit_config a second time must not disable JIT. The probe - # tries to register a one-shot ``_jit_probe`` UDT and previously - # interpreted the duplicate-name error as a probe failure. + # used to register a one-shot ``_jit_probe`` UDT, and took the + # duplicate-name error as a probe failure. gb.ss.config["jit_c_compiler_name"] = "/nonexistent/path/to/cc" result2 = gb.ss.fix_jit_config() assert result2 is True @@ -2046,15 +2329,11 @@ def test_eq_ne_udt_cfunc_works_for_unrepresentable_dtypes(): @pytest.mark.skipif("not supports_udfs") -def test_udt_scalar_broadcast_skips_jit(): - """Mixed UDT+scalar arithmetic and eq/ne stay on the cfunc path. - - The JIT codegen helpers (``set_jit_c_on_op`` and - ``set_jit_c_comparison_on_op``) only run when ``dtype == dtype2``. A - mixed pair would need a per-side type in the C signature plus one C - name per (op, left_type, right_type) triple; for now we keep the - cfunc fallback for broadcast and skip JIT. This test pins that - routing so a future change doesn't accidentally regress it (either +def test_udt_scalar_broadcast_jit_routing(): + """Mixed UDT and scalar arithmetic gets a C definition with a type for each of z, x + and y; eq/ne on a mixed pair, and ops that have no C form, keep the cfunc. + + This test pins that routing so a future change doesn't accidentally regress it (either direction is fine, but it should be a conscious decision). """ if not _has_jit_set: @@ -2068,14 +2347,195 @@ def test_udt_scalar_broadcast_skips_jit(): # Same-type pairs do JIT. assert binary.plus[record, record].jit_c_source is not None assert binary.eq[record, record].jit_c_source is not None - # Mixed pairs do not JIT (cfunc handles broadcast correctly; this is - # the documented limit, not an oversight). - assert binary.plus[record, dtypes.INT64].jit_c_source is None - assert binary.plus[dtypes.INT64, record].jit_c_source is None + # Mixed arithmetic does, with the scalar read once for every leaf + assert binary.plus[record, dtypes.INT64].jit_c_source == ( + "void plus_11__BcastJitUV_int64_t (_BcastJitUV *z, const _BcastJitUV *x, " + "const int64_t *y) { z->bcast_u = (x->bcast_u) + (((double) ((*y)))) ; " + "z->bcast_v = (x->bcast_v) + (((double) ((*y)))) ; }" + ) + assert binary.plus[dtypes.INT64, record].jit_c_source is not None + # Comparisons on a mixed pair keep the cfunc assert binary.eq[record, dtypes.INT64].jit_c_source is None assert binary.ne[record, dtypes.FP64].jit_c_source is None +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_c_source(): + """The C definitions of lifted ops on a scalar and an array UDT, and the cases with no C + form (an array operand that broadcasts, integer truediv, an array leaf in a record). + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + arr = dtypes.register_anonymous(np.dtype((np.float64, (13,))), "_MixedArr13") + # (the C name of whichever test registered this numpy type first: ``test_op`` has it too) + c = arr.jit_c_name + assert binary.times[dtypes.FP64, arr].jit_c_source == ( + f"void times_6_double_{c} ({c} *z, const double *x, " + f"const {c} *y) {{ for (int64_t i = 0 ; i < 13 ; i++) " + "{ z->v[i] = ((*x)) * (y->v[i]) ; } }" + ) + assert binary.times[dtypes.FP64, arr].jit_c_name == f"times_6_double_{c}" + assert binary.times[arr, dtypes.FP64].jit_c_name == f"times_{len(c)}_{c}_double" + # plus_times on these types: its monoid is the same-type plus + semi = gb.semiring.plus_times[dtypes.FP64, arr] + assert semi.jit_c_source == binary.times[dtypes.FP64, arr].jit_c_source + assert semi.monoid.jit_c_source is not None + # Operands converted to the type the op computes in, as the cfunc converts them + assert "(((double) ((*x)))) - (y->v[i])" in binary.minus[dtypes.INT32, arr].jit_c_source + arr_i8 = dtypes.register_anonymous(np.dtype((np.int8, (13,))), "_MixedArrI8") + assert "(((double) (x->v[i]))) + (y->v[i])" in binary.plus[arr_i8, arr].jit_c_source + # No C form: these keep the cfunc + arr1 = dtypes.register_anonymous(np.dtype((np.float64, (1,))), "_MixedArr1") + assert binary.plus[arr1, arr].jit_c_source is None # (broadcasts) + assert binary.truediv[dtypes.INT64, arr_i8].jit_c_source is None + nested = dtypes.register_anonymous( + np.dtype([("mx_w", np.float64, (2,)), ("mx_c", np.int64)]), "_MixedArrLeaf" + ) + assert binary.plus[nested, dtypes.INT64].jit_c_source is None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_c_name_is_its_own(): + """A mixed op's C name is not a one-type op's on a UDT named like the pair, nor another + pair's. SuiteSparse finds a compiled kernel by its op's name, so the second op ran the + first's kernel: wrong values, or writes past the end of a smaller element. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + a = dtypes.register_anonymous(np.dtype((np.float64, (19,))), "_CnA") + b = dtypes.register_anonymous(np.dtype((np.int64, (19,))), "_CnB") + ab = dtypes.register_anonymous(np.dtype((np.int8, (19,))), f"{a.jit_c_name}_{b.jit_c_name}") + assert binary.plus[ab, ab].jit_c_name != binary.plus[a, b].jit_c_name + n = 50 + # Values that vary, or SuiteSparse before 10.5 builds iso vectors and computes an iso + # result without the kernel + k = np.broadcast_to(np.arange(n)[:, None] % 7, (n, 19)) + vab = Vector.from_coo(np.arange(n), k.astype(np.int8), dtype=ab) + va = Vector.from_coo(np.arange(n), k + 1.5, dtype=a) + vb = Vector.from_coo(np.arange(n), 10 * k, dtype=b) + with _jit_mode("on"): + same = vab.ewise_mult(vab, binary.plus).new() + mixed = va.ewise_mult(vb, binary.plus).new() + assert (same.to_dense() == 2 * k).all() + assert (mixed.to_dense() == 11 * k + 1.5).all() + + +def _mixed_parity_cases(): + rng = np.random.default_rng(17) + arr_f = dtypes.register_anonymous(np.dtype((np.float64, (5,))), "_MixedParF") + arr_i = dtypes.register_anonymous(np.dtype((np.int64, (5,))), "_MixedParI") + arr_c = dtypes.register_anonymous(np.dtype((np.complex128, (5,))), "_MixedParC") + rec = dtypes.register_anonymous( + np.dtype([("mp_a", np.float64), ("mp_b", np.int32)], align=True), "_MixedParRec" + ) + n = 64 + vf = rng.standard_normal((n, 5)) + vf[0, :2] = [np.nan, np.inf] + vi = rng.integers(-9, 10, (n, 5)) + vc = rng.standard_normal((n, 5)) + 1j * rng.standard_normal((n, 5)) + vr = np.zeros(n, rec.np_type) + vr["mp_a"] = rng.standard_normal(n) + vr["mp_b"] = rng.integers(-9, 10, n) + sf = rng.standard_normal(n) + sf[1] = 0.0 + si = rng.integers(-3, 4, n) + su = rng.integers(0, 2**63, n, dtype=np.uint64) + idx = np.arange(n) + u = { + "FP64[5]": Vector.from_coo(idx, vf, dtype=arr_f), + "INT64[5]": Vector.from_coo(idx, vi, dtype=arr_i), + "FC64[5]": Vector.from_coo(idx, vc, dtype=arr_c), + "record": Vector.from_coo(idx, vr, dtype=rec), + } + s = { + "FP64": Vector.from_coo(idx, sf), + "INT64": Vector.from_coo(idx, si), + "UINT64": Vector.from_coo(idx, su), + } + real_ops = ["plus", "minus", "times", "truediv", "floordiv", "min", "max"] + return ( + [ + ("FP64[5]", "FP64", real_ops), + ("FP64[5]", "INT64", real_ops), + ("INT64[5]", "UINT64", real_ops), # (computed in FP64, as on built-in vectors) + ("INT64[5]", "INT64", ["plus", "minus", "times", "floordiv", "min", "max"]), + ("FC64[5]", "FP64", ["plus", "minus", "times", "truediv"]), + ("record", "FP64", real_ops), + ("record", "INT64", ["plus", "minus", "times", "min", "max"]), + ], + u, + s, + ) + + +@pytest.mark.slow +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_cfunc_parity(): + """A mixed op's C kernel and its Numba cfunc give the same bytes, the UDT on either side.""" + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + cases, u, s = _mixed_parity_cases() + checked = 0 + for ukey, skey, ops in cases: + for op in ops: + for left, right in [(u[ukey], s[skey]), (s[skey], u[ukey])]: + typed = getattr(binary, op)[left.dtype, right.dtype] + if typed.jit_c_source is None: + continue + with _jit_mode("on"): + jit = left.ewise_mult(right, getattr(binary, op)).new() + with _jit_mode("off"): + cf = left.ewise_mult(right, getattr(binary, op)).new() + assert jit.dtype == cf.dtype + assert jit.nvals == cf.nvals + assert ( + jit.to_dense().tobytes() == cf.to_dense().tobytes() + ), f"binary.{op}[{left.dtype}, {right.dtype}]: JIT and cfunc disagree" + checked += 1 + assert checked > 60 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_semiring_runs_jit(capfd): + """The scalar x array semiring compiles as one JIT kernel (burble), with the result of the + generic kernel. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + arr = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_MixedSemi6") + rng = np.random.default_rng(23) + A = gb.Matrix.from_coo( + rng.integers(0, 30, 200), + rng.integers(0, 20, 200), + rng.random(200), + nrows=30, + ncols=20, + dup_op=binary.plus, + ) + U = rng.random((20, 6)) + u = Vector.from_coo(np.arange(20), U, dtype=arr) + semi = gb.semiring.plus_times[dtypes.FP64, arr] + capfd.readouterr() + with burble(): + w = A.mxv(u, semi).new() + w.wait() + out = capfd.readouterr().out + assert "GrB_mxv" in out + assert "generic" not in out + # (and burble names the generic kernel when it runs) + with _jit_mode("off"), burble(): + w_generic = A.mxv(u, semi).new() + w_generic.wait() + assert "generic" in capfd.readouterr().out + idx, vals = w.to_coo() + gidx, gvals = w_generic.to_coo() + assert_array_equal(idx, gidx) + np.testing.assert_allclose(vals, gvals, rtol=1e-13) + np.testing.assert_allclose(vals, (A.to_dense(0) @ U)[idx], rtol=1e-13) + + def test_context_importable(): if _IS_SSGB7: with pytest.raises(ImportError, match="Context was added"): @@ -2085,3 +2545,94 @@ def test_context_importable(): from graphblas.core.ss.context import global_context assert gb.ss.global_context is global_context + + +def test_fix_from_conda_finds_windows_compilers(tmp_path): + """The conda repair looks in ``bin`` and, as Windows conda installs them, ``Library/bin`` + with ``.exe`` names; a baked-in Windows path's basename is found too. Nothing there: None. + """ + from graphblas.core.ss.jit_config import _fix_from_conda + + def cfg_for(name): + return {"jit_c_compiler_name": name, "jit_c_compiler_flags": "-O3"} + + win = tmp_path / "win" + (win / "Library" / "bin").mkdir(parents=True) + (win / "Library" / "bin" / "clang.exe").write_text("") + cfg = cfg_for(r"C:\bld\_h_env\Library\bin\cl.exe") + assert _fix_from_conda(cfg, str(win)) is True + assert cfg["jit_c_compiler_name"] == str(win / "Library" / "bin" / "clang.exe") + (win / "Library" / "bin" / "cl.exe").write_text("") + cfg = cfg_for(r"C:\bld\_h_env\Library\bin\cl.exe") + assert _fix_from_conda(cfg, str(win)) is True + assert cfg["jit_c_compiler_name"] == str(win / "Library" / "bin" / "cl.exe") + posix = tmp_path / "posix" + (posix / "bin").mkdir(parents=True) + (posix / "bin" / "gcc").write_text("") + cfg = cfg_for("/home/conda/feedstock_root/build_artifacts/bin/x86_64-conda-linux-gnu-cc") + assert _fix_from_conda(cfg, str(posix)) is True + assert cfg["jit_c_compiler_name"] == str(posix / "bin" / "gcc") + assert _fix_from_conda(cfg_for("/nowhere/cc"), str(tmp_path / "empty")) is None + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_mixed_jit_c_name_fits_suitesparse(): + """A mixed op's C name from two long UDT names still fits SuiteSparse's 127 characters. + + SuiteSparse refuses a longer name and takes the definition under an empty one; the kernel + it compiles then calls no function and leaves its output uninitialized. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + rec = [("ln_x", np.float64), ("ln_y", np.float64)] + rec32 = [("ln_x", np.float32), ("ln_y", np.float32)] + a = dtypes.register_new("LongNameA" + "a" * 60, np.dtype(rec, align=True)) + b = dtypes.register_new("LongNameB" + "b" * 60, np.dtype(rec32, align=True)) + op = binary.plus[a, b] + assert op.jit_c_source is not None + assert len(op.jit_c_name) < gb.core.lib.GxB_MAX_NAME_LEN + n = 40 + va = np.zeros(n, a.np_type) + va["ln_x"] = np.arange(n) + va["ln_y"] = 2 * np.arange(n) + vb = np.zeros(n, b.np_type) + vb["ln_x"] = 0.5 + vb["ln_y"] = np.arange(n) % 3 + u = Vector.from_coo(np.arange(n), va, dtype=a) + w = Vector.from_coo(np.arange(n), vb, dtype=b) + with _jit_mode("on"): + r = u.ewise_mult(w, binary.plus).new() + _, vals = r.to_coo() + np.testing.assert_array_equal(vals["ln_x"], np.arange(n) + 0.5) + np.testing.assert_array_equal(vals["ln_y"], 2 * np.arange(n) + np.arange(n) % 3) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_jit_c_name_too_long_keeps_the_cfunc(): + """An op whose C name SuiteSparse refuses (128 characters or more) runs its cfunc. + + Setting the definition after a refused name would compile a kernel that calls no + function, so the op gets no C form at all. + """ + if not _has_jit_set: + pytest.skip("requires SuiteSparse:GraphBLAS >= 9") + _require_jit_on() + rec = np.dtype([("tl_x", np.float64), ("tl_y", np.int64)], align=True) + t = dtypes.register_new("TooLongName" + "t" * 120, rec) + n = 40 + vals = np.zeros(n, t.np_type) + vals["tl_x"] = np.arange(n) / 4 + vals["tl_y"] = np.arange(n) + v = Vector.from_coo(np.arange(n), vals, dtype=t) + plus_op = binary.plus[t] + eq_op = binary.eq[t] + assert plus_op.jit_c_source is None + assert eq_op.jit_c_source is None + with _jit_mode("on"): + r = v.ewise_mult(v, binary.plus).new() + same = v.ewise_mult(v, binary.eq).new() + _, got = r.to_coo() + np.testing.assert_array_equal(got["tl_x"], np.arange(n) / 2) + np.testing.assert_array_equal(got["tl_y"], 2 * np.arange(n)) + assert same.reduce(monoid.land).new() diff --git a/scripts/jit_diagnostics.py b/scripts/jit_diagnostics.py index 3c94f19bd..6959bfb2d 100755 --- a/scripts/jit_diagnostics.py +++ b/scripts/jit_diagnostics.py @@ -68,7 +68,7 @@ def main(): print("\nSuiteSparse:GraphBLAS 7.x: JIT not available") return 0 - print("\n--- GraphBLAS JIT defaults (from compiled C library) ---") + print("\n--- GraphBLAS JIT settings after graphblas init ---") for key in [ "jit_c_control", "jit_c_compiler_name",