From 8b25e48daaad9560ecf04db60f4143991d40e04e Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 12 Oct 2026 12:19:41 -0500 Subject: [PATCH 1/3] Add a fast path for assigning a Python scalar under a mask w(mask, replace) << x with a Python bool, int or float x goes through _update, a second Updater, an index resolver for `...`, a Scalar built from x and _prep_for_assign before it reaches GrB_Vector_assign_ (or GrB_Matrix_assign_) over all indices. That is about 7 us of Python for a call SuiteSparse finishes in about 0.5 us, and level-synchronous traversals make it once per level: graphblas-algorithms marks each BFS level visited with seen(q.S) << True. Make that call directly when nothing needs the full path: x of type bool, int (within INT64) or float, a non-UDT target, no accum, no opts, no input_mask, a Mask (structural or value, complemented or not) of an object of the target's own type and shape, or no mask without replace, and no active Recorder. x takes the dtype the full path gives its Scalar, and GraphBLAS casts it as before. Everything else, including masks the full path converts or rejects, ints beyond INT64 and numpy scalars, takes the full path unchanged. Measured with no test suite running (load average about 5.3 from other work), mlx, SuiteSparse 10.5.0, best of 5, two rounds each: seen(q.S) << True, 10000-element Vector 7.7 us -> 2.3 to 2.4 us bfs_edges, 10k-node road ball, 134 levels 5.1 ms -> 4.2 to 4.3 ms bfs_edges, 50k-node road ball, 352 levels 16.5 to 16.9 ms -> 14.7 to 14.9 ms with the same edges. The tests compare the fast path with the full path (forced by an active Recorder) over 7 dtypes, 9 values (bool, 0, -1, 300, 2**40, 1.5, -0.0, nan), 5 mask kinds and replace for Vector, and a smaller grid for Matrix, plus the errors of a wrong-kind or wrong-size mask and an int beyond INT64. Breaking the fast path's replace, complement, structure or value each fails between 216 and 423 cases. --- graphblas/core/base.py | 69 +++++++++++++++++++++++++++++++++- graphblas/tests/test_matrix.py | 22 +++++++++++ graphblas/tests/test_vector.py | 36 ++++++++++++++++++ 3 files changed, 126 insertions(+), 1 deletion(-) diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 3499af531..fd171401e 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -3,7 +3,7 @@ from .. import backend, config from .. import replace as replace_singleton 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 +26,54 @@ 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. + """ + from ..dtypes import lookup_dtype + + 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 +406,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..e8e49c371 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -200,6 +200,42 @@ 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 + + def test_from_coo_default_values(): u = Vector.from_coo([1, 3]) v = Vector.from_coo([1, 3], [1.0, 1.0]) From 7b4644440076b7242bfeb62f1ce9b0de7de19ddc Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 18 Oct 2026 01:00:49 -0500 Subject: [PATCH 2/3] Import lookup_dtype once for the masked literal assign _assign_literal imported lookup_dtype on every call. base.py can import it at module level (literals.py, which base.py imports, already does), and the import was about a sixth of what the fast path still costs: seen(q.S) << True on a 10000-element Vector, best of 5, goes from 2.28 us to 1.9 us (the GrB_Vector_assign_BOOL call itself is 0.42 us). --- graphblas/core/base.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/graphblas/core/base.py b/graphblas/core/base.py index fd171401e..283a6276e 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -2,6 +2,7 @@ from .. import backend, config from .. import replace as replace_singleton +from ..dtypes import lookup_dtype from ..exceptions import check_status from . import NULL, ffi, lib from .descriptor import lookup as descriptor_lookup @@ -47,8 +48,6 @@ def _assign_literal(output, value, mask, replace): ``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. """ - from ..dtypes import lookup_dtype - vdtype = lookup_dtype(type(value), value) cvalue = ffi.new(f"{vdtype.c_type}*") cvalue[0] = value # cffi coercion, identical to the Scalar.value setter From 3a7a9082bac835e343b11f23e00210b6c6d893fe Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 20 Oct 2026 07:49:27 -0500 Subject: [PATCH 3/3] Test that the masked literal assign raises a GraphBLAS error as the full path does _assign_literal checks the status of GrB_*_assign_ itself, and no test reached that check, since Python validates the mask and value first. A Vector whose data was exported with give_ownership=True has no GraphBLAS object left, so the assign fails: both the fast path and the full path (forced by a Recorder) raise NullPointer. --- graphblas/tests/test_vector.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index e8e49c371..1f0ef108d 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -23,6 +23,7 @@ IndexOutOfBound, InvalidObject, InvalidValue, + NullPointer, OutputNotEmpty, UdfParseError, ) @@ -236,6 +237,18 @@ def test_update_with_a_literal_keeps_the_full_path_errors(): 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])