diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 3499af531..283a6276e 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -2,8 +2,9 @@ from .. import backend, config from .. import replace as replace_singleton +from ..dtypes import lookup_dtype from ..exceptions import check_status -from . import NULL +from . import NULL, ffi, lib from .descriptor import lookup as descriptor_lookup from .expr import AmbiguousAssignOrExtract, Updater from .literals import _literal_store_dtype @@ -26,6 +27,52 @@ def _is_recording(): return _recorder.get(_prev_recorder) is not None +_ASSIGN_LITERAL_TYPES = (bool, int, float) + + +def _is_plain_mask(mask, output): + """Whether ``mask`` needs nothing from the full assign path: a Mask of an object + of ``output``'s own type and shape, and not a value mask of a UDT. + """ + return ( + isinstance(mask, Mask) + and type(mask.parent) is type(output) + and mask.parent.shape == output.shape + and (mask.structure or not mask.parent.dtype._is_udt) + ) + + +def _assign_literal(output, value, mask, replace): + """``output(mask, replace) << value`` with GrB_*_assign_ over all indices. + + ``value`` takes the dtype the full path gives its Scalar, and GraphBLAS casts + it to ``output``'s dtype, as the full path's assign does. + """ + vdtype = lookup_dtype(type(value), value) + cvalue = ffi.new(f"{vdtype.c_type}*") + cvalue[0] = value # cffi coercion, identical to the Scalar.value setter + if mask is None: + mask_carg = NULL + desc = None + else: + mask_carg = mask.parent._carg + desc = descriptor_lookup( + output_replace=replace, + mask_complement=mask.complement, + mask_structure=mask.structure, + ) + desc = NULL if desc is None else desc._carg + if output.ndim == 1: + args = (output._carg, mask_carg, NULL, cvalue[0], lib.GrB_ALL, output._size, desc) + err_code = libget(f"GrB_Vector_assign_{vdtype.name}")(*args) + else: + nrows, ncols = output._nrows, output._ncols + args = (output._carg, mask_carg, NULL, cvalue[0], lib.GrB_ALL, nrows, lib.GrB_ALL, ncols) + err_code = libget(f"GrB_Matrix_assign_{vdtype.name}")(*args, desc) + if err_code: + check_status(err_code, output) + + def call(cfunc_name, args): call_args = [getattr(x, "_carg", x) if x is not None else NULL for x in args] cfunc = libget(cfunc_name) @@ -358,6 +405,25 @@ def update(self, expr, **opts): return self._update(expr, opts=opts) def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, *, opts): + # Fast path for `w(mask, replace) << x` with an exact-fit Python bool, int or + # float x: the assign the Updater path ends in, without building a second + # Updater, an index resolver and a Scalar (about 7 us of Python a call). + # Ints beyond INT64 (which that path types as the target's dtype), a UDT + # target, accum, opts, any mask the path would convert or reject, and an + # active Recorder (so the call is recorded) take the full path. + if ( + type(expr) in _ASSIGN_LITERAL_TYPES + and not self._is_scalar + and accum is None + and input_mask is None + and not opts + and (type(expr) is not int or -(2**63) <= expr < 2**63) + and not self.dtype._is_udt + and (not replace if mask is None else _is_plain_mask(mask, self)) + and not _is_recording() + ): + _assign_literal(self, expr, mask, replace) + return if not isinstance(expr, BaseExpression): if isinstance(expr, AmbiguousAssignOrExtract): if expr._is_scalar and self._is_scalar: diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 34b05b2ff..9f9fbab36 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -213,6 +213,28 @@ def test_ss_export_full_of_a_matrix_that_is_not_full(method, fmt): assert A.isequal(Matrix.from_coo([0], [1], [2.0], nrows=2, ncols=2)) +@pytest.mark.parametrize("dtype", ["BOOL", "INT64", "FP64"]) +@pytest.mark.parametrize("value", [True, 7, 2.5]) +@pytest.mark.parametrize("mask_kind", [None, "S", "V", "~S", "~V"]) +@pytest.mark.parametrize("replace", [False, True]) +def test_update_with_a_literal_matches_the_full_path(dtype, value, mask_kind, replace): + # C(mask, replace) << x takes a fast path; an active Recorder sends the same + # update down the full path to compare with. + C = Matrix.from_coo([0, 1], [1, 2], [1, 0], dtype=dtype, nrows=3, ncols=3) + M = Matrix.from_coo([0, 1, 2], [0, 2, 1], [True, False, True], nrows=3, ncols=3) + mask = {None: None, "S": M.S, "V": M.V, "~S": ~M.S, "~V": ~M.V}[mask_kind] + if mask is None and replace: + return + fast = C.dup() + full = C.dup() + fast(mask=mask, replace=replace) << value + with gb.Recorder(): + full(mask=mask, replace=replace) << value + assert fast.isequal(full, check_dtype=True) + with pytest.raises(TypeError, match="Vector mask"): + fast(Vector(bool, 3).S) << value + + def test_from_coo_default_values(): B = Matrix.from_coo([1, 3], [2, 4]) C = Matrix.from_coo([1, 3], [2, 4], [1.0, 1.0]) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 509b74173..1f0ef108d 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -23,6 +23,7 @@ IndexOutOfBound, InvalidObject, InvalidValue, + NullPointer, OutputNotEmpty, UdfParseError, ) @@ -200,6 +201,54 @@ def test_ss_export_full_of_a_vector_that_is_not_full(method): assert v.isequal(Vector.from_coo([1], [2.0], size=3)) +@pytest.mark.parametrize("dtype", ["BOOL", "INT8", "UINT8", "INT64", "UINT64", "FP32", "FP64"]) +@pytest.mark.parametrize("value", [True, False, 0, -1, 300, 2**40, 1.5, -0.0, float("nan")]) +@pytest.mark.parametrize("mask_kind", [None, "S", "V", "~S", "~V"]) +@pytest.mark.parametrize("replace", [False, True]) +def test_update_with_a_literal_matches_the_full_path(dtype, value, mask_kind, replace): + # w(mask, replace) << x for a Python bool, int or float takes a fast path; an + # active Recorder sends the same update down the full path to compare with. + w = Vector.from_coo([0, 2, 4], [1, 0, 1], dtype=dtype, size=6) + m = Vector.from_coo([1, 2, 3, 5], [True, False, True, True], size=6) + mask = {None: None, "S": m.S, "V": m.V, "~S": ~m.S, "~V": ~m.V}[mask_kind] + if mask is None and replace: + return + fast = w.dup() + full = w.dup() + fast(mask=mask, replace=replace) << value + with gb.Recorder(): + full(mask=mask, replace=replace) << value + assert fast.dtype == full.dtype + fast_indices, fast_values = fast.to_coo() + full_indices, full_values = full.to_coo() + np.testing.assert_array_equal(fast_indices, full_indices) + np.testing.assert_array_equal(fast_values, full_values) # nan equals nan here + np.testing.assert_array_equal(np.signbit(fast_values), np.signbit(full_values)) + + +def test_update_with_a_literal_keeps_the_full_path_errors(): + v = Vector(int, 3) + with pytest.raises(TypeError, match="Mask object must be type Vector"): + v(Matrix(bool, 3, 3).S) << 1 + with pytest.raises(DimensionMismatch): + v(Vector(bool, 4).S) << 1 + u = Vector(dtypes.UINT64, 2) + u << 2**64 - 1 # beyond INT64: typed as the target's dtype by the full path + assert u.get(0) == 2**64 - 1 + + +@pytest.mark.skipif("not suitesparse") +def test_update_with_a_literal_raises_graphblas_errors(): + # export(give_ownership=True) leaves no GraphBLAS object behind, so the assign + # itself fails, and the fast path raises what the full path does + v = Vector(int, 3) + v.ss.export(give_ownership=True) + with pytest.raises(NullPointer): + v(Vector(bool, 3).S) << 1 + with pytest.raises(NullPointer), gb.Recorder(): + v(Vector(bool, 3).S) << 1 + + def test_from_coo_default_values(): u = Vector.from_coo([1, 3]) v = Vector.from_coo([1, 3], [1.0, 1.0])