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..517ba72e1 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -44,11 +44,32 @@ from .udf import ( _bool_to_int8, + _compile_for_dtype, _compile_udf_for_udt, _finalize_udt_op, _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 +236,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 +462,124 @@ 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, as the eager loop does. + 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,31 @@ 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) + 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) + _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 +1242,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 +1258,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/indexbinary.py b/graphblas/core/operator/indexbinary.py index 4539387a9..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,9 +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: + 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 a8eca2460..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,9 +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: + 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/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/core/operator/udf.py b/graphblas/core/operator/udf.py index 30b3e7ff4..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. @@ -91,11 +138,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: + _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) raise UdfParseError( @@ -106,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 e8f0870e0..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,9 +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: + 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_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 64bc1880c..df5896e88 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -647,6 +647,98 @@ 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_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.*: 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): @@ -7051,3 +7143,104 @@ 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 + + +@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 + +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