From 1ae85cc93bbaf9ad6b757f2bd9c3334749532951 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 3 Aug 2026 15:01:57 -0500 Subject: [PATCH 1/5] Learn built-in UDF return types by typing alone, defer the compiles First touch of a built-in UDF binop compiled it for all 11 to 13 sample dtypes, and threw nearly all of that away. The compiles existed only to read back a return type per dtype for op.types; a program that then used one dtype paid for the rest. Get the return types from numba's typing pipeline alone (run_frontend plus type_inference_stage, no lowering), apply the existing downcast heuristic over the full result set, and defer each dtype's lowering, cfunc wrapper, and GrB_BinaryOp_new to the first op[dtype] access via _build_deferred/_materialize_deferred on OpBase. op.types is unchanged: dumped for all five ops before and after and diffed, identical. Numba exposes no public way to ask for a return type without also compiling for it, so _infer_ret_types_typing_only reaches into numba.core. We support numba back to 0.57 and this was verified against 0.66.0 only, so the helper returns None on a failed import or any surprise from the call and _build falls back to the eager loop. Verified by forcing _HAS_TYPING_ONLY = False: the suite passes with only test_builtin_udf_types_precede_compilation failing, which is the intended signal. A numba that moves these internals turns into a red test rather than a silent return to the old cost. Only the built-in UDFs take this path, keyed off the same flag that marks them cacheable. Typing is more permissive than lowering, so a user function whose dtype types but fails to lower would be wrongly advertised in .types; register_new and register_anonymous stay eager and keep validating lowerability per dtype. Semiring construction iterates _typed_ops directly, so it materializes a deferred multiplier first. Measured here, fresh process, warm numba cache, medians of 3, load average 4 to 5 on a shared machine so these are indicative and the before-numbers are the noisier half. First-touch floordiv: 474 -> 248 ms. First-touch of all five: 753 -> 289 ms. The deferred per-dtype build does not vanish, it moves to first use of that dtype. This does not speed up a program that ends up using every dtype; it stops programs that use one or two from paying for all of them. It does not touch UDT auto-lift, which was already compiled on demand. _finalize_typed_binaryop keeps error_model="numpy" on the cfunc, as the call site it was extracted from had it and every other cfunc in the package sets it. No test here distinguishes it: with it removed from the cfunc, from the njit, or from both, and with the numba cache cleared, integer floordiv by zero still returns numpy's answer. Deferral surfaced one more direct consumer of _typed_ops: the typed commutes_to properties asked whether the PARTNER op had already compiled this dtype, so in a fresh process floordiv[INT64].commutes_to answered None until rfloordiv happened to be built, and the answer depended on access order. The membership test now consults .types, whose keyset is identical to _typed_ops once materialized, letting __getitem__ drive the deferred build. test_deferred_commutes_to pins it in a subprocess, where access order is controlled. --- graphblas/core/operator/base.py | 20 +++ graphblas/core/operator/binary.py | 264 ++++++++++++++++++++-------- graphblas/core/operator/semiring.py | 3 + graphblas/tests/test_op.py | 73 ++++++++ 4 files changed, 282 insertions(+), 78 deletions(-) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index b9cd34e52..f895c7a87 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -375,6 +375,23 @@ def __init__(self, name, *, anonymous=False): def __repr__(self): return f"{self._modname}.{self.name}" + def _build_deferred(self, type_): + """Build a typed op on demand for ``type_``, or return None if this op + has no deferred builds. Overridden where ``.types`` is populated before + the per-dtype numba compilation (see ``BinaryOp``). + """ + return + + def _materialize_deferred(self): + """Build every typed op that ``.types`` advertises but that has not been + compiled yet. A no-op unless ``_build_deferred`` defers builds; call it + before code that iterates ``_typed_ops`` directly (e.g. building a + Semiring over this op). + """ + for type_ in list(self.types): + if type_ not in self._typed_ops: + self._build_deferred(type_) + def __getitem__(self, type_): if type(type_) is tuple: from .utils import get_typed_op @@ -386,6 +403,9 @@ def __getitem__(self, type_): if not self._is_udt: type_ = lookup_dtype(type_) if type_ not in self._typed_ops: + op = self._build_deferred(type_) + if op is not None: + return op if self._udt_types is None: if self.is_positional: return self._typed_ops[UINT64] diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 6abccb39e..a976b446a 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -49,6 +49,26 @@ _get_udt_wrapper, _resolve_udt_return_type, ) + + try: + # Typing-only return-type inference for the module-level built-in UDFs: + # learn return types without lowering or codegen, so ``.types`` can be + # populated up front while the njit compile is deferred to first use. + # Numba's public API offers no way to ask for a return type without + # also compiling for it, so this reaches into ``numba.core``. That is + # internal, and we support numba back to 0.57, so treat every part of + # it as breakable: a failed import or any surprise from the call leaves + # ``_infer_ret_types_typing_only`` returning None and ``_build`` on the + # eager full-compile loop, which is what it did before this existed. + from numba.core import compiler as _numba_compiler + from numba.core import typed_passes as _numba_typed_passes + from numba.core.registry import cpu_target as _numba_cpu_target + + _HAS_TYPING_ONLY = True + except Exception: # pragma: no cover - numba internals moved + _HAS_TYPING_ONLY = False +else: + _HAS_TYPING_ONLY = False if _supports_complex: from ...dtypes import FC32, FC64 @@ -215,13 +235,13 @@ def monoid(self): @property def commutes_to(self): commutes_to = self.parent.commutes_to - if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt): + if commutes_to is not None and (self.type in commutes_to.types or self.type._is_udt): return commutes_to[self.type] @property def _semiring_commutes_to(self): commutes_to = self.parent._semiring_commutes_to - if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt): + if commutes_to is not None and (self.type in commutes_to.types or self.type._is_udt): return commutes_to[self.type] @property @@ -441,6 +461,125 @@ def _pair_dtype(op, dtype, dtype2): return op[INT64] +def _adjust_ret_type(type_, ret_type, return_types): + """Downcast a UDF's inferred return type toward the input type when that is + the intent (INT->INT, FP->FP, FC->FC, and the UINT64/BOOL special cases). + + ``return_types`` holds the results already decided for earlier sample dtypes; + the UINT64 and BOOL rules consult INT64 and INT8, which are decided first. + Shared by the eager compile loop and the typing-only inference path so both + produce the same ``.types``. + """ + if ret_type != type_ and ( + ("INT" in ret_type.name and "INT" in type_.name) + or ("FP" in ret_type.name and "FP" in type_.name) + or ("FC" in ret_type.name and "FC" in type_.name) + or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64) + ): + # This is what users want most of the time, but we can't make a perfect + # rule. There should be a way for users to be explicit. + return type_ + if type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: + return INT8 + return ret_type + + +def _finalize_typed_binaryop(parent, binary_udf, name, type_, ret_type): + """Compile the cfunc wrapper for one dtype and register the ``GrB_BinaryOp``. + + Shared by the eager build loop and the deferred per-dtype build. ``ret_type`` + is already resolved (heuristic applied); this lowers ``binary_udf`` for the + dtype (via the cfunc), wires up the GraphBLAS op, and records the typed op. + For deferred built-ins this is where that dtype pays its compilation cost. + """ + nt = numba.types + input_type = _bool_to_int8(type_) + return_type = _bool_to_int8(ret_type) + + # Build wrapper because GraphBLAS wants pointers and void return + wrapper_sig = nt.void( + nt.CPointer(return_type.numba_type), + nt.CPointer(input_type.numba_type), + nt.CPointer(input_type.numba_type), + ) + + if type_ == BOOL: + if ret_type == BOOL: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) + + else: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = binary_udf(bool(x[0]), bool(y[0])) + + elif ret_type == BOOL: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = bool(binary_udf(x[0], y[0])) + + else: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = binary_udf(x[0], y[0]) + + binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(binary_wrapper) + new_binary = ffi_new("GrB_BinaryOp*") + check_status_carg( + lib.GrB_BinaryOp_new( + new_binary, + binary_wrapper.cffi, + ret_type.gb_obj, + type_.gb_obj, + type_.gb_obj, + ), + "BinaryOp", + new_binary[0], + ) + op = TypedUserBinaryOp(parent, name, type_, ret_type, new_binary[0]) + parent._add(op) + return op + + +def _infer_ret_types_typing_only(func): + """Infer ``{DataType: return DataType}`` for the sample dtypes without + lowering ``func``, or return None on any failure so the caller falls back to + the eager full-compile loop. + + Used only for the module-level built-in UDFs, whose typing and lowering are + known to agree (verified: typing-only reproduces their live ``.types`` + exactly). Typing is more permissive than lowering in general, so this must + not be used for arbitrary user funcs, where a dtype that types but fails to + lower would be wrongly reported as supported. + """ + if not _HAS_TYPING_ONLY: + return None + try: + typingctx = _numba_cpu_target.typing_context + targetctx = _numba_cpu_target.target_context + typingctx.refresh() + targetctx.refresh() + return_types = {} + for type_ in _sample_values: + try: + interp = _numba_compiler.run_frontend(func) + result = _numba_typed_passes.type_inference_stage( + typingctx, targetctx, interp, [type_.numba_type, type_.numba_type], None + ) + numba_ret_type = result.return_type + except numba.TypingError: + # This dtype does not type-check; skip it, matching the eager + # loop's ``except numba.TypingError: continue``. + continue + ret_type = _adjust_ret_type(type_, lookup_dtype(numba_ret_type), return_types) + return_types[type_] = ret_type + except Exception: # pragma: no cover - unexpected numba result shape/API + return None + else: + return return_types + + if _has_numba: from .udt_codegen import ( _compile_codegen, @@ -468,6 +607,7 @@ class BinaryOp(OpBase): "_is_udt", "_numba_func", "_custom_dtype", + "_defer_builds", ) _module = binary _modname = "binary" @@ -579,11 +719,11 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): if name is None: name = getattr(func, "__name__", "") success = False - # The error model has to be set here, not only on the cfunc wrapper - # below: ``.compile(sig)`` further down builds the specialization the - # wrapper then reuses, and a Dispatcher keeps one compilation per - # signature. Whichever compile happens first fixes the model, so - # setting it only on the wrapper leaves ``x // 0`` raising + # The error model has to be set here as well as on the cfunc wrapper in + # ``_finalize_typed_binaryop``. A Dispatcher keeps one compilation per + # signature, so whichever compile happens first fixes the model for + # that signature, and the eager loop's ``.compile(sig)`` gets there + # before the wrapper does. Miss either one and ``x // 0`` raises # ZeroDivisionError inside a cfunc, where Numba prints the traceback # and returns, handing GraphBLAS an element it never wrote. # @@ -603,78 +743,33 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): binary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) return_types = {} - nt = numba.types if not is_udt: - for type_ in _sample_values: - sig = (type_.numba_type, type_.numba_type) - try: - binary_udf.compile(sig) - except numba.TypingError: - continue - ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) - if ret_type != type_ and ( - ("INT" in ret_type.name and "INT" in type_.name) - or ("FP" in ret_type.name and "FP" in type_.name) - or ("FC" in ret_type.name and "FC" in type_.name) - or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64) - ): - # Downcast `ret_type` to `type_`. - # This is what users want most of the time, but we can't make a perfect rule. - # There should be a way for users to be explicit. - ret_type = type_ - elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: - ret_type = INT8 - - input_type = _bool_to_int8(type_) - return_type = _bool_to_int8(ret_type) - - # Build wrapper because GraphBLAS wants pointers and void return - wrapper_sig = nt.void( - nt.CPointer(return_type.numba_type), - nt.CPointer(input_type.numba_type), - nt.CPointer(input_type.numba_type), - ) - - if type_ == BOOL: - if ret_type == BOOL: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) - - else: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(bool(x[0]), bool(y[0])) - - elif ret_type == BOOL: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(x[0], y[0])) - - else: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(x[0], y[0]) - - binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( - binary_wrapper - ) - new_binary = ffi_new("GrB_BinaryOp*") - check_status_carg( - lib.GrB_BinaryOp_new( - new_binary, - binary_wrapper.cffi, - ret_type.gb_obj, - type_.gb_obj, - type_.gb_obj, - ), - "BinaryOp", - new_binary[0], - ) - op = TypedUserBinaryOp(new_type_obj, name, type_, ret_type, new_binary[0]) - new_type_obj._add(op) - success = True - return_types[type_] = ret_type + # ``cache=True`` marks the module-level built-in UDFs (floordiv and + # friends). For those, infer return types without lowering so + # ``.types`` is fully populated up front, then defer each dtype's + # njit lowering + cfunc wrapper to first use (see ``_build_deferred``). + # register_new/register_anonymous user funcs keep the eager loop + # below, which validates lowerability for every sample dtype. If + # inference fails for any reason it returns None and we fall back to + # the eager loop too. + inferred = _infer_ret_types_typing_only(func) if cache else None + if inferred is not None: + new_type_obj.types.update(inferred) + return_types.update(inferred) + new_type_obj._defer_builds = True + success = bool(inferred) + else: + for type_ in _sample_values: + sig = (type_.numba_type, type_.numba_type) + try: + binary_udf.compile(sig) + except numba.TypingError: + continue + ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) + ret_type = _adjust_ret_type(type_, ret_type, return_types) + _finalize_typed_binaryop(new_type_obj, binary_udf, name, type_, ret_type) + success = True + return_types[type_] = ret_type if success or is_udt: return new_type_obj raise UdfParseError("Unable to parse function using Numba") @@ -1149,6 +1244,10 @@ def __init__( self._is_udt = is_udt self.is_positional = is_positional self._custom_dtype = None + # Set True in ``_build`` for the built-in UDFs whose ``.types`` was + # populated by typing-only inference; ``_build_deferred`` then compiles + # each dtype's typed op on first request. + self._defer_builds = False if is_udt: self._udt_types = {} # {(dtype, dtype): DataType} self._udt_ops = {} # {(dtype, dtype): TypedUserBinaryOp} @@ -1161,6 +1260,15 @@ def __init__( self._udt_ops = {} self._custom_dtype = _udt_dtype + def _build_deferred(self, type_): + # For built-ins whose ``.types`` was populated by typing-only inference, + # compile and register the typed op for ``type_`` the first time it is + # requested. ``ret_type`` is read back from ``.types`` (never re-derived + # per dtype: the UINT64/BOOL heuristic needs the full sample pass). + if not self._defer_builds or type_ not in self.types: + return None + return _finalize_typed_binaryop(self, self._numba_func, self.name, type_, self.types[type_]) + __call__ = TypedBuiltinBinaryOp.__call__ is_commutative = TypedBuiltinBinaryOp.is_commutative commutes_to = ParameterizedBinaryOp.commutes_to diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index f9d8603db..0f3ab7310 100644 --- a/graphblas/core/operator/semiring.py +++ b/graphblas/core/operator/semiring.py @@ -245,6 +245,9 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): new_type_obj = cls(name, monoid, binaryop, anonymous=anonymous) if binaryop._is_udt: return new_type_obj + # A built-in UDF multiplier (e.g. floordiv) may defer its per-dtype + # builds; force them so the iteration below sees every typed op. + binaryop._materialize_deferred() for binary_in, binary_func in binaryop._typed_ops.items(): binary_out = binary_func.return_type # Unfortunately, we can't have user-defined monoids over bools yet diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 64bc1880c..60ede25a1 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7051,3 +7051,76 @@ def test_builtin_udf_compiles_uncached_where_numba_cannot_cache(monkeypatch): assert isinstance(op._numba_func._cache, caching.NullCache) v = Vector.from_coo([0, 1], [-3, 4]) assert v.ewise_mult(v, op).new().isequal(Vector.from_coo([0, 1], [1, 1])) + + +# The built-in UDF binops advertise every dtype in ``.types`` up front but +# compile none of them until asked. Run in a subprocess: any earlier test may +# already have materialized them in this process. +_DEFERRED_BUILD_PROBE = """ +import graphblas as gb + +op = gb.binary.floordiv +print("package: " + gb.__file__) +print("types: %d" % len(op.types)) +print("compiled before: %d" % len(op._typed_ops)) +op[gb.dtypes.INT64] +print("compiled after: %d" % len(op._typed_ops)) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udf_types_precede_compilation(): + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + result = subprocess.run( + [sys.executable, "-c", _DEFERRED_BUILD_PROBE], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + + assert int(lines["types"]) == len(binary.floordiv.types), report + assert int(lines["compiled before"]) == 0, report + assert int(lines["compiled after"]) == 1, report + + +_DEFERRED_COMMUTES_PROBE = """ +import graphblas as gb + +print("package: " + gb.__file__) +ct = gb.binary.floordiv[gb.dtypes.INT64].commutes_to +print("floordiv_ok: " + str(ct is gb.binary.rfloordiv[gb.dtypes.INT64])) +ct = gb.binary.absfirst[gb.dtypes.INT64].commutes_to +print("absfirst_ok: " + str(ct is gb.binary.abssecond[gb.dtypes.INT64])) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_deferred_commutes_to(): + # A deferred partner op must still answer commutes_to. The membership + # test consults .types, not ._typed_ops: with the latter, a fresh process + # answered None for floordiv[INT64].commutes_to until rfloordiv happened + # to be compiled, so the answer depended on access order. Subprocess for + # the same reason as test_initialize_does_not_build_lazy_udfs: in this + # process the partners may already be built. + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + result = subprocess.run( + [sys.executable, "-c", _DEFERRED_COMMUTES_PROBE], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + assert lines.get("floordiv_ok") == "True", report + assert lines.get("absfirst_ok") == "True", report From 2aa679e89a83394da4739b06fcfdc84c9ef1697a Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 07:43:02 -0500 Subject: [PATCH 2/5] Test deferred built-in UDF builds in the test process conftest touches every delayed op and semiring before the tests run, and building a semiring compiles every dtype of its binary op, so in the test process OpBase.__getitem__ never built a deferred dtype; only the subprocess tests did, and they count compiled dtypes without using the op they get back. Build a new op from floordiv's function as _initialize does, index it, and check that it compiles that dtype alone, returns it on the next index too, and floors toward minus infinity. Without the return of the deferred op, indexing gives a UDT-lifted floordiv instead. A second test sets _HAS_TYPING_ONLY to False, as on a Numba whose internals moved, and checks that the eager build advertises the same types as the typing-only one. --- graphblas/tests/test_op.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 60ede25a1..5ececdc3e 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -7090,6 +7090,34 @@ def test_builtin_udf_types_precede_compilation(): assert int(lines["compiled after"]) == 1, report +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udf_builds_one_dtype_when_indexed(): + # conftest builds every dtype of binary.floordiv before the tests run, so + # build a new op from the same function, as _initialize does + floordiv = BinaryOp._build("floordiv", operator.binary._floordiv, anonymous=True, cache=True) + assert floordiv.types[INT64] == INT64 + assert floordiv.types[BOOL] == INT8 + assert not floordiv._typed_ops + typed = floordiv[INT64] + assert floordiv._typed_ops == {INT64: typed} + assert floordiv[INT64] is typed + v = Vector.from_coo([0, 1, 2], [7, -7, 7]) + w = Vector.from_coo([0, 1, 2], [2, 2, -2]) + assert v.ewise_mult(w, typed).new().isequal(Vector.from_coo([0, 1, 2], [3, -4, -4])) + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udf_types_without_typing_only(monkeypatch): + # A Numba whose typing internals moved compiles every dtype up front instead, + # and must advertise the same types + func = operator.binary._floordiv + deferred = BinaryOp._build("floordiv", func, anonymous=True, cache=True) + monkeypatch.setattr(operator.binary, "_HAS_TYPING_ONLY", False) + eager = BinaryOp._build("floordiv", func, anonymous=True, cache=True) + assert eager.types == deferred.types + assert eager._typed_ops.keys() == eager.types.keys() + + _DEFERRED_COMMUTES_PROBE = """ import graphblas as gb From 6c8e1007dfc281af5c47231eb1711a07394adcd5 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 27 Oct 2026 09:54:49 -0500 Subject: [PATCH 3/5] Skip a dtype a UDF types for but Numba cannot lower Behavior fix. UnaryOp, BinaryOp, IndexUnaryOp (and so SelectOp) and IndexBinaryOp compile a UDF for each sample dtype and skip the dtypes it fails to type for, catching numba.TypingError. Numba types some ops it has no lowering for, such as `>` on complex values, which `max` uses, and raises a bare NotImplementedError from lowering. That escaped the loops, so the whole registration failed: BinaryOp.register_anonymous(lambda x, y: max(x, y)) before: NotImplementedError: No definition for lowering (complex64, complex64) -> bool after: a BinaryOp for BOOL, the ints and the floats; no FC32 or FC64 The same for UnaryOp, IndexUnaryOp, SelectOp and IndexBinaryOp. Catch NotImplementedError beside TypingError in all four loops. Numba's own NumbaNotImplementedError is a TypingError already (numba 0.66 and 0.68). --- graphblas/core/operator/binary.py | 9 ++++--- graphblas/core/operator/indexbinary.py | 3 ++- graphblas/core/operator/indexunary.py | 3 ++- graphblas/core/operator/unary.py | 3 ++- graphblas/tests/test_indexbinary.py | 15 +++++++++++ graphblas/tests/test_op.py | 37 ++++++++++++++++++++++++++ 6 files changed, 64 insertions(+), 6 deletions(-) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index a976b446a..b1b943fe4 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -569,8 +569,7 @@ def _infer_ret_types_typing_only(func): ) numba_ret_type = result.return_type except numba.TypingError: - # This dtype does not type-check; skip it, matching the eager - # loop's ``except numba.TypingError: continue``. + # This dtype does not type-check; skip it, as the eager loop does. continue ret_type = _adjust_ret_type(type_, lookup_dtype(numba_ret_type), return_types) return_types[type_] = ret_type @@ -763,7 +762,11 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): sig = (type_.numba_type, type_.numba_type) try: binary_udf.compile(sig) - except numba.TypingError: + except (numba.TypingError, NotImplementedError): + # Numba types some ops it cannot lower, such as ``>`` (so + # ``max``) on complex values, and then raises a bare + # NotImplementedError. Skip that dtype like one that fails + # to type. continue ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) ret_type = _adjust_ret_type(type_, ret_type, return_types) diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 4539387a9..5d0d2e557 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -380,7 +380,8 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): ) try: indexbinary_udf.compile(sig) - except numba.TypingError: + except (numba.TypingError, NotImplementedError): + # A dtype Numba types but cannot lower; see ``BinaryOp._build``. continue ret_type = lookup_dtype(indexbinary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index a8eca2460..30c8d7926 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -129,7 +129,8 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): sig = (type_.numba_type, UINT64.numba_type, UINT64.numba_type, type_.numba_type) try: indexunary_udf.compile(sig) - except numba.TypingError: + except (numba.TypingError, NotImplementedError): + # A dtype Numba types but cannot lower; see ``BinaryOp._build``. continue ret_type = lookup_dtype(indexunary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index e8f0870e0..caa55aa82 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -188,7 +188,8 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): sig = (type_.numba_type,) try: unary_udf.compile(sig) - except numba.TypingError: + except (numba.TypingError, NotImplementedError): + # A dtype Numba types but cannot lower; see ``BinaryOp._build``. continue ret_type = lookup_dtype(unary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 333c35fcc..857160282 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -607,6 +607,21 @@ def test_bad_udf(): indexbinary.register_anonymous(lambda x, ix, jx, y, iy, jy, theta: result) # noqa: F821 +def test_udf_skips_dtypes_numba_cannot_lower(): + """Numba types ``max`` on complex values but cannot lower it; build the other dtypes.""" + + def max_vals(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return max(x, y) + + op = indexbinary.register_anonymous(max_vals) + assert dtypes.FP64 in op.types + if dtypes._supports_complex: + assert dtypes.FC64 not in op.types + A = Matrix.from_coo([0, 1], [0, 1], [-1.5, 2.5]) + B = Matrix.from_coo([0, 1], [0, 1], [1.0, 1.0]) + assert A.ewise_mult(B, op(0.0)).new().isequal(Matrix.from_coo([0, 1], [0, 1], [1.0, 2.5])) + + def test_bad_type(): with pytest.raises(TypeError, match="UDF argument must be a function"): indexbinary.register_anonymous(42) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 5ececdc3e..c256bc822 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -647,6 +647,43 @@ def times_minus_sum(x, y): assert w.isequal(result) +@pytest.mark.skipif("not supports_udfs") +def test_udf_skips_dtypes_numba_cannot_lower(): + """A UDF is built for the dtypes Numba can lower it for, and skips the rest. + + Numba types ``max`` on complex values, then has no lowering for the ``>`` in + it and raises a bare NotImplementedError. The per-dtype build loops caught + only TypingError, so registering such a UDF failed for every dtype. + """ + + def unary_max(x): + return max(x, 0) # pragma: no cover (numba) + + def binary_max(x, y): + return max(x, y) # pragma: no cover (numba) + + def indexunary_max(x, row, col, thunk): + return max(x, thunk) # pragma: no cover (numba) + + def select_max(x, row, col, thunk): + return max(x, thunk) > 0 # pragma: no cover (numba) + + unary_op = UnaryOp.register_anonymous(unary_max) + binary_op = BinaryOp.register_anonymous(binary_max) + indexunary_op = IndexUnaryOp.register_anonymous(indexunary_max) + select_op = SelectOp.register_anonymous(select_max) + for udf in [unary_op, binary_op, indexunary_op, select_op]: + assert FP64 in udf.types + if dtypes._supports_complex: + assert FC64 not in udf.types + v = Vector.from_coo([0, 1], [-1.5, 2.5]) + w = Vector.from_coo([0, 1], [1.0, 1.0]) + assert v.apply(unary_op).new().isequal(Vector.from_coo([0, 1], [0.0, 2.5])) + assert v.ewise_mult(w, binary_op).new().isequal(Vector.from_coo([0, 1], [1.0, 2.5])) + assert v.apply(indexunary_op, 1.0).new().isequal(Vector.from_coo([0, 1], [1.0, 2.5])) + assert v.select(select_op, 0.0).new().isequal(Vector.from_coo([1], [2.5], size=2)) + + @pytest.mark.skipif("not supports_udfs") def test_monoid_udf(): def plus_plus_one(x, y): From 168839e6bfea511f525a0e786c25f5440d823534 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 27 Oct 2026 16:18:58 -0500 Subject: [PATCH 4/5] Say an is_udt UDF does not work with a dtype Numba cannot lower Behavior fix. An is_udt=True UDF compiles for each dtype on first use, through _compile_udf_for_udt, which turns a NumbaError into UdfParseError ("does not work with"). Numba types some ops it has no lowering for, such as `>` on complex values, which `max` uses, and raises a bare NotImplementedError from lowering. That escaped, from `in` too: op = BinaryOp.register_anonymous(lambda x, y: max(x, y), is_udt=True) op["FC64"] # before: NotImplementedError: No definition for lowering # (complex128, complex128) -> bool "FC64" in op # before: the same NotImplementedError op["FC64"] # after: UdfParseError: binary. does not work with # (FC64, FC64): No definition for lowering ... "FC64" in op # after: False The same for UnaryOp, IndexUnaryOp and SelectOp, which share the helper. Catch NotImplementedError there too, as the per-dtype build loops now do. --- graphblas/core/operator/udf.py | 6 ++++-- graphblas/tests/test_op.py | 23 +++++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 30b3e7ff4..4c0b45c43 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -91,11 +91,13 @@ def _compile_udf_for_udt(numba_func, sig, *, op_kind, op_name, dtypes): Catches the full ``NumbaError`` hierarchy (TypingError, LoweringError, UnsupportedError, ...) so any compilation failure produces a - single-line UDT diagnostic. + single-line UDT diagnostic. Also catches the bare NotImplementedError + Numba raises for an op it types but cannot lower, such as ``>`` (so + ``max``) on complex values. """ try: numba_func.compile(sig) - except NumbaError as exc: + except (NumbaError, NotImplementedError) as exc: dtypes_str = ", ".join(str(d) for d in dtypes) snippet = _summarize_numba_typing_error(exc) raise UdfParseError( diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index c256bc822..a7a44f0b9 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -684,6 +684,29 @@ def select_max(x, row, col, thunk): assert v.select(select_op, 0.0).new().isequal(Vector.from_coo([1], [2.5], size=2)) +@pytest.mark.skipif("not supports_udfs") +def test_udt_udf_on_a_dtype_numba_cannot_lower(): + """An ``is_udt=True`` UDF on a dtype Numba types but cannot lower does not work with it. + + It used to raise Numba's bare NotImplementedError, from ``in`` too. + """ + + def unary_max(x): + return max(x, 0) # pragma: no cover (numba) + + def binary_max(x, y): + return max(x, y) # pragma: no cover (numba) + + unary_op = UnaryOp.register_anonymous(unary_max, is_udt=True) + binary_op = BinaryOp.register_anonymous(binary_max, is_udt=True) + for udf in [unary_op, binary_op]: + assert FP64 in udf + if dtypes._supports_complex: + with pytest.raises(UdfParseError, match=r"does not work with \(FC64.*lowering"): + udf[FC64] + assert FC64 not in udf + + @pytest.mark.skipif("not supports_udfs") def test_monoid_udf(): def plus_plus_one(x, y): From 0cc181637a6dcfda99f5ede53ee29e9af6fc3197 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 28 Oct 2026 04:13:17 -0500 Subject: [PATCH 5/5] Skip a dtype Numba cannot lower on Numba before 0.60 too Behavior fix. Registering a UDF skips the dtypes Numba types but has no lowering for, such as complex for `max`, by catching the bare NotImplementedError from lowering. Numba before 0.60 captures errors "old_style" by default and raises a LoweringError from it instead, and 0.58 and 0.59 first warn NumbaPendingDeprecationWarning, which `-W error` (as in our tests) raises in its place. CI draws Numba 0.57 to 0.68: BinaryOp.register_anonymous(lambda x, y: max(x, y)) # Numba 0.57 to 0.59 before: LoweringError: Failed in nopython mode pipeline (step: native lowering); 0.58 and 0.59 also warn, and under -W error raise NumbaPendingDeprecationWarning instead after: a BinaryOp with no FC32 or FC64, as on 0.60 and later The same for UnaryOp, IndexUnaryOp, SelectOp and IndexBinaryOp. Their per-dtype loops now share _compile_for_dtype, which also skips a LoweringError raised from a NotImplementedError, and raises any other. An is_udt UDF on such a dtype already raised UdfParseError, except under -W error on 0.58 and 0.59. _numba_compile keeps that warning out of both paths, and the message now names the missing lowering on every version: op = BinaryOp.register_anonymous(lambda x, y: max(x, y), is_udt=True) op["FC64"] # Numba 0.57 to 0.59 before: UdfParseError: ... does not work with (FC64, FC64): Failed in nopython mode pipeline (step: native lowering) after: UdfParseError: ... does not work with (FC64, FC64): No definition for lowering (complex128, complex128) -> bool --- graphblas/core/operator/binary.py | 9 +---- graphblas/core/operator/indexbinary.py | 6 +-- graphblas/core/operator/indexunary.py | 6 +-- graphblas/core/operator/udf.py | 52 +++++++++++++++++++++++++- graphblas/core/operator/unary.py | 6 +-- graphblas/tests/test_op.py | 34 ++++++++++++++++- 6 files changed, 91 insertions(+), 22 deletions(-) diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index b1b943fe4..517ba72e1 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -44,6 +44,7 @@ from .udf import ( _bool_to_int8, + _compile_for_dtype, _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, @@ -760,13 +761,7 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): else: for type_ in _sample_values: sig = (type_.numba_type, type_.numba_type) - try: - binary_udf.compile(sig) - except (numba.TypingError, NotImplementedError): - # Numba types some ops it cannot lower, such as ``>`` (so - # ``max``) on complex values, and then raises a bare - # NotImplementedError. Skip that dtype like one that fails - # to type. + if not _compile_for_dtype(binary_udf, sig): continue ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) ret_type = _adjust_ret_type(type_, ret_type, return_types) diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 5d0d2e557..0e0eff04d 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -19,6 +19,7 @@ from .udf import ( _bool_to_int8, + _compile_for_dtype, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, _resolve_udt_return_type, @@ -378,10 +379,7 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): UINT64.numba_type, type_.numba_type, ) - try: - indexbinary_udf.compile(sig) - except (numba.TypingError, NotImplementedError): - # A dtype Numba types but cannot lower; see ``BinaryOp._build``. + if not _compile_for_dtype(indexbinary_udf, sig): continue ret_type = lookup_dtype(indexbinary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 30c8d7926..7d187bfe7 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -14,6 +14,7 @@ from .udf import ( _bool_to_int8, + _compile_for_dtype, _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, @@ -127,10 +128,7 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if not is_udt: for type_ in _sample_values: sig = (type_.numba_type, UINT64.numba_type, UINT64.numba_type, type_.numba_type) - try: - indexunary_udf.compile(sig) - except (numba.TypingError, NotImplementedError): - # A dtype Numba types but cannot lower; see ``BinaryOp._build``. + if not _compile_for_dtype(indexunary_udf, sig): continue ret_type = lookup_dtype(indexunary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index 4c0b45c43..ab5f0b8a2 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -11,16 +11,19 @@ Built-in operators that lift to UDTs are generated by ``udt_codegen``, which uses these too. ``_bool_to_int8``, the numba#5395 workaround, is here as well: every -Numba UDF on BOOL uses it, not only those on UDTs. +Numba UDF on BOOL uses it, not only those on UDTs. So is ``_compile_for_dtype``, +which every UDF's per-dtype build uses. Importing this module needs Numba. """ import re +import warnings import numba import numpy as np from numba import NumbaError +from numba.core.errors import LoweringError, NumbaPendingDeprecationWarning from ...dtypes import BOOL, INT8, UINT64, lookup_dtype from ...exceptions import UdfParseError, check_status_carg @@ -45,6 +48,50 @@ def _bool_to_int8(dtype): return INT8 if dtype == BOOL else dtype +def _numba_compile(numba_func, sig): + """``numba_func.compile(sig)``, without Numba's warning about how it captures errors. + + Numba before 0.60 captures errors "old_style" by default: a bare exception + from lowering, such as the NotImplementedError for an op Numba types but + cannot lower, comes out as a LoweringError raised from it. 0.58 and 0.59 + also warn NumbaPendingDeprecationWarning that code may depend on this, and + ``-W error`` (as in our tests) raises that warning instead of the error. + The callers here handle both styles, so the warning only gets in the way. + """ + # MAINT 2026-10-28: once Numba 0.61 is the minimum (0.60 made "new_style" the + # default; 0.61 removed "old_style"), drop this warning filter and the + # LoweringError branch in _compile_for_dtype. + if getattr(numba.core.config, "CAPTURED_ERRORS", None) != "old_style": + numba_func.compile(sig) + return + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + "Code using Numba extension API maybe depending on 'old_style'", + NumbaPendingDeprecationWarning, + ) + numba_func.compile(sig) + + +def _compile_for_dtype(numba_func, sig): + """Compile ``sig`` and return whether Numba could; raise any other error. + + It cannot when the function does not type for ``sig``, or when Numba types + an op it has no lowering for, such as ``>`` (so ``max``) on complex values. + Numba raises a bare NotImplementedError for the latter, or before 0.60 a + LoweringError from one. + """ + try: + _numba_compile(numba_func, sig) + except (numba.TypingError, NotImplementedError): + return False + except LoweringError as exc: + if isinstance(exc.__context__, NotImplementedError): + return False + raise + return True + + def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, typed_user_cls): """Compile the cfunc, allocate the ``GrB`` op, wrap it, and cache it. @@ -96,7 +143,7 @@ def _compile_udf_for_udt(numba_func, sig, *, op_kind, op_name, dtypes): ``max``) on complex values. """ try: - numba_func.compile(sig) + _numba_compile(numba_func, sig) except (NumbaError, NotImplementedError) as exc: dtypes_str = ", ".join(str(d) for d in dtypes) snippet = _summarize_numba_typing_error(exc) @@ -108,6 +155,7 @@ def _compile_udf_for_udt(numba_func, sig, *, op_kind, op_name, dtypes): # Numba prefixes its actionable diagnostic line with one of these. _NUMBA_DIAG_PREFIXES = ( "No implementation of function", + "No definition for lowering", "No conversion from", "Field ", "Cannot infer", diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index caa55aa82..b6e615ca5 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -37,6 +37,7 @@ from .udf import ( _bool_to_int8, + _compile_for_dtype, _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, @@ -186,10 +187,7 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): if not is_udt: for type_ in _sample_values: sig = (type_.numba_type,) - try: - unary_udf.compile(sig) - except (numba.TypingError, NotImplementedError): - # A dtype Numba types but cannot lower; see ``BinaryOp._build``. + if not _compile_for_dtype(unary_udf, sig): continue ret_type = lookup_dtype(unary_udf.overloads[sig].signature.return_type) if ret_type != type_ and ( diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index a7a44f0b9..df5896e88 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -702,11 +702,43 @@ def binary_max(x, y): for udf in [unary_op, binary_op]: assert FP64 in udf if dtypes._supports_complex: - with pytest.raises(UdfParseError, match=r"does not work with \(FC64.*lowering"): + with pytest.raises( + UdfParseError, match=r"does not work with \(FC64.*: No definition for lowering" + ): udf[FC64] assert FC64 not in udf +@pytest.mark.skipif("not supports_udfs") +def test_compile_for_dtype_raises_other_lowering_errors(): + """Only a missing lowering skips a dtype, in either of Numba's error styles. + + Numba before 0.60 raises a LoweringError from whatever lowering raised, such + as the NotImplementedError for a missing lowering; later Numba raises that. + """ + from numba.core.errors import LoweringError + + from graphblas.core.operator.udf import _compile_for_dtype + + class Dispatcher: + def __init__(self, exc, old_style): + self.exc = exc + self.old_style = old_style + + def compile(self, sig): + if self.old_style: + # Numba raises it inside the ``except`` that caught the error + err = LoweringError(str(self.exc)) + err.__context__ = self.exc + raise err + raise self.exc + + for old_style in [False, True]: + assert not _compile_for_dtype(Dispatcher(NotImplementedError("none"), old_style), ()) + with pytest.raises(LoweringError if old_style else KeyError): + _compile_for_dtype(Dispatcher(KeyError("bug"), old_style), ()) + + @pytest.mark.skipif("not supports_udfs") def test_monoid_udf(): def plus_plus_one(x, y):