From 5f48083370bd0525431382384e3e291aa396500c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 21 Oct 2026 13:08:06 -0500 Subject: [PATCH 01/11] ss.fix_jit_config: find a conda compiler on Windows (Library/bin, .exe names) The conda repair looked for {baked-in name, cc, clang, gcc} in $CONDA_PREFIX/bin only. Conda on Windows installs compilers in %CONDA_PREFIX%\Library\bin, and their names end in .exe, so the repair found nothing there and left SuiteSparse's build-time compiler, which a user's machine does not have. It now looks in both places, with and without .exe, and takes a Windows path's basename with either slash. Tested with a fake prefix (no Windows machine at hand). --- graphblas/core/ss/jit_config.py | 17 +++++++++++------ graphblas/tests/test_ssjit.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/graphblas/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 677b979d8..5a044cd39 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -127,14 +127,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: diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 9e083a2a1..d984e7cf2 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -2085,3 +2085,31 @@ 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 From 7f0c7fe2e7bbf252955eb62397ed8b8cd5e9e1f2 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 21 Oct 2026 03:21:31 -0500 Subject: [PATCH 02/11] JIT C definitions for lifted ops on a scalar and a UDT, or on two UDT types A lifted binary op (plus, minus, times, truediv, floordiv, min, max) whose operands are a scalar and an array or record UDT, in either order, or two UDTs of one layout with other element types, ran only through its Numba cfunc: the C definition was set for same-type pairs alone, with one C type for x, y and z. So a semiring such as plus_times[FP64, FP64[R]] ran SuiteSparse's generic kernel with a function call per multiply. Such an op now gets a C definition with a type for each of z, x and y (udt_codegen._make_jit_c_mixed_definition), named op_N_x_y with N the length of x's C name, so that no two ops share a name (SuiteSparse finds a compiled kernel in memory by its op's name): a scalar operand is read as (*x) for every element or leaf, and each operand is converted in C to the type the op computes in, as the cfunc converts it (udt._udt_op_types). Ops with no such form keep the cfunc: an array operand that broadcasts, an element not computed in the result's type (integer truediv), a record with an array leaf, and eq/ne on a mixed pair. The whole mxv with plus_times[FP64, FP64[R]] now compiles as one JIT kernel (burble: dot2, jit). test_udt_scalar_broadcast_skips_jit becomes test_udt_scalar_broadcast_jit_routing: mixed arithmetic now JITs, mixed eq/ne still does not. --- graphblas/core/operator/binary.py | 29 +++- graphblas/core/operator/udt_codegen.py | 105 +++++++++++++ graphblas/tests/test_ssjit.py | 198 +++++++++++++++++++++++-- 3 files changed, 312 insertions(+), 20 deletions(-) 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..e866348aa 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -1419,6 +1419,97 @@ 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): + if leaf.shape: + 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}" + 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]``. @@ -1520,6 +1611,20 @@ def set_jit_c_comparison_on_op(gb_obj, op_name, dtype, set_string_func, *, is_eq 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 + _set_jit_c_strings(gb_obj, *result, lib.GrB_BinaryOp_set_String) + return result + + def set_jit_c_on_op(gb_obj, op_name, py_op, dtype, set_string_func, arity=2): """Generate and set JIT C name+definition on a GrB operator, if possible. diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index d984e7cf2..3ae667fa4 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -2046,15 +2046,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 +2064,190 @@ 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 + vab = Vector.from_coo(np.arange(n), np.ones((n, 19), np.int8), dtype=ab) + va = Vector.from_coo(np.arange(n), np.full((n, 19), 1.5), dtype=a) + vb = Vector.from_coo(np.arange(n), np.full((n, 19), 10), 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).all() + assert (mixed.to_dense() == 11.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 "(jit: " in out + assert "generic" not in out + with _jit_mode("off"): + w_generic = A.mxv(u, semi).new() + 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"): From b2ea4109bf8ef64b9ab01283e64d75ea21c37c01 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 21 Oct 2026 03:33:05 -0500 Subject: [PATCH 03/11] Test the mixed semiring's JIT kernel by burble's generic, absent and present The burble names a JIT kernel only when it compiles or loads it, so a kernel loaded earlier in the process prints no '(jit: ...)'. A UDT product has no kernel but the JIT one and the generic one, which burble names, so the test checks that 'generic' is absent with the JIT on and present with it off. --- graphblas/tests/test_ssjit.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 3ae667fa4..72ce02c3a 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -2237,10 +2237,12 @@ def test_udt_mixed_semiring_runs_jit(capfd): w.wait() out = capfd.readouterr().out assert "GrB_mxv" in out - assert "(jit: " in out assert "generic" not in out - with _jit_mode("off"): + # (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) From c4de50b63c4312d6371d17728d51a7c90e196a14 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 00:57:24 -0500 Subject: [PATCH 04/11] Vary the mixed JIT C name test's inputs, so its kernel runs on every SuiteSparse test_udt_mixed_jit_c_name_is_its_own built its vectors from constant values. SuiteSparse before 10.5 makes a vector built from equal values iso, and an ewise_mult of iso vectors gives an iso result computed without the JIT kernel the test is about, so on those versions it checked nothing. On them its to_dense also raised "TypeError: unsupported operand type(s) for +: 'int' and 'tuple'" (wave 2 CI, 17 jobs, e.g. 122633175811 and 122633175893), the iso array-UDT export bug fixed in #591. The values now differ by index. --- graphblas/tests/test_ssjit.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 72ce02c3a..e2eaf5ad0 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -2124,14 +2124,17 @@ def test_udt_mixed_jit_c_name_is_its_own(): 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 - vab = Vector.from_coo(np.arange(n), np.ones((n, 19), np.int8), dtype=ab) - va = Vector.from_coo(np.arange(n), np.full((n, 19), 1.5), dtype=a) - vb = Vector.from_coo(np.arange(n), np.full((n, 19), 10), dtype=b) + # 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).all() - assert (mixed.to_dense() == 11.5).all() + assert (same.to_dense() == 2 * k).all() + assert (mixed.to_dense() == 11 * k + 1.5).all() def _mixed_parity_cases(): From caf6dfc32714fcb89c6417335ce7d7450b7badfd Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 21:25:57 -0500 Subject: [PATCH 05/11] Set the SuiteSparse JIT control at init, from a new config key Behavior change. Before: python-graphblas left jit_c_control where the SuiteSparse:GraphBLAS build started it ("run" in the conda-forge and PyPI builds, "on" in upstream's own), and UDT work raised "run" to "on" when it needed C compiled, whoever had set it. Turning the JIT on for everything meant setting gb.ss.config["jit_c_control"] after init. After: when the "suitesparse" backend initializes, it sets jit_c_control from the new config key graphblas.config["jit_c_control"]. The key is unset by default, which sets "run": nothing changes for the conda-forge and PyPI builds, and a build that starts at "on" now starts at "run" too, so every build has one documented default. GRAPHBLAS_JIT_C_CONTROL, gb.config.set before init, or a YAML file in ~/.config/graphblas/ set it; the values are "on", "load", "run", "pause" and "off" (YAML's unquoted on/off arrive as bools and are taken as such), and any other value raises ValueError before anything is initialized. UDT work still raises the default "run" to "on" when it needs C compiled, but now keeps a "run" the user set, as it already kept "off", "pause" and "load". A "run" set before init counts: GRAPHBLAS_JIT_C_CONTROL, gb.config.set or a YAML file, one key that donfig merges, so they could not be told apart and need not be; the environment variable is also the form that reaches every worker of a pool or cluster, which is what a "run" is for (GraphBLAS #449, below). A later gb.ss.config["jit_c_control"] = "run" counts when it changes the configured control (a "run" over a configured "on", say), and keeps a configured "run"; with nothing configured it is the default "run" again. The comparison is with the configured control, not the live one, which UDT work may have raised to "on": code that saves the default "run", sets something else and restores it would otherwise turn the default into a chosen "run" by accident. So, at runtime, "load" or "off" keep UDT work from switching to "on", as does GRAPHBLAS_JIT_C_CONTROL=run set before starting. gb.ss.config["jit_c_control"] = None is the default whatever is configured. Under a "run" that was set, UDT ops run their Numba functions (no compiled kernel, typically 2-3x slower for elementwise ops), and types and ops registered from C source (dtypes.ss.register_new and the op modules' ss.register_new) fail with a message that says the control was set to "run" and how to opt in: gb.ss.config["jit_c_control"] = "on", or None for the default. "on" makes SuiteSparse compile kernels it has no precompiled one for: Matrix.reduce_rowwise with lor on BOOL or plus on UINT64 measured 2.7-8x faster than the generic kernel. It stays opt-in because from SuiteSparse 9.4 on a failed compile returns GxB_JIT_ERROR, and processes that share one cache directory race (GraphBLAS #449; 10.1.1 dropped the file locks): on 10.5.2, 127 of 280 concurrent processes compiling the same new kernels into one cache got errors, invalidated or cleared outputs, wrong values or hangs, where 8.2.1, which locked the cache, had none. When "on" is asked for, it is set only where it works (SuiteSparse 7 has no JIT, and nothing is set there); elsewhere init sets "run", with no error or warning: a library built without the JIT (psg's wheels clamp writes to "run"), no compiler on disk after the conda repair, a cache directory it cannot write to (a probe cached there still loads, and the first new kernel would then fail to compile), or a compiler that fails a probe. Without the probe a broken compiler would fail the first user operation that needs a new kernel (a slice, say) with a bare JitError. The probe compiles a trivial type through GxB_Type_new; the kernel is cached, so after the first time per cache directory a probe is a load, about 2 ms against 0.3 s. The type is named after a hash of the compiler name, flags and libraries, so another toolchain sharing the cache gets a probe of its own. A failed probe sets "run", not the "load" SuiteSparse leaves, since "load" still reads kernels from the cache. UDT work makes the same writability check before it raises the default "run" to "on": without it (as before this change), a read-only cache that already held the probe let the raise through, and the first new UDT kernel failed with a bare JitError; now the UDT op runs its Numba function and the control stays "run". _probe_jit now goes through the C API for fix_jit_config and UDT work too, so it registers no dtypes.ss._jit_probe and needs no re-entrancy guard. Init can run in the middle of importing a core module, so it reads and writes these settings through the C API (_CConfig) instead of gb.ss.config, and jit_config imports ffi and lib inside the functions that use them: init imports it, so importing it first would otherwise start init halfway through its own import. Importing gb.ss still changes nothing (#630). --- graphblas/__init__.py | 32 +++ graphblas/core/ss/jit_config.py | 290 ++++++++++++++++++++------- graphblas/graphblas.yaml | 5 + graphblas/ss/__init__.py | 10 +- graphblas/ss/_core.py | 19 ++ graphblas/tests/test_core.py | 21 ++ graphblas/tests/test_dtype.py | 13 +- graphblas/tests/test_ssjit.py | 344 +++++++++++++++++++++++++++++--- scripts/jit_diagnostics.py | 2 +- 9 files changed, 620 insertions(+), 116 deletions(-) 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/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 5a044cd39..4805c738b 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -18,16 +18,25 @@ ``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`). """ 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 +62,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 +85,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 +106,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 @@ -200,39 +212,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(): @@ -242,12 +373,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: @@ -256,57 +389,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 @@ -371,6 +508,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..deceb1dc0 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -1,2 +1,7 @@ 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. +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..c99c9eb4b 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -254,6 +254,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_ssjit.py b/graphblas/tests/test_ssjit.py index e2eaf5ad0..1a37627c2 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,112 @@ 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. + 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 +301,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 +313,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 +352,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 +364,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 +412,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 +428,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 +619,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 +638,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 +650,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 +689,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 +745,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 +787,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 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", From ff6956c3b25896afdefd4902f8f926ce6656f6b6 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 21:26:38 -0500 Subject: [PATCH 06/11] Docs: the SuiteSparse JIT control, how to opt in, and when "on" is unsafe A section in the initialization page of the user guide, about a screen long: what "on" buys (the 2.7-8x reduce numbers), that the default is "run" and that UDT work raises it but leaves a control the user set, how to opt in (GRAPHBLAS_JIT_C_CONTROL, gb.config, gb.ss.config) and the values, that unset is the default and a "run" set later in gb.ss.config is then that default again (at runtime "load" or "off", or GRAPHBLAS_JIT_C_CONTROL=run before starting, keep UDT work from switching to "on"), where "on" falls back to "run" with no error, and the costs of "on". The first is plain: processes running at the same time must not share a cache directory on SuiteSparse 10.1.1 and later, which no longer lock the cache (GraphBLAS #449: JitError, invalid or cleared outputs, wrong values returned with success, hangs). The simple remedy for multiprocessing, dask or pytest-xdist is GRAPHBLAS_JIT_C_CONTROL=run in the workers' environment, which holds for UDT work too; it costs UDT operators their compiled kernels (they run their Numba functions, typically 2-3x slower elementwise) and makes C-source registration fail. A cache per process keeps the kernels. The default will change once SuiteSparse publishes kernels atomically. Then the first-use compile latency, the cache keyed only on the version, and what to do on trouble: set GRAPHBLAS_JIT_C_CONTROL=run, delete the cache, report it. The UDT page says the control starts at the config value, "run" by default, that UDT work and C-source registration raise that default but keep a control the user set (a "run" when set before init), and what a "run" that was set costs there, and links to the section; graphblas.yaml, jit_config's module docstring and gb.ss.config's jit_c_control entry point there, and the entry says what None does. --- docs/user_guide/init.rst | 69 +++++++++++++++++++++++++++++++++ docs/user_guide/udt.rst | 27 ++++++++----- graphblas/core/ss/jit_config.py | 2 + graphblas/graphblas.yaml | 1 + graphblas/ss/_core.py | 8 +++- 5 files changed, 96 insertions(+), 11 deletions(-) diff --git a/docs/user_guide/init.rst b/docs/user_guide/init.rst index ffb6a3463..a54221d78 100644 --- a/docs/user_guide/init.rst +++ b/docs/user_guide/init.rst @@ -57,3 +57,72 @@ 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. + +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/core/ss/jit_config.py b/graphblas/core/ss/jit_config.py index 4805c738b..62d78a44a 100644 --- a/graphblas/core/ss/jit_config.py +++ b/graphblas/core/ss/jit_config.py @@ -24,6 +24,8 @@ 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 diff --git a/graphblas/graphblas.yaml b/graphblas/graphblas.yaml index deceb1dc0..93a6748ef 100644 --- a/graphblas/graphblas.yaml +++ b/graphblas/graphblas.yaml @@ -4,4 +4,5 @@ mapnumpy: True # 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/_core.py b/graphblas/ss/_core.py index c99c9eb4b..db6ffa4c2 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -143,7 +143,13 @@ 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. Only available for SuiteSparse:GraphBLAS >=8 From e309709a9265767cfcfb75d41b78492d75ffa522 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 29 Oct 2026 04:02:55 -0500 Subject: [PATCH 07/11] Exclude the mixed op codegen's array leaf check from coverage A record with an array leaf has no C typedef, so the mixed op's C definition returns before it looks at the leaves. --- graphblas/core/operator/udt_codegen.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index e866348aa..6b650e227 100644 --- a/graphblas/core/operator/udt_codegen.py +++ b/graphblas/core/operator/udt_codegen.py @@ -1477,7 +1477,8 @@ def assign(target, x_path, y_path, x_stored, y_stored, types, out): 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): - if leaf.shape: + # 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]) From f13eba2de2cdc099355de51b5f57f99bd087d0b8 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 2 Nov 2026 20:59:15 -0800 Subject: [PATCH 08/11] test_global_config: None sets jit_c_control to the default "run" The test expected every key outside GlobalConfig._defaults to refuse None, but "Set the SuiteSparse JIT control at init, from a new config key" gave jit_c_control a setter in which None is the default "run" whatever graphblas.config has, and not a "run" the user chose (UDT work may still raise it to "on"). Check both, then put the control back as the test found it, and restore the user-set flag, which the round trip earlier in the test can change too. --- graphblas/tests/test_ss_utils.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) 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 From 6b9e3a0f2170b422e74066183c32c07425b4f058 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 04:02:46 -0800 Subject: [PATCH 09/11] Skip the unusable-compiler probe check where SuiteSparse compiles with cmake test_set_jit_c_control_at_init points jit_c_compiler_name at a directory and expects the probe to fail. SuiteSparse built with MSVC always compiles through cmake, and cmake's Visual Studio generator chooses the compiler itself, so on Windows the probe succeeds and the control stays 'on'. Every Windows job failed on it. --- graphblas/tests/test_ssjit.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 1a37627c2..56baec2b9 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -267,11 +267,14 @@ def test_set_jit_c_control_at_init(monkeypatch): 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. - 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 + # 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 = [] From 2d06bff84bed535df79a20d6d29f386bcc09d5fc Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 07:00:22 -0800 Subject: [PATCH 10/11] Say that the JIT stays at "run" on Windows, where SuiteSparse compiles with cmake SuiteSparse built with MSVC, as conda-forge builds it, compiles JIT kernels only through cmake. python-graphblas checks the configured C compiler path before it sets "on", and on conda-forge Windows that is the build machine's cl.exe, which no user has, so a configured "on" becomes "run" there. The user guide's JIT section and the jit_use_cmake entry of gb.ss.config say so. --- docs/user_guide/init.rst | 4 +++- graphblas/ss/_core.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/user_guide/init.rst b/docs/user_guide/init.rst index a54221d78..141e1f250 100644 --- a/docs/user_guide/init.rst +++ b/docs/user_guide/init.rst @@ -96,7 +96,9 @@ later to ``gb.ss.config`` is that same default, so to keep UDT work from compili 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. +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: diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index db6ffa4c2..40c7695f0 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -152,6 +152,8 @@ class GlobalConfig(BaseConfig): 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. From 838163a5ce1cdcdecc2be12485abe3248142f50e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 3 Nov 2026 11:28:11 -0800 Subject: [PATCH 11/11] Keep a UDT op's JIT C name within SuiteSparse's 127 characters SuiteSparse refuses a JIT C name of GxB_MAX_NAME_LEN (128) characters or more, but the definition was set anyway and taken under the empty name. A mixed op on two UDTs whose names add up to about 118 characters then compiled a kernel that calls no function: every output value was left uninitialized, with no error. A single-type op on a UDT named with 122 or more characters (on main too) failed with an empty JitError. Shorten a long mixed-op name to a digest (a length of 0 cannot come from the readable form, so the name stays the pair's own), and set the definition only when SuiteSparse took the name; otherwise the op keeps its cfunc. Found by review (r652f). --- graphblas/core/operator/udt_codegen.py | 23 ++++++++-- graphblas/tests/test_ssjit.py | 63 ++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 4 deletions(-) diff --git a/graphblas/core/operator/udt_codegen.py b/graphblas/core/operator/udt_codegen.py index 6b650e227..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 @@ -1507,6 +1508,10 @@ def assign(target, x_path, y_path, x_stored, y_stored, types, out): # 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} }}" @@ -1591,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): @@ -1608,7 +1620,8 @@ 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 @@ -1622,7 +1635,8 @@ def set_jit_c_mixed_on_op(gb_obj, op_name, py_op, dtype, dtype2, ret_type, input result = _make_jit_c_mixed_definition(op_name, py_op, dtype, dtype2, ret_type, input_types) if result is None: return None - _set_jit_c_strings(gb_obj, *result, lib.GrB_BinaryOp_set_String) + if not _set_jit_c_strings(gb_obj, *result, lib.GrB_BinaryOp_set_String): + return None return result @@ -1638,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/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index 56baec2b9..3721eee1f 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -2573,3 +2573,66 @@ def cfg_for(name): 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()