Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
9dcd74d
Add fast path to Vector.get and Matrix.get for plain integer indices
eriknw Jul 31, 2026
a56a575
Add fast paths for __contains__, scalar __setitem__, and ScalarIndexE…
eriknw Jul 31, 2026
0b42ded
Add plain-int fast lane to parse_index
eriknw Jul 31, 2026
4a957ed
Add parity tests for the scalar-access fast paths
eriknw Jul 31, 2026
985734c
Add single-extract fast path for v[i].value / float(v[i]) scalar reads
eriknw Jul 31, 2026
2e1138a
Gate the scalar fast lanes on real integers, matching parse_index
eriknw Aug 3, 2026
12e52c7
Skip the setitem UDT fallback test without UDF support
eriknw Aug 4, 2026
64fba19
Keep the scalar-element fast path a GrB scalar under bizarro scalars
eriknw Aug 4, 2026
f90544e
Leave ints beyond INT64 to the full setitem path
eriknw Oct 8, 2026
d0099bb
Update `scalar._literal_store_dtype` to `literals._literal_store_dtype`
eriknw Oct 14, 2026
ba1d1c1
Teach codespell the bu in the fast-path parity tests
eriknw Oct 18, 2026
458aa9d
Stop the Recorders two tests leave running, and fail a test that does
eriknw Oct 20, 2026
8c054ab
Report GraphBLAS's error when extracting an element into a C scalar
eriknw Oct 20, 2026
69b5bb7
Report GraphBLAS's error from nvals, nrows, ncols, size and Scalar reads
eriknw Oct 20, 2026
2d13242
Test that the scalar fast paths raise GraphBLAS errors as the full pa…
eriknw Oct 20, 2026
49257e5
Cover the row and column case of Matrix in; exclude failure-only helpers
eriknw Oct 20, 2026
36a8610
Give a scalar index expression's other attributes a GrB_Scalar after …
eriknw Nov 3, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 36 additions & 1 deletion graphblas/core/automethods.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,46 @@

from .. import config

# Scalar read-only attributes whose value is a plain read of the underlying C
# scalar. For an index-extract expression (v[i], A[i, j]) these resolve with a
# single extractElement straight into a cscalar via ScalarIndexExpr._extract_fast,
# skipping the extra GrB_Scalar round-trip that `.new()` + Scalar.value performs.
# Only ScalarIndexExpr defines the hook, and no Vector/Matrix expression routes
# these attrs through _get_value, so their resolution short-circuits on the
# membership test.
_fast_scalar_attrs = frozenset(
{
"value",
"__float__",
"__int__",
"__complex__",
"__index__",
"__bool__",
"__array__",
"is_empty",
"_is_empty",
}
)


def _get_value(self, attr=None, default=None):
if config.get("autocompute"):
if self._value is None:
self._value = self.new()
if (
attr in _fast_scalar_attrs
and (extract_fast := getattr(self, "_extract_fast", None)) is not None
):
self._value = extract_fast()
else:
self._value = self.new()
elif (
attr not in _fast_scalar_attrs
and getattr(self, "_extract_fast", None) is not None
and self._value._is_cscalar != self._is_cscalar
):
# A value read cached a C scalar (_extract_fast); give the other attributes
# the kind of Scalar this expression reports, holding the same value.
self._value = self._value.dup(is_cscalar=self._is_cscalar, name=self._value.name)
if attr is None:
return self._value
return getattr(self._value, attr)
Expand Down
5 changes: 5 additions & 0 deletions graphblas/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,11 @@ def record_raw(text):
rec.record_raw(text)


def _is_recording():
"""Whether a Recorder is active; fast paths that bypass ``call`` must check this."""
return _recorder.get(_prev_recorder) is not None


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)
Expand Down
15 changes: 15 additions & 0 deletions graphblas/core/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,21 @@ def parse_indices(self, indices, shape):
def parse_index(self, index, typ, size):
from .scalar import _as_scalar

if typ is int:
# Fast lane for a plain Python int, the overwhelmingly common case.
# The two np.issubdtype checks below cost ~125ns each (measured), a
# meaningful slice of single-element index parsing. output_type maps
# only the exact `int` type to `int` (bool -> bool, numpy ints -> their
# own type), so bool and numpy scalars never enter here and keep their
# existing handling. A plain int is always signed, so the negative
# branch always applies, matching the signedinteger branch below.
if index >= size:
raise IndexError(f"Index out of range: index={index}, size={size}")
if index < 0:
index = index + size
if index < 0:
raise IndexError(f"Index out of range: index={index - size}, size={size}")
return AxisIndex(None, _as_scalar(index, _INDEX, is_cscalar=True), None, size)
if np.issubdtype(typ, np.integer):
if index >= size:
raise IndexError(f"Index out of range: index={index}, size={size}")
Expand Down
155 changes: 149 additions & 6 deletions graphblas/core/matrix.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,16 @@

