diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index d43331257..d37653250 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -256,12 +256,13 @@ Anonymous UDTs (no ``name=`` argument) still take the JIT path: a synthetic identifier. Pass ``name=`` only for readable JIT cache filenames and introspection. -The first time auto-lift produces a non-JIT'd op for a given ``(op, dtype)`` -pair in a process, a ``graphblas.exceptions.NoJITWarning`` (a subclass of -``UserWarning``) is emitted with the cause and the remediation. The warning -fires once per ``(op, dtype)`` rather than once per process, so distinct -fallback causes each surface a warning. Silence it by category or by -message:: +When the cause is the UDT itself (the first four cases above), the first time +auto-lift produces a non-JIT'd op for a given ``(op, dtype)`` pair in a +process, a ``graphblas.exceptions.NoJITWarning`` (a subclass of +``UserWarning``) names the cause. An unusable compiler does not warn; +``gb.ss.fix_jit_config()`` returns ``False`` when it cannot compile, and +``gb.ss.jit_compiler_is_usable()`` is the cheap check. Silence the warning by +category or by message:: import warnings from graphblas.exceptions import NoJITWarning @@ -281,16 +282,25 @@ machines. With the bogus default, SuiteSparse emits the JIT ``.c`` source 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 auto-fixes this at import. 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), and ``jit_c_control`` is -bumped from the SS default ``'run'`` (run cached kernels only; no compile, -no load from disk) to ``'on'`` (compile, load, and run). When the default -config is already valid, only the mode bump applies. +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 +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) when +a UDT operator gets C source for SuiteSparse to compile, such as the first +``binary.plus[udt]``, or when ``gb.dtypes.ss.register_new`` registers a type +from a C typedef. Any other setting 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 UDT operation. Call the helper manually to re-fix or verify:: - gb.ss.fix_jit_config() # repair compiler path (full probe) + gb.ss.fix_jit_config() # repair compiler path, set 'on', probe gb.ss.jit_compiler_is_usable() # cheap check: True iff path exists Pickle and serialize diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index b55a17f13..ede25b467 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -956,7 +956,14 @@ def _set_jit_c_strings(gb_obj, c_name, c_defn, set_string_func): Both strings are one-shot on SuiteSparse: a second set returns ``GrB_ALREADY_SET`` silently. The return is not checked here because callers always pass a freshly-allocated handle. + + Arming an op with C source is the point at which this process has + committed to wanting a JIT compiler, so it is also where SuiteSparse's + non-compiling default gets raised. """ + 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) set_string_func(gb_obj, ffi.new("char[]", c_defn.encode()), lib.GxB_JIT_C_DEFINITION) diff --git a/graphblas/core/ss/dtypes.py b/graphblas/core/ss/dtypes.py index 1b0d93506..eeca52d0b 100644 --- a/graphblas/core/ss/dtypes.py +++ b/graphblas/core/ss/dtypes.py @@ -40,6 +40,12 @@ def register_new(name, jit_c_definition, *, np_type=None): if "struct" not in jit_c_definition: raise ValueError("Only struct typedefs are currently allowed for JIT dtypes") + # Registering a type by its C typedef is a request to compile C, so it is + # a fair place to raise SuiteSparse's non-compiling default. + from .jit_config import _enable_jit_for_udt + + _enable_jit_for_udt() + gb_obj = ffi.new("GrB_Type*") status = lib.GxB_Type_new( gb_obj, 0, ffi_new("char[]", name.encode()), ffi_new("char[]", jit_c_definition.encode()) diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index f77b55651..8cdc19e25 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -62,8 +62,8 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): Replaces the baked-in compiler path (which often points at a conda-build host that doesn't exist in user environments) with one from - ``$CONDA_PREFIX/bin/``, and strips build-time-only flags (``-isysroot``, - ``-fdebug-prefix-map``). + ``$CONDA_PREFIX/bin/``, strips build-time-only flags (``-isysroot``, + ``-fdebug-prefix-map``), and sets ``jit_c_control`` to ``'on'``. Parameters ---------- @@ -73,10 +73,10 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): 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. SuiteSparse auto-flips ``jit_c_control`` - from ``'on'`` to ``'load'`` on a failed compile; the probe absorbs - that first failure so user-visible ops afterwards see a stable - ``'load'`` (cache-only) state and punt to generic cleanly. + 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. Returns ------- @@ -84,30 +84,47 @@ def fix_jit_config(*, use_sysconfig=True, probe=True): Fix applied and (if ``probe``) verified working. False Fix attempted but the probe failed. ``jit_c_control`` is now - whatever SuiteSparse left it at (typically ``'load'``). + whatever SuiteSparse left it at. None No environment available to fix from. There's no ``$CONDA_PREFIX``, and either ``use_sysconfig=False`` or no sysconfig compiler is set. """ cfg = _ss_config() - if conda_prefix := os.environ.get("CONDA_PREFIX", ""): - rv = _fix_from_conda(cfg, conda_prefix) - elif use_sysconfig: - rv = _fix_from_sysconfig(cfg) - else: - return None + 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``. + 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 # *still* falls back to cfunc (different cause: UDT layout, etc.) # deserves a fresh notification rather than silent suppression. _warned_no_jit_for.clear() + # Re-arm the cached probe answer for the same reason: a pre-repair False + # is stale now, so let the next UDT op re-derive it against the fixed + # toolchain instead of trusting the old verdict. + global _jit_enabled_for_udt + _jit_enabled_for_udt = None if not probe: return True return _probe_jit(cfg) +def _repair_jit_compiler(cfg, *, use_sysconfig=True): + """Point the JIT compiler settings at something that exists on this host. + + Only rewrites the compiler name and flags; ``jit_c_control`` is left + alone, so on its own this changes nothing about what SuiteSparse will + do. Callers that want compilation enabled set the control themselves. + """ + if conda_prefix := os.environ.get("CONDA_PREFIX", ""): + return _fix_from_conda(cfg, conda_prefix) + if use_sysconfig: + return _fix_from_sysconfig(cfg) + return None + + def _fix_from_conda(cfg, conda_prefix): """Conda-aware fix; swap the compiler path with one from $CONDA_PREFIX/bin/.""" jit_cc = cfg["jit_c_compiler_name"] @@ -122,7 +139,6 @@ def _fix_from_conda(cfg, conda_prefix): else: return None # nothing usable _fix_compiler_flags(cfg) - cfg["jit_c_control"] = "on" return True @@ -139,7 +155,6 @@ def _fix_from_sysconfig(cfg): cfg["jit_c_compiler_flags"] = f"{cflags} -I{include}" if libs := sysconfig.get_config_var("LIBS"): cfg["jit_c_libraries"] = libs - cfg["jit_c_control"] = "on" return True @@ -182,11 +197,13 @@ def _strip_mismatched_arch(flags): def _probe_jit(cfg): """Probe a trivial JIT compile to verify the config works. - On failure, SuiteSparse will have flipped ``jit_c_control`` from - ``'on'`` to ``'load'`` (its built-in response to a failed compile); - we leave that state alone. Both ``'load'`` and ``'off'`` cause - downstream ops to punt to the generic kernel, but ``'load'`` preserves - any pre-compiled kernels in the cache. + ``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. """ from ... import dtypes as _dtypes @@ -197,6 +214,8 @@ def _probe_jit(cfg): 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: @@ -205,34 +224,94 @@ def _probe_jit(cfg): # validation ``ValueError``s. The probe's contract is "did this # work?", so absorb every failure mode here. return False + finally: + _probing_jit = False return True def _auto_fix_jit_at_import(): - """Run :func:`fix_jit_config` at ``gb.ss`` import; designed not to raise. - - Called unguarded from ``graphblas/ss/__init__.py``, so any exception - here breaks ``import graphblas.ss``. The body sticks to dict ops and - delegates the failure-prone work to ``_probe_jit``, which catches - everything internally. - - The probe is the load-bearing piece: without it, SS would surface - ``JitError`` on the first user-triggered JIT compile (a failed - compile is only converted to a silent ``'load'`` fallback on - subsequent calls). + """Repair the JIT compiler path at ``gb.ss`` import; designed not to raise. + + Called unguarded from ``graphblas/ss/__init__.py``, so any exception here + 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. """ cfg = _ss_config() if "jit_c_control" not in cfg: return - if jit_compiler_is_usable(): - if cfg["jit_c_control"] in ("run", "load"): - cfg["jit_c_control"] = "on" - else: - fix_jit_config(use_sysconfig=True, probe=False) - if jit_compiler_is_usable() and cfg["jit_c_control"] in ("run", "load"): + if not jit_compiler_is_usable(): + _repair_jit_compiler(cfg) + + +# Tri-state: ``None`` until a UDT first asks for a JIT kernel 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 + + +def _enable_jit_for_udt(): + """Enable JIT compilation the first time a UDT needs a kernel built. + + 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). + + This is where that bump belongs, rather than at import: registering a + type from its C typedef 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. + """ + 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. + 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(): cfg["jit_c_control"] = "on" - if cfg.get("jit_c_control") == "on": - _probe_jit(cfg) + # A library built without the JIT clamps that write back down to + # ``'run'``. Otherwise the probe is load-bearing: a failed compile + # drops the control to ``'load'`` and, from SuiteSparse 9.4 on, + # also raises ``JitError``, so without the probe that one error + # would reach the first user operation that needs a kernel. + _jit_enabled_for_udt = cfg["jit_c_control"] == "on" and _probe_jit(cfg) + elif _jit_enabled_for_udt and control == "run": + # The cache answers "can this process compile?", which is settled once. + # The control is separate state that anything may have lowered since, + # and a kernel armed now still needs it raised. + cfg["jit_c_control"] = "on" + return _jit_enabled_for_udt # Keyed by ``(op_name, dtype_name)`` so each distinct pair warns once. @@ -244,8 +323,15 @@ def _auto_fix_jit_at_import(): def _maybe_warn_no_jit(*, op_name="", dtype_name=""): """Emit a ``NoJITWarning`` (once per ``(op_name, dtype_name)``) when UDT auto-lift falls back. - The most likely cause (bogus compiler path, ``jit_c_control`` off, or - UDT not C-expressible) is named in the message along with the remediation. + The only caller is ``udt_utils._maybe_warn_jit_skipped``, which fires + after C codegen returned nothing. Codegen reads the dtype and neither the + compiler path nor ``jit_c_control``, so the dtype is the cause here even + on a host that could not have compiled a kernel anyway. Do not re-derive + the cause from the live config; that is not the state codegen read. + + The introspection properties (``DataType.jit_c_definition`` and + ``TypedUserBinaryOp.jit_c_source``) show what was generated, or ``None`` + when codegen was skipped. """ key = (op_name, dtype_name) if key in _warned_no_jit_for: @@ -253,41 +339,20 @@ def _maybe_warn_no_jit(*, op_name="", dtype_name=""): _warned_no_jit_for.add(key) import warnings as _warnings - cfg = _ss_config() - if not jit_compiler_is_usable(): - cause = ( - "the JIT compiler path is not usable " - f"({cfg.get('jit_c_compiler_name', '')!r}); " - "call ``gb.ss.fix_jit_config()`` to repair it" - ) - elif cfg.get("jit_c_control") != "on": - cause = ( - f"jit_c_control is {cfg.get('jit_c_control')!r} (must be 'on' to compile); " - "set ``gb.ss.config['jit_c_control'] = 'on'`` to enable compilation" - ) - else: - # Compiler is usable and mode is OK; the UDT itself isn't C-expressible. - # The introspection properties (``DataType.jit_c_definition`` and - # ``TypedUserBinaryOp.jit_c_source``) show what was generated, or - # ``None`` when codegen was skipped. - loc = f" (op={op_name!r}, dtype={dtype_name!r})" if op_name else "" - cause = ( - "this UDT is not expressible as a C struct" - f"{loc} (a field name is a C reserved word or stdlib macro, a " - "field type isn't in the numpy-to-C map, a field is array-typed, " - "or the record has a packed layout). The op still works via the " - "Numba cfunc path; only the JIT speedup is lost" - ) from ...exceptions import NoJITWarning + loc = f" (op={op_name!r}, dtype={dtype_name!r})" if op_name else "" _warnings.warn( - f"UDT operator running without JIT compilation: {cause}. " - f"Operations will use the Numba function-pointer fallback " - f"(typically 2-3x slower for elementwise ops, since SuiteSparse " - f"can't inline the kernel into its eWise and reduce templates). " - f"This warning fires once per (op, dtype) per process; silence with " - f"``warnings.filterwarnings('ignore', category=gb.exceptions.NoJITWarning)`` " - f"or by message match.", + f"UDT operator running without JIT compilation: this UDT is not " + f"expressible as a C struct{loc} (a field name is a C reserved word or " + f"stdlib macro, a field type isn't in the numpy-to-C map, a field is " + f"array-typed, or the record has a packed layout). The op still works " + f"through the Numba function-pointer fallback; only the JIT speedup is " + f"lost (typically 2-3x slower for elementwise ops, since SuiteSparse " + f"can't inline the kernel into its eWise and reduce templates). This " + f"warning fires once per (op, dtype) per process; silence with " + f"``warnings.filterwarnings('ignore', " + f"category=gb.exceptions.NoJITWarning)`` or by message match.", NoJITWarning, stacklevel=3, ) diff --git a/graphblas/exceptions.py b/graphblas/exceptions.py index b78023db3..df674b55e 100644 --- a/graphblas/exceptions.py +++ b/graphblas/exceptions.py @@ -109,10 +109,9 @@ class NoJITWarning(UserWarning): Emitted from :func:`graphblas.core.ss.jit_config._maybe_warn_no_jit` when an op like ``binary.plus[udt]`` would otherwise have JIT-compiled, - but couldn't because the JIT compiler is unusable, ``jit_c_control`` - is off, or the UDT isn't expressible as a C struct. Fires once per - ``(op, dtype)`` pair per process, so a user who registers several - UDTs gets one warning per pair regardless of cause. + but the UDT has no C form, so no kernel source was generated. Fires once + per ``(op, dtype)`` pair per process, so a user who registers several UDTs + gets one warning per pair. Inherits from :class:`UserWarning` so existing ``UserWarning`` filters still match it; users can also filter by category for a tight scope: diff --git a/graphblas/ss/__init__.py b/graphblas/ss/__init__.py index 394dcc391..33a342373 100644 --- a/graphblas/ss/__init__.py +++ b/graphblas/ss/__init__.py @@ -11,12 +11,14 @@ from ..core.ss.jit_config import _auto_fix_jit_at_import as _auto_fix from ..core.ss.jit_config import fix_jit_config, jit_compiler_is_usable - # Auto-fix the JIT config 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). When the compiler is already usable, the - # call only bumps ``jit_c_control`` from SS's default ``run`` to ``on`` so - # first-time UDT auto-lift actually JIT-compiles. Skipped on vanilla: the - # auto-fix reads ``gb.ss.config``, which needs GxB callables vanilla strips. + # 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. if backend == "suitesparse": _auto_fix() del _auto_fix diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 2e3e17d87..a9550ff51 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -461,8 +461,12 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): # this test sees the warning regardless of test ordering (the same # ``(op, dtype)`` may have been warned about in an earlier test). jit_config._warned_no_jit_for.discard(("plus", udt.name)) - with pytest.warns(NoJITWarning, match="without JIT"): + 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. + 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 # comparisons), so pin its skip independently. @@ -478,6 +482,39 @@ def test_udt_with_c_reserved_field_name_falls_back_to_cfunc(): assert tuple(result[1].new().value) == (202, 2020) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not _has_jit_set") +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 + 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``. + """ + from graphblas.core.ss import jit_config + from graphblas.exceptions import NoJITWarning + + # ``switch`` is a C reserved word; the field names are unique to this test + # because ``register_anonymous`` caches one DataType per ``np.dtype``. + record_dtype = np.dtype([("switch", np.int64), ("plain", np.int64)], align=True) + udt = dtypes.register_anonymous(record_dtype, "_ProbeFailedUdt") + + prev_control = gb.ss.config["jit_c_control"] + prev_enabled = jit_config._jit_enabled_for_udt + try: + gb.ss.config["jit_c_control"] = "run" + jit_config._jit_enabled_for_udt = False + jit_config._warned_no_jit_for.discard(("plus", udt.name)) + with pytest.warns(NoJITWarning, match="not expressible as a C struct") as record: + gb.binary.plus[udt] + finally: + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + assert not any("jit_c_control" in str(w.message) for w in record) + + @pytest.mark.skipif("not _has_jit_set") def test_udt_jit_c_info_pinned_at_first_register(): """SuiteSparse's ``GxB_JIT_C_DEFINITION`` is one-shot per ``GrB_Type``. diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index fc174c768..34e4767d6 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -1,6 +1,9 @@ import contextlib import os +import subprocess +import sys import sysconfig +from pathlib import Path import numpy as np import pytest @@ -23,6 +26,7 @@ from graphblas.core.operator.indexbinary import _has_idxbinop from graphblas.core.operator.udt_utils import _has_jit_set from graphblas.core.ss import _IS_SSGB7 +from graphblas.core.ss.jit_config import _enable_jit_for_udt from .conftest import autocompute, burble @@ -41,13 +45,12 @@ # 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) -# Capture the post-import JIT state before the autouse ``_setup_jit`` fixture -# runs. ``_auto_fix_jit_at_import`` probes once; ``jit_c_control`` is ``'on'`` -# iff the env can actually JIT-compile. Tests that assert on the import-time -# state read this constant; the fixture may transiently mutate the live config. -_JIT_WORKS_AT_IMPORT = ( - not _IS_SSGB7 and backend == "suitesparse" and gb.ss.config["jit_c_control"] == "on" -) +# ``_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() @pytest.fixture(scope="module", autouse=True) @@ -57,9 +60,8 @@ def _setup_jit(): Strategy: 1. _fix_jit_config(): fix conda-baked compiler paths and probe. - Returns True: JIT works, proceed. - - Returns False: probe failed. SuiteSparse will have left - ``jit_c_control = 'load'`` (compile disabled, cache loading - still allowed); we leave that state untouched. + - Returns False: probe failed. Whatever ``jit_c_control`` + SuiteSparse was left at is left untouched. - Returns None: no conda env, try sysconfig instead. 2. Sysconfig fallback: for non-conda installs (pure pip). """ @@ -71,10 +73,9 @@ def _setup_jit(): prev = gb.ss.config["jit_c_control"] result = _fix_jit_config() - # ``True``: conda JIT configured and verified. ``False``: probe failed - # (``_probe_jit`` leaves ``jit_c_control`` at ``'load'``; don't try - # sysconfig, since if the conda compiler can't build GraphBLAS JIT - # kernels, Python's sysconfig compiler won't either). Both done. + # ``True``: conda JIT configured and verified. ``False``: probe failed; + # don't try sysconfig, since if the conda compiler can't build GraphBLAS + # JIT kernels, Python's sysconfig compiler won't either. Both done. if result is None: # No conda env. Try sysconfig for non-conda installs. cc = sysconfig.get_config_var("CC") @@ -111,26 +112,233 @@ def _jit_mode(mode): gb.ss.config["jit_c_control"] = prev +def _run_in_fresh_process(script): + """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. + """ + repo_root = Path(__file__).resolve().parents[2] + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + env=dict(os.environ, PYTHONPATH=str(repo_root)), + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + return lines, report + + +# Run in a subprocess: the import under test happens once per process, and by +# the time any test executes it is long past. ``jit_c_control`` has to be read +# through the C API rather than ``gb.ss.config``, because reaching for +# ``gb.ss`` is itself the thing under test. +_SS_IMPORT_PROBE = """ +import graphblas as gb + +gb.Vector # initialize the backend, but not gb.ss +from graphblas.core import ffi, lib + +def control(): + val_ptr = ffi.new("int32_t*") + assert lib.GxB_Global_Option_get_INT32(lib.GxB_JIT_C_CONTROL, val_ptr) == lib.GrB_SUCCESS + return val_ptr[0] + +before = control() +gb.ss.about["library_version"] # the line a user writes for a bug report +print("package: " + gb.__file__) +print("before: " + str(before)) +print("after: " + str(control())) +""" + + @pytest.mark.skipif("_IS_SSGB7") -def test_auto_fix_jit_at_import_left_compiler_usable(): - """After ``import graphblas.ss``, a probe-confirmed compiler implies - ``jit_c_control == 'on'``. +def test_ss_import_leaves_jit_c_control_alone(): + """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. """ + lines, report = _run_in_fresh_process(_SS_IMPORT_PROBE) + assert lines["after"] == lines["before"], report + + +# Run in a subprocess: this is about the very first call in a process, and by +# the time any test executes, plenty of JIT setup has already happened. +_FIX_JIT_TWICE_PROBE = """ +import graphblas as gb + +print("package: " + gb.__file__) +print("first: " + str(gb.ss.fix_jit_config())) +print("second: " + str(gb.ss.fix_jit_config())) +""" + + +@pytest.mark.skipif("_IS_SSGB7") +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. + """ + lines, report = _run_in_fresh_process(_FIX_JIT_TWICE_PROBE) + assert lines["first"] == lines["second"], report + + +@pytest.mark.skipif("_IS_SSGB7") +def test_auto_fix_jit_at_import_left_compiler_usable(): + """The import-time repair leaves a compiler the UDT path can actually use.""" if not gb.ss.jit_compiler_is_usable(): pytest.skip("sandboxed env without a usable compiler; auto-fix had nothing to repair") if not _JIT_WORKS_AT_IMPORT: - # Compiler file exists but the import-time probe failed (e.g., the - # baked-in flags target a different arch than the host). After a - # compile failure SuiteSparse drops ``jit_c_control`` to a - # non-compiling mode so downstream ops punt to generic cleanly; the - # probe absorbs that first failure. Which mode it lands in (``'load'`` - # or ``'run'``) varies by SuiteSparse version. + # Compiler file exists but the probe failed (e.g., the baked-in flags + # target a different arch than the host). After a compile failure + # SuiteSparse drops ``jit_c_control`` to a non-compiling mode so + # downstream ops punt to generic cleanly; the probe absorbs that first + # failure. Which mode it lands in varies by SuiteSparse version. assert gb.ss.config["jit_c_control"] in {"load", "run"} pytest.skip("compiler present but JIT probe failed in this env") assert _JIT_WORKS_AT_IMPORT assert gb.ss.config["jit_c_control"] == "on" +@pytest.mark.skipif("_IS_SSGB7") +def test_enable_jit_for_udt_re_raises_a_lowered_control(monkeypatch): + """A cached "yes" must still raise a lowered ``'run'`` before arming a kernel. + + The cache answers "can this process compile?", which is settled once. The + control is separate state that anything may have moved since, so returning + the cached answer without re-raising it would leave an op armed with C + source that SuiteSparse never builds, and nothing would warn: codegen + succeeded. + + Driven through a stand-in config, so the assertion holds on every host: a + library built without the JIT clamps writes down to ``'run'``, and would + never read back the ``'on'`` this is about. The live round trip is pinned + separately below, where the clamp is not in the way. + """ + from graphblas.core.ss import jit_config + + fake = {} + monkeypatch.setattr(jit_config, "_ss_config", lambda: fake) + monkeypatch.setattr(jit_config, "_jit_enabled_for_udt", True) + # 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 + + +@pytest.mark.skipif("_IS_SSGB7") +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 + good, so every later UDT op would run without the JIT and nothing would + say so. The same goes for ``'pause'`` and ``'load'``. + """ + from graphblas.core.ss import jit_config + + fake = {} + monkeypatch.setattr(jit_config, "_ss_config", lambda: fake) + 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"): + 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 + 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.""" + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; the control need not hold 'on'") + + from graphblas.core.ss import jit_config + + prev_control = gb.ss.config["jit_c_control"] + prev_enabled = jit_config._jit_enabled_for_udt + try: + jit_config._jit_enabled_for_udt = True + gb.ss.config["jit_c_control"] = "run" + assert _enable_jit_for_udt() is True + assert gb.ss.config["jit_c_control"] == "on" + finally: + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + + +@pytest.mark.skipif("_IS_SSGB7") +@pytest.mark.skipif("not supports_udfs") +def test_failed_compile_is_not_retried_at_the_next_arming(): + """A compile that fails is tried once, not again at every later UDT op. + + SuiteSparse lowers ``jit_c_control`` to ``'load'`` after a failed compile + so the failure is not repeated. Arming another op must leave it there: + raising it back to ``'on'`` would run the failing compiler again at the + next kernel, and on SuiteSparse 9.4 and later raise ``JitError`` again, + with ``'off'`` the only way out. + """ + if not _JIT_WORKS_AT_IMPORT: + pytest.skip("JIT not compiled in or not working here; nothing compiles to fail") + from graphblas.core.ss import jit_config + from graphblas.exceptions import JitError + + # Field names unique to this test: ``register_anonymous`` caches one + # DataType per ``np.dtype``. + udt = dtypes.register_anonymous( + np.dtype([("nrt_a", np.int64), ("nrt_b", np.int64)], align=True), "_NoRetryUdt" + ) + v = Vector.from_coo([0, 1], np.array([(1, 2), (3, 4)], dtype=udt.np_type), dtype=udt) + # Before 9.4 SuiteSparse does not report a failed compile: it returns + # ``GrB_NO_VALUE`` ("use the generic kernel"), though it still lowers the + # control. + reported = gb.ss.about["library_version"] >= (9, 4) + prev_cc = gb.ss.config["jit_c_compiler_name"] + prev_control = gb.ss.config["jit_c_control"] + 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" + # 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. + gb.ss.config["jit_c_compiler_name"] = "/nonexistent/path/to/cc" + jit_config._jit_enabled_for_udt = True + plus = binary.plus[udt] + assert plus.jit_c_source is not None + assert gb.ss.config["jit_c_control"] == "on" + with pytest.raises(JitError) if reported else contextlib.nullcontext(): + v.ewise_add(v, plus).new() + assert gb.ss.config["jit_c_control"] == "load" + assert binary.minus[udt].jit_c_source is not None + assert gb.ss.config["jit_c_control"] == "load" + result = v.ewise_add(v, plus).new() # the generic kernel, without raising + assert tuple(result[1].new().value) == (6, 8) + finally: + gb.ss.config["jit_c_compiler_name"] = prev_cc + gb.ss.config["jit_c_control"] = prev_control + jit_config._jit_enabled_for_udt = prev_enabled + + @pytest.mark.skipif("_IS_SSGB7") def test_public_fix_jit_config_repairs_a_broken_compiler(): """``gb.ss.fix_jit_config()`` repairs a clobbered compiler path and returns ``True``."""