diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index a893ba3ad..287e525d5 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -189,7 +189,7 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary ops; the rest pass both. Returns the cached ``TypedUser*Op``. """ - wrapper = numba.cfunc(wrapper_sig, nopython=True)(wrapper) + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) c_typename = _GB_OBJ_C_TYPENAME[typed_user_cls.opclass] error_label = c_typename.removeprefix("GrB_").removeprefix("GxB_") gb_obj = ffi.new(f"{c_typename}*") diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 97e0a9e70..5329b4415 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -525,7 +525,14 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - binary_udf = numba.njit(func) + # 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 + # ZeroDivisionError inside a cfunc, where Numba prints the traceback + # and returns, handing GraphBLAS an element it never wrote. + 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 @@ -581,7 +588,9 @@ def binary_wrapper(z, x, y): # pragma: no cover (numba) 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)(binary_wrapper) + 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( @@ -981,7 +990,7 @@ def _initialize(cls): (binary.any, _second), ]: binop.orig_func = func - binop._numba_func = numba.njit(func) if _has_numba else None + binop._numba_func = numba.njit(func, error_model="numpy") if _has_numba else None binop._udt_types = {} binop._udt_ops = {} binary.any._numba_func = binary.second._numba_func diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index a19f5ffec..de2ca5f71 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -246,7 +246,9 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - indexbinary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexbinary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexbinary_udf ) @@ -333,7 +335,9 @@ def indexbinary_wrapper( ): # pragma: no cover (numba) z[0] = indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0]) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) new_idxbinop = ffi_new("GxB_IndexBinaryOp*") check_status_carg( lib.GxB_IndexBinaryOp_new( @@ -392,7 +396,9 @@ def _compile_udt(self, dtype, dtype2): numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) new_idxbinop = ffi_new("GxB_IndexBinaryOp*") check_status_carg( lib.GxB_IndexBinaryOp_new( diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index c7af8ed78..5820e66c5 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -116,7 +116,9 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): if name is None: name = getattr(func, "__name__", "") success = False - indexunary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexunary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexunary_udf ) @@ -176,7 +178,9 @@ def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) z[0] = indexunary_udf(x[0], row, col, y[0]) - indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexunary_wrapper) + indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexunary_wrapper + ) new_indexunary = ffi_new("GrB_IndexUnaryOp*") check_status_carg( lib.GrB_IndexUnaryOp_new( diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index fc33ea8c0..3e2a60d45 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -558,7 +558,7 @@ def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scala func_name="_op", source_label=f"", ) - return numba.njit(op_func) + return numba.njit(op_func, error_model="numpy") def _make_array_wrapper( size, diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 22822921a..64f68eb48 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -176,7 +176,9 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): if name is None: name = getattr(func, "__name__", "") success = False - unary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + unary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=unary_udf) return_types = {} nt = numba.types @@ -231,7 +233,9 @@ def unary_wrapper(z, x): def unary_wrapper(z, x): z[0] = unary_udf(x[0]) # pragma: no cover (numba) - unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) + unary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + unary_wrapper + ) new_unary = ffi_new("GrB_UnaryOp*") check_status_carg( lib.GrB_UnaryOp_new( @@ -453,7 +457,7 @@ def _initialize(cls): ]: unop.orig_func = func if _has_numba: - unop._numba_func = numba.njit(func) + unop._numba_func = numba.njit(func, error_model="numpy") else: unop._numba_func = None unop._udt_types = {} diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index e439fc149..de1c01b88 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1359,6 +1359,52 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) assert result.isequal(w) +@pytest.mark.skipif("not supports_udfs") +def test_udf_division_by_zero_follows_numpy(): + """Dividing by zero in a UDF returns numpy's answer instead of losing the element. + + Under Numba's default error model the division raises ZeroDivisionError + inside the cfunc, where Numba prints the traceback and returns, so + GraphBLAS keeps whatever was in the output element (in practice the + previous element's value). ``error_model="numpy"`` fixes that, but only if + it is set on the ``njit`` Dispatcher: a Dispatcher holds one compilation + per signature, and ``_build`` calls ``.compile(sig)`` before the wrapper + exists, so setting it on the ``cfunc`` alone comes too late to matter. + """ + + def _idiv(x, y): # pragma: no cover (numba) + return x // y + + op = BinaryOp.register_anonymous(_idiv, "_udf_zero_idiv") + v = Vector.from_coo([0, 1], [10, 20], dtype=dtypes.INT64) + w = Vector.from_coo([0, 1], [2, 0], dtype=dtypes.INT64) + assert op(v & w).new().to_coo()[1].tolist() == [5, 0] + + def _tdiv(x, y): # pragma: no cover (numba) + return x / y + + op = BinaryOp.register_anonymous(_tdiv, "_udf_zero_tdiv") + v = Vector.from_coo([0, 1], [1.0, 1.0], dtype=dtypes.FP64) + w = Vector.from_coo([0, 1], [2.0, 0.0], dtype=dtypes.FP64) + assert op(v & w).new().to_coo()[1].tolist() == [0.5, float("inf")] + + # Same guarantee for a UDT UDF, which reaches the cfunc by another route. + udt = dtypes.register_anonymous( + np.dtype([("dz_a", np.int64), ("dz_b", np.int64)], align=True), "_UdfDivZeroRec" + ) + + def _rec_idiv(x, y): # pragma: no cover (numba) + return (x["dz_a"] // y["dz_a"], x["dz_b"]) + + op = BinaryOp.register_anonymous(_rec_idiv, "_udf_zero_rec", is_udt=True) + v = Vector(udt, size=1) + v[0] = (10, 5) + w = Vector(udt, size=1) + w[0] = (0, 1) + got = v.ewise_mult(w, op).new()[0].new().value + assert (got["dz_a"], got["dz_b"]) == (0, 5) + + @pytest.mark.skipif("not supports_udfs") def test_select_op_outlives_source_indexunary(): """A SelectOp keeps alive the IndexUnaryOp whose GraphBLAS handle it borrows.