from .. import backend, binary, monoid, select, semiring
from ..dtypes import _INDEX, FP64, INT64, lookup_dtype, unify
from ..exceptions import DimensionMismatch, InvalidValue, NoValue, check_status
from ..exceptions import (
DimensionMismatch,
GrB_NO_VALUE,
InvalidValue,
NoValue,
check_status,
check_status_carg,
)
from . import _supports_udfs, automethods, ffi, lib, utils
from .base import BaseExpression, BaseType, call
from .base import BaseExpression, BaseType, _is_recording, call
from .descriptor import lookup as descriptor_lookup
from .dtypes import _raise_dtype_or_arraylike
from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater
Expand Down Expand Up @@ -372,6 +379,53 @@ def __setitem__(self, keys, expr, **opts):
M[0, 0:3] = 17

"""
# Fast path for `A[i, j] = scalar`: a plain (row, col) integer pair and
# an exact-fit Python scalar, with no mask/accum/opts, a non-UDT dtype,
# and no active Recorder. Mirrors Updater -> _assign_element for a single
# element while skipping the resolver, Updater, and Scalar objects. Only
# int/float/bool/complex are taken here so dtype inference and cffi
# coercion match _assign_element exactly; everything else (slices, fancy
# indexing, numpy or Scalar values, `A(mask)[i, j] << x`) falls back to
# the full assign path, leaving mask/accum and coercion unchanged. So
# does an int beyond INT64, which that path types as the matrix's dtype
# (``uint64_mat[i, j] = 2**63``; literals._literal_store_dtype).
if (
not opts
and type(keys) is tuple
and len(keys) == 2
and type(expr) in (int, float, bool, complex)
and (type(expr) is not int or -(2**63) <= expr < 2**63)
and not self.dtype._is_udt
and not _is_recording()
):
row, col = keys
# Only int and np.integer keys take the fast lane (bools excluded),
# matching parse_index; other __index__ objects fall through to the
# expression path and its canonical errors.
if (type(row) is int or isinstance(row, np.integer)) and (
type(col) is int or isinstance(col, np.integer)
):
rowidx = row.__index__()
colidx = col.__index__()
nrows = self._nrows
ncols = self._ncols
if rowidx < 0:
rowidx += nrows
if rowidx < 0 or rowidx >= nrows:
raise IndexError(f"Index out of range: index={row}, size={nrows}")
if colidx < 0:
colidx += ncols
if colidx < 0 or colidx >= ncols:
raise IndexError(f"Index out of range: index={col}, size={ncols}")
vdtype = lookup_dtype(type(expr), expr)
cvalue = ffi_new(f"{vdtype.c_type}*")
cvalue[0] = expr # cffi coercion, identical to the Scalar.value setter
err_code = utils.libget(f"GrB_Matrix_setElement_{vdtype.name}")(
self.gb_obj[0], cvalue[0], rowidx, colidx
)
if err_code:
check_status_carg(err_code, "Matrix", self.gb_obj[0])
return
Updater(self, opts=opts)[keys] = expr

def __contains__(self, index):
Expand All @@ -384,6 +438,46 @@ def __contains__(self, index):
(10, 15) in M

"""
# Fast path for a plain (row, col) integer pair: probe with
# GrB_Matrix_extractElement directly instead of building an extract
# expression and Scalar. An out-of-range index falls through to the
# expression path so it raises the same IndexError as the slow path.
# Fall back for a UDT dtype and an active Recorder (so the call is
# recorded), mirroring Matrix.get. TransposedMatrix reuses this method.
if (
type(index) is tuple
and len(index) == 2
and not self.dtype._is_udt
and not _is_recording()
):
row, col = index
# Only int and np.integer take the fast lane (bools excluded),
# matching parse_index; other __index__ objects fall through to the
# expression path and its canonical errors.
if (type(row) is int or isinstance(row, np.integer)) and (
type(col) is int or isinstance(col, np.integer)
):
rowidx = row.__index__()
colidx = col.__index__()
nrows = self._nrows
ncols = self._ncols
if rowidx < 0:
rowidx += nrows
if colidx < 0:
colidx += ncols
if 0 <= rowidx < nrows and 0 <= colidx < ncols:
if self._is_transposed:
rowidx, colidx = colidx, rowidx
dtype = self.dtype
res = ffi_new(f"{dtype.c_type}*")
err_code = utils.libget(f"GrB_Matrix_extractElement_{dtype.name}")(
res, self.gb_obj[0], rowidx, colidx
)
if err_code:
if err_code == GrB_NO_VALUE:
return False
check_status_carg(err_code, "Matrix", self.gb_obj[0])
return True
extractor = self[index]
if not extractor._is_scalar:
raise TypeError(
Expand Down Expand Up @@ -507,14 +601,18 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False, **opts
def nrows(self):
"""Number of rows in the Matrix."""
scalar = _scalar_index("s_nrows")
call("GrB_Matrix_nrows", [_Pointer(scalar), self])
# TransposedMatrix shares this, and reports against its Matrix (no error function)
matrix = self._matrix if self._is_transposed else self
call("GrB_Matrix_nrows", [_Pointer(scalar, exc_arg=matrix), self])
return scalar.gb_obj[0]

@property
def ncols(self):
"""Number of columns in the Matrix."""
scalar = _scalar_index("s_ncols")
call("GrB_Matrix_ncols", [_Pointer(scalar), self])
# TransposedMatrix shares this, and reports against its Matrix (no error function)
matrix = self._matrix if self._is_transposed else self
call("GrB_Matrix_ncols", [_Pointer(scalar, exc_arg=matrix), self])
return scalar.gb_obj[0]

@property
Expand All @@ -526,7 +624,9 @@ def shape(self):
def nvals(self):
"""Number of non-empty values in the Matrix."""
scalar = _scalar_index("s_nvals")
call("GrB_Matrix_nvals", [_Pointer(scalar), self])
# TransposedMatrix shares this, and reports against its Matrix (no error function)
matrix = self._matrix if self._is_transposed else self
call("GrB_Matrix_nvals", [_Pointer(scalar, exc_arg=matrix), self])
return scalar.gb_obj[0]

@property
Expand Down Expand Up @@ -857,6 +957,46 @@ def get(self, row, col, default=None):
Python scalar

"""
# Fast path for plain integer indices: call GrB_Matrix_extractElement
# directly instead of building an extract expression, which costs ~10x
# more than the C call for single-element access. Fall back when a
# Recorder is active (so the call is recorded) and for UDTs (whose
# values need numpy-based conversion in Scalar.value).
# Only int and np.integer take the fast lane (bools excluded), matching
# parse_index; other __index__ objects fall through to the expression
# path and its canonical errors.
if (
(type(row) is int or isinstance(row, np.integer))
and (type(col) is int or isinstance(col, np.integer))
and not self.dtype._is_udt
and not _is_recording()
):
rowidx = row.__index__()
colidx = col.__index__()
nrows = self._nrows
ncols = self._ncols
if rowidx < 0:
rowidx += nrows
if rowidx < 0 or rowidx >= nrows:
raise IndexError(f"Index out of range: index={row}, size={nrows}")
if colidx < 0:
colidx += ncols
if colidx < 0 or colidx >= ncols:
raise IndexError(f"Index out of range: index={col}, size={ncols}")
if self._is_transposed:
# TransposedMatrix reuses this method; gb_obj is the
# untransposed parent, so extract the mirrored element.
rowidx, colidx = colidx, rowidx
dtype = self.dtype
res = ffi_new(f"{dtype.c_type}*")
err_code = utils.libget(f"GrB_Matrix_extractElement_{dtype.name}")(
res, self.gb_obj[0], rowidx, colidx
)
if err_code:
if err_code == GrB_NO_VALUE:
return default
check_status_carg(err_code, "Matrix", self.gb_obj[0])
return res[0]
expr = self[row, col]
if expr._is_scalar:
rv = expr.new().value
Expand Down Expand Up @@ -3098,10 +3238,13 @@ def _extract_element(
desc = descriptor_lookup(**opts) # noqa: F841 (keep desc in scope for context)
if is_cscalar:
dtype_name = "UDT" if dtype._is_udt else dtype.name
# Report an error against this Matrix (not a TransposedMatrix, which has no
# error function): the empty C scalar has no GraphBLAS object to ask.
out = _Pointer(result, exc_arg=self._matrix if self._is_transposed else self)
if (
call(
f"GrB_Matrix_extractElement_{dtype_name}",
[_Pointer(result), self, rowidx.index, colidx.index],
[out, self, rowidx.index, colidx.index],
)
is not NoValue
):
Expand Down
54 changes: 50 additions & 4 deletions graphblas/core/scalar.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from ..dtypes import _INDEX, FP64, _index_dtypes, lookup_dtype, unify
from ..exceptions import EmptyObject, check_status
from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils
from .base import BaseExpression, BaseType, call
from .base import BaseExpression, BaseType, _is_recording, call
from .expr import AmbiguousAssignOrExtract
from .literals import (
_SINGLE_PRECISION,
Expand Down Expand Up @@ -396,7 +396,8 @@ def value(self):
else:
scalar = Scalar(self.dtype, is_cscalar=True, name="s_temp")
dtype_name = "UDT" if is_udt else self.dtype.name
call(f"GrB_Scalar_extractElement_{dtype_name}", [_Pointer(scalar), self])
out = _Pointer(scalar, exc_arg=self)
call(f"GrB_Scalar_extractElement_{dtype_name}", [out, self])
if is_udt:
return _element_value(self.dtype, scalar.gb_obj)
return scalar.gb_obj[0]
Expand Down Expand Up @@ -431,7 +432,7 @@ def nvals(self):
if self._is_cscalar:
return 0 if self._empty else 1
scalar = _scalar_index("s_nvals")
call("GrB_Scalar_nvals", [_Pointer(scalar), self])
call("GrB_Scalar_nvals", [_Pointer(scalar, exc_arg=self), self])
return scalar.gb_obj[0]

@property
Expand Down Expand Up @@ -1060,7 +1061,37 @@ class ScalarIndexExpr(AmbiguousAssignOrExtract):
def new(self, dtype=None, *, is_cscalar=None, name=None, **opts):
if is_cscalar is None:
is_cscalar = False
return self.parent._extract_element(
parent = self.parent
# Fast path for the default `expr.new()`: extract a single element
# straight into a fresh GrB_Scalar via GrB_*_extractElement_Scalar,
# skipping the `call` wrapper's per-arg _carg marshalling. Falls back
# for a dtype cast, cscalar output, opts, UDTs, and an active Recorder
# (so the call is recorded). Result is a GrB_Scalar (is_cscalar=False),
# empty exactly when the element is missing, same as _extract_element.
if (
dtype is None
and not is_cscalar
and not opts
and not parent.dtype._is_udt
and not _is_recording()
):
indices = self.resolved_indexes.indices
result = Scalar(parent.dtype, is_cscalar=False, name=name) # pragma: is_grbscalar
if len(indices) == 1:
err_code = lib.GrB_Vector_extractElement_Scalar(
result.gb_obj[0], parent.gb_obj[0], indices[0].index._carg
)
else:
rowidx, colidx = indices
if parent._is_transposed:
rowidx, colidx = colidx, rowidx
err_code = lib.GrB_Matrix_extractElement_Scalar(
result.gb_obj[0], parent.gb_obj[0], rowidx.index._carg, colidx.index._carg
)
if err_code:
check_status(err_code, [result])
return result
return parent._extract_element(
self.resolved_indexes, dtype, opts, is_cscalar=is_cscalar, name=name
)

Expand All @@ -1072,6 +1103,21 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts):
return Scalar(dtype, is_cscalar=is_cscalar, name=name)
return self.new(dtype, is_cscalar=is_cscalar, name=name, **opts)

def _extract_fast(self):
"""Resolve a value read (``.value``, ``float(...)``, ...) with one extract.

Those readers only need the raw element, so extract it straight into a
cscalar and skip the extra GrB_Scalar round-trip that ``.new()`` followed
by ``Scalar.value`` would perform. Defer to the full ``.new()`` for UDTs
(whose values need numpy conversion in ``Scalar.value``) and while a
Recorder is active (so it observes the same calls as the expression path).
``automethods._get_value`` consults this hook for ``_fast_scalar_attrs``.
"""
parent = self.parent
if parent.dtype._is_udt or _is_recording():
return self.new()
return parent._extract_element(self.resolved_indexes, None, {}, is_cscalar=True)

is_cscalar = Scalar.is_cscalar
is_grbscalar = Scalar.is_grbscalar
__hash__ = None
Expand Down
Loading
Loading