From 4cc8cac66dbc50ebee8f09bd63f0fc34dae5227c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 31 Jul 2026 12:17:50 -0500 Subject: [PATCH 1/7] Pass and receive numpy views in array-UDT UDFs A UDF over an array UDT wrote past the end of its output element. The cfunc wrapper passed the raw element pointer and stored the return value with `z_ptr[0] = ...`, so what landed in the buffer was Numba's NestedArray descriptor (data pointer, shape, strides) rather than the element payload. The descriptor is the wider of the two, so the store ran off the end of the element into memory SuiteSparse owns. Measured against the code this replaces, driving the wrapper's cfunc through ctypes over a guarded buffer: (6,) np itemsize 48 | numba value type size 56 | overrun 8 bytes (2, 3) np itemsize 48 | numba value type size 72 | overrun 24 bytes z[:6] = [100.0, 101.0, 102.0, 103.0, 104.0, 105.0] z[6:] = [5e-324, -1.0, -1.0, ...] # -1.0 is the guard byte pattern indices written past the element: [6] Both operands and the output are now bound with `numba.carray(ptr, shape)` in the UDT's declared shape, so the UDF receives a numpy view it can index (`x[i, j]`, `x.shape`) and the wrapper slice-assigns the result back with `z[:] = ...`. Array-typed record leaves slice-assign for a second reason: Numba's record-field setitem copies the destination's extent regardless of the source's, so a short source was read past its end. The return-type resolution has to move with the wrapper. Once the UDF receives a view, a UDF that builds its result (`x + y`) instead of returning an operand types as a plain Numba `Array`, which `lookup_dtype` does not recognize; `_resolve_udt_return_type` gains an `Array` branch matching it back to an input array UDT by base element type and rank. Splitting the two apart would leave array-UDT UDFs that build a result broken outright. An `Array` type carries `ndim` but not its extents, so a return whose shape does not fit the UDT is not a type error. It is now memory-safe, but still silent: the slice-assign's ValueError is raised inside a cfunc, which Numba prints and swallows, leaving the element as SuiteSparse found it. --- docs/user_guide/udt.rst | 20 +++++ graphblas/core/operator/base.py | 130 +++++++++++++++++++++++---- graphblas/core/operator/udt_utils.py | 8 +- graphblas/tests/test_op.py | 127 ++++++++++++++++++++++++++ 4 files changed, 264 insertions(+), 21 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 18af2cf3e..67ec2cc5d 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -173,6 +173,26 @@ For nested record UDTs the tuple is *flat over the leaves*. Given return ``(id, x, y)``, not ``(id, (x, y))``. Returning an existing record value (e.g., one of the inputs) is also fine and preserves the nested shape. +For array UDTs each operand arrives as a numpy view of that element's values, in +the UDT's declared shape: a ``np.dtype((np.float64, (2, 4)))`` UDT hands the UDF +a 2-by-4 array, indexable as ``x[i, j]``. Array expressions work as written, and +the UDF may return one of its operands or build a new array of the same shape: + +.. code-block:: python + + def midpoint(x, y): + return (x + y) / 2 + + op = binary.register_new("midpoint", midpoint, is_udt=True) + + a = Vector(point3, size=1) + a[0] = [0.0, 2.0, 4.0] + b = Vector(point3, size=1) + b[0] = [10.0, 20.0, 30.0] + + c = a.ewise_mult(b, op[point3]).new() + # c[0] = [5.0, 11.0, 17.0] + If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced from Numba's typing pass instead of a 200-line traceback. diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 287e525d5..6b3976c5d 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -329,17 +329,80 @@ def _leaves(d): f"shape {shape}. Return a numpy array (e.g., ``np.array(...)``) or a " f"scalar; tuple returns are only matched to record UDTs." ) + elif isinstance(numba_ret_type, numba.core.types.Array): + # A UDF over an array UDT may build its result (``x + y``) instead + # of returning an operand. Numba types that as a plain Array, which + # ``lookup_dtype`` doesn't recognize, so match it back to an array + # UDT input by base element type and dimensionality. + candidates = [ + d + for d in dtypes + if d._is_udt + and d.np_type.subdtype is not None + and d.numba_type.dtype == numba_ret_type.dtype + and len(d.numba_type.shape) == numba_ret_type.ndim + ] + # An Array type carries ``ndim`` but not its extents, so operands + # that differ only in length are indistinguishable here. Guessing + # would hand SuiteSparse an element of the wrong size, so say so. + # Compare the UDTs rather than their shapes: a flat ``FP64[2, 3]`` + # and a layered ``FP64[3][2]`` are separate DataTypes with separate + # GrB_Type handles, yet Numba collapses both to the same shape. + unique = [] + for d in candidates: + if not any(d is seen for seen in unique): + unique.append(d) + if len(unique) > 1: + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches more than one " + f"input array UDT ({', '.join(str(d) for d in unique)}). " + f"Return one of the operands, or make the operands the same type." + ) + if unique: + return unique[0] + # An array UDT went in and an array came out, but not one that + # fits: name the mismatch rather than fall through to the generic + # "unsupported type", whose advice the user already followed. + array_inputs = [d for d in dtypes if d._is_udt and d.np_type.subdtype is not None] + if array_inputs: + d = array_inputs[0] + nested = d.numba_type + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches no input array " + f"UDT: {d} elements are {nested.dtype} with shape {nested.shape}. " + f"Return an array of that dtype and rank, or one of the operands." + ) raise UdfParseError( f"UDT UDF returned an unsupported type {numba_ret_type!r}. " f"Return a scalar, a tuple matching a record UDT's fields, or a numpy array " f"matching an array UDT's shape." ) + def _array_udt_view(dtype): + """Return ``(base_element_numba_type, shape)`` for an array UDT. + + Array UDTs are addressed as a ``carray`` over their base elements, in + the UDT's declared shape, rather than as Numba's ``NestedArray``. Numba + models a ``NestedArray`` *value* as an array descriptor (data pointer, + shape, strides, ...), so loading or storing one through a ``CPointer`` + moves the descriptor rather than the element payload, corrupting + whatever follows it (and overrunning the element outright once the + descriptor is the wider of the two). + + Read from ``numba_type`` rather than ``np_type.subdtype`` so this + agrees with the type the UDF was compiled against: numpy keeps nested + subarray dtypes layered, e.g. ``FP64[5]`` inside ``[6]`` stays + ``(dtype((' Date: Thu, 6 Aug 2026 00:46:02 -0700 Subject: [PATCH 2/7] Pin the array-UDT any/first/second wrapper to one element's bytes Binary ops outside _BUILTIN_UDT_BINARY_OPS (any, first, second) compile through the generic _numba_func branch of BinaryOp._compile_udt. The old wrapper there loaded and stored an array-UDT operand as a NestedArray value, which Numba models as its full array descriptor (meminfo, parent, nitems, itemsize, data, shape, strides): 56 bytes on 64-bit for a 1-D element, regardless of the 32-byte payload. SuiteSparse's generic reduce keeps a UDT accumulator in a stack array sized to the element, so each fold overflowed it by 24 bytes and clobbered a spilled GrB_Type pointer: segfault on some builds, SIGBUS or a silently wrong answer on others. The carray-based wrapper rework on this branch writes exactly itemsize bytes; this test keeps it that way. The test builds the wrapper for binary.any the same way _compile_udt does, compiles it with numba.cfunc, and calls it via ctypes on heap buffers with slack, so a regression trips an assert instead of corrupting a stack frame. The z guard sentinel (0xAB) must differ from the sentinel in y's trailing slack (0xCD): the descriptor load/store is a byte-preserving copy of the source element plus its trailing bytes, so with one shared sentinel the overflow would rewrite z's guard with identical values and go undetected. A public-path any-reduce smoke runs after the byte-level checks, so a regression fails the assert before reaching the code that can crash the process. Verified both directions: passes here (host numba 0.65.1 and a conda env with numba 0.66 plus python-suitesparse-graphblas 10.0.1.1), and on an export of main the assert trips with exactly 24 overflow bytes while the payload check still passes, demonstrating the blind spot a same-sentinel check would have. --- graphblas/tests/test_op.py | 71 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 0a83d42d7..a3eba185f 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1971,6 +1971,77 @@ def _second(x, y): # pragma: no cover (numba) assert all(val == guard for val in list(z)[size:]), f"wrote past the UDT element: {list(z)}" +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_any_wrapper_stays_within_element(): + """``binary.any`` on an array UDT must write one element, not Numba's array descriptor. + + ``any``, ``first``, and ``second`` are not in ``_BUILTIN_UDT_BINARY_OPS``, + so they compile through the generic ``_numba_func`` branch of + ``BinaryOp._compile_udt``. The wrapper there used to load and store the + operand as a ``NestedArray`` *value*, which Numba models as its full array + descriptor (meminfo, parent, nitems, itemsize, data, shape, strides): 56 + bytes on 64-bit for a 1-D element, regardless of payload size. SuiteSparse's + generic reduce keeps a UDT accumulator in a stack array sized to the element + (32 bytes here), so each fold overflowed it by 24 bytes; depending on the + build that clobbered a spilled pointer (segfault or SIGBUS) or silently + produced a wrong answer. + + Drive the compiled wrapper directly, over heap buffers with slack, so a + regression trips an assert instead of corrupting a stack frame. The two + sentinels must differ: the descriptor load/store is a byte-preserving copy + of the source element plus its trailing bytes, so if ``y``'s slack held the + same sentinel as ``z``'s guard, the overflow would rewrite ``z``'s guard + bytes with identical values and the check would be blind to it. + """ + import ctypes + + import numba + + from graphblas.core.operator.base import _get_udt_wrapper + + # ``register_anonymous`` caches per np.dtype, so this may return the same + # DataType as other float64[4] tests, renamed. That is fine here: the + # wrapper below is compiled fresh and nothing asserts on cached JIT state. + udt = dtypes.register_anonymous(np.dtype((np.float64, (4,))), "_AnyOverflowArr") + + # Mirror the generic ``_numba_func`` branch of ``BinaryOp._compile_udt``. + numba_func = binary.any._numba_func + sig = (udt.numba_type, udt.numba_type) + numba_func.compile(sig) + numba_ret_type = numba_func.overloads[sig].signature.return_type + wrapper, wrapper_sig = _get_udt_wrapper( + numba_func, udt, udt, udt, numba_ret_type=numba_ret_type + ) + cfunc = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + call = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p)(cfunc.address) + + itemsize = udt.np_type.itemsize + slack = 128 # the descriptor overran by 24 bytes; leave generous headroom + z = np.full(itemsize + slack, 0xAB, dtype=np.uint8) + x = np.full(itemsize + slack, 0xCD, dtype=np.uint8) + y = np.full(itemsize + slack, 0xCD, dtype=np.uint8) + xvals = np.array([1.0, 2.0, 3.0, 4.0]) + yvals = np.array([10.0, 20.0, 30.0, 40.0]) + x[:itemsize] = xvals.view(np.uint8) + y[:itemsize] = yvals.view(np.uint8) + call(z.ctypes.data, x.ctypes.data, y.ctypes.data) + + # ``any`` uses ``_second`` semantics, so the payload must be ``y``'s. + np.testing.assert_array_equal(z[:itemsize].view(np.float64), yvals) + overrun = np.flatnonzero(z[itemsize:] != 0xAB) + assert overrun.size == 0, f"wrote {overrun.size} bytes past the element at offsets {overrun}" + + # Public-path smoke: the reduce whose stack accumulator the old wrapper + # overflowed. Kept after the byte-level checks so a regression fails the + # assert above instead of reaching code that may crash the process. + v = Vector(udt, size=3) + rows = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]] + for i, row in enumerate(rows): + v[i] = row + res = v.reduce(monoid.any).new() + assert any(np.array_equal(res.value, row) for row in rows) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_array_udf_returns_new_array(): From ed7ed4432fdf76a07569125dd38ae9067bf0d982 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 00:37:27 -0500 Subject: [PATCH 3/7] Resolve array-UDT UDF returns to the operand they came from A UDF that returns an array-UDT operand as-is types as that operand's NestedArray, which went through lookup_dtype. That registry keys on the Numba type, which a flat FP64[2, 3] and a layered FP64[3][2] share, so it named whichever UDT was registered last. binary.second, monoid.any, and assignment back into the operand's vector then failed with DomainMismatch. Match a NestedArray return to its operand by identity first, and raise if two distinct operands share its type. A NestedArray that matches no operand, such as a record's array field, is no longer matched by rank alone to an operand of a different length. The "matches no input array UDT" error now describes every array operand. Fold the two wrapper byte-guard tests into one, over a 1-D and a 2-D element. A user UDF and binary.any build the same wrapper, and only distinct sentinels see the old overrun reliably: it copied y's trailing bytes, so the dropped test caught it only because the memory after its buffer happened not to match the guard. --- graphblas/core/operator/base.py | 107 +++++++++++++++------- graphblas/tests/test_op.py | 157 +++++++++++++++++--------------- 2 files changed, 157 insertions(+), 107 deletions(-) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 6b3976c5d..bac82df93 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -267,13 +267,51 @@ def _summarize_numba_typing_error(exc): return line return "Numba could not compile the function for these input types" + def _array_udt_operands(dtypes): + """Return the distinct array-UDT operands in ``dtypes``, order preserved. + + ``x`` and ``y`` are usually the same UDT, hence the dedup. It is by + identity, not by shape: a flat ``FP64[2, 3]`` and a layered + ``FP64[3][2]`` are separate DataTypes with separate GrB_Type handles, + yet Numba collapses both to the same ``nestedarray(float64, (2, 3))``. + """ + operands = [] + for d in dtypes: + if ( + d._is_udt + and d.np_type.subdtype is not None + and not any(d is seen for seen in operands) + ): + operands.append(d) + return operands + def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. When a UDF returns a tuple, Numba infers ``Tuple(...)`` rather than a Record type. Match by field count, preferring a candidate whose field - types align with the Tuple's element types. + types align with the Tuple's element types. An array return is matched + to an array-UDT operand by exact Numba type when the extents survived + (``NestedArray``), and by element type and rank when they did not. """ + if isinstance(numba_ret_type, numba.core.types.NestedArray): + # An array-UDT operand returned as-is keeps its own Numba type. + # Match it to that operand before ``lookup_dtype``, which keys on + # the Numba type and so answers with whichever UDT that shares + # it was registered last: a flat ``FP64[2, 3]`` and a layered + # ``FP64[3][2]`` both map to ``nestedarray(float64, (2, 3))``. + # Naming the other one gives the op an output type that a monoid, + # or an assignment back into the operand's vector, rejects as a + # domain mismatch. + same = [d for d in _array_udt_operands(dtypes) if d.numba_type == numba_ret_type] + if len(same) > 1: + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches more than one " + f"input array UDT ({', '.join(str(d) for d in same)}). " + f"Make the operands the same type." + ) + if same: + return same[0] try: return lookup_dtype(numba_ret_type) except (ValueError, TypeError): @@ -334,43 +372,43 @@ def _leaves(d): # of returning an operand. Numba types that as a plain Array, which # ``lookup_dtype`` doesn't recognize, so match it back to an array # UDT input by base element type and dimensionality. - candidates = [ - d - for d in dtypes - if d._is_udt - and d.np_type.subdtype is not None - and d.numba_type.dtype == numba_ret_type.dtype - and len(d.numba_type.shape) == numba_ret_type.ndim - ] - # An Array type carries ``ndim`` but not its extents, so operands - # that differ only in length are indistinguishable here. Guessing - # would hand SuiteSparse an element of the wrong size, so say so. - # Compare the UDTs rather than their shapes: a flat ``FP64[2, 3]`` - # and a layered ``FP64[3][2]`` are separate DataTypes with separate - # GrB_Type handles, yet Numba collapses both to the same shape. - unique = [] - for d in candidates: - if not any(d is seen for seen in unique): - unique.append(d) - if len(unique) > 1: + array_inputs = _array_udt_operands(dtypes) + if isinstance(numba_ret_type, numba.core.types.NestedArray): + # Its extents are known and matched no operand above, so a + # same-rank operand would be the wrong size (e.g., returning a + # record's 3-long array field over an ``FP64[4]`` operand). + candidates = [] + else: + # An Array type carries ``ndim`` but not its extents, so + # operands that differ only in length are indistinguishable + # here. Guessing would hand SuiteSparse an element of the + # wrong size, so say so. + candidates = [ + d + for d in array_inputs + if d.numba_type.dtype == numba_ret_type.dtype + and len(d.numba_type.shape) == numba_ret_type.ndim + ] + if len(candidates) > 1: raise UdfParseError( f"UDT UDF returned {numba_ret_type!r}, which matches more than one " - f"input array UDT ({', '.join(str(d) for d in unique)}). " + f"input array UDT ({', '.join(str(d) for d in candidates)}). " f"Return one of the operands, or make the operands the same type." ) - if unique: - return unique[0] + if candidates: + return candidates[0] # An array UDT went in and an array came out, but not one that # fits: name the mismatch rather than fall through to the generic # "unsupported type", whose advice the user already followed. - array_inputs = [d for d in dtypes if d._is_udt and d.np_type.subdtype is not None] if array_inputs: - d = array_inputs[0] - nested = d.numba_type + described = "; ".join( + f"{d} elements are {d.numba_type.dtype} with shape {d.numba_type.shape}" + for d in array_inputs + ) raise UdfParseError( f"UDT UDF returned {numba_ret_type!r}, which matches no input array " - f"UDT: {d} elements are {nested.dtype} with shape {nested.shape}. " - f"Return an array of that dtype and rank, or one of the operands." + f"UDT: {described}. Return an array matching one of those, or one of " + f"the operands." ) raise UdfParseError( f"UDT UDF returned an unsupported type {numba_ret_type!r}. " @@ -402,7 +440,8 @@ def _input_operand(dtype, var): """Return ``(setup_line, deref_expr, ptr_arg_type)`` for one input operand. - ``setup_line`` is the optional ``var = numba.carray(var_ptr, n)`` line - to add to the wrapper body (empty for non-UDT cases). + to add to the wrapper body, where ``n`` is 1 for a record UDT and the + element's shape for an array UDT (empty for non-UDT cases). - ``deref_expr`` is the value to pass to ``numba_func`` for this operand. - ``ptr_arg_type`` is the Numba ``CPointer(...)`` type for the wrapper signature. """ @@ -432,10 +471,10 @@ def _output_handler(return_type, numba_ret_type): ``write_kind`` is ``"record_fields"`` when the UDF returns a Tuple to be unpacked into a record output, ``"array_elements"`` for an array UDT - output, otherwise ``"direct"``. ``write_info`` is the tuple-unpack field - tuple for ``"record_fields"`` and a ``(BL, BR, zname)`` 3-tuple for - ``"direct"``; ``"array_elements"`` needs none, since ``setup_line`` - already binds ``z`` to the right shape. + output, otherwise ``"direct"``. ``write_info`` is one + ``(access_path, is_array)`` pair per leaf for ``"record_fields"`` and a + ``(BL, BR, zname)`` 3-tuple for ``"direct"``; ``"array_elements"`` needs + none, since ``setup_line`` already binds ``z`` to the right shape. """ nt = numba.types ztype = INT8 if return_type == BOOL else return_type @@ -488,7 +527,7 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): For record returns, the wrapper writes leaf-by-leaf; Numba does not compile a nested-tuple assignment to an outer record field. Array returns and array-typed record leaves slice-assign into the caller's - buffer, so a wrong-shape return cannot overrun it. It does not reach + buffer, so a wrong-shape return cannot overrun it. The error does not reach the caller either: the resulting ``ValueError`` is raised inside a cfunc, which Numba prints and swallows, so the write stops there and every leaf from that point on keeps whatever SuiteSparse had in the diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index a3eba185f..d000e3f1d 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1927,71 +1927,27 @@ def test_udt_auto_monoid(): @pytest.mark.skipif("not supports_udfs") -def test_udt_array_wrapper_stays_within_element(): - """The array-UDT cfunc wrapper writes the element payload and nothing more. - - Numba represents a ``NestedArray`` *value* as an array descriptor (data - pointer, shape, strides, ...), so the wrapper used to store that descriptor - where GraphBLAS expected only the elements: 56 bytes into the 48 an - ``FP64[6]`` element gets. GraphBLAS owns the buffer it hands the cfunc, so - the extra bytes land on memory the library allocated for something else. - - Pinned at the codegen level rather than end-to-end, and deliberately so: - nothing in the suite was ever observed to fault on the old codegen, so an - end-to-end test would not catch a regression. Measured directly instead, - driving the wrapper through ctypes over a guard-filled buffer: shape - ``(6,)`` wrote 8 bytes past a 48-byte element, shape ``(2, 3)`` wrote 24. - """ - import ctypes - - import numba - - from graphblas.core.operator.base import _get_udt_wrapper - - size = 6 - udt = dtypes.register_anonymous(np.dtype((np.float64, (size,))), "_ArrWrapPin") - - @numba.njit - def _second(x, y): # pragma: no cover (numba) - return y - - wrapper, wrapper_sig = _get_udt_wrapper(_second, udt, udt, udt) - cfunc = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) - call = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p)(cfunc.address) - - guard = -1.0 - xvals = [float(i) for i in range(size)] - yvals = [100.0 + i for i in range(size)] - z = (ctypes.c_double * (4 * size))(*([guard] * (4 * size))) - x = (ctypes.c_double * size)(*xvals) - y = (ctypes.c_double * size)(*yvals) - call(ctypes.byref(z), ctypes.byref(x), ctypes.byref(y)) - - assert list(z)[:size] == yvals - assert all(val == guard for val in list(z)[size:]), f"wrote past the UDT element: {list(z)}" - - -@pytest.mark.skipif("not supports_udfs") -def test_udt_array_any_wrapper_stays_within_element(): - """``binary.any`` on an array UDT must write one element, not Numba's array descriptor. - - ``any``, ``first``, and ``second`` are not in ``_BUILTIN_UDT_BINARY_OPS``, - so they compile through the generic ``_numba_func`` branch of - ``BinaryOp._compile_udt``. The wrapper there used to load and store the - operand as a ``NestedArray`` *value*, which Numba models as its full array - descriptor (meminfo, parent, nitems, itemsize, data, shape, strides): 56 - bytes on 64-bit for a 1-D element, regardless of payload size. SuiteSparse's - generic reduce keeps a UDT accumulator in a stack array sized to the element - (32 bytes here), so each fold overflowed it by 24 bytes; depending on the +@pytest.mark.parametrize("shape", [(4,), (3, 2)]) +def test_udt_array_wrapper_stays_within_element(shape): + """The array-UDT cfunc wrapper writes one element's bytes and nothing more. + + The wrapper used to load and store each operand as a ``NestedArray`` + *value*, which Numba models as its full array descriptor (meminfo, parent, + nitems, itemsize, data, shape, strides): 56 bytes on 64-bit for a 1-D + element and 72 for a 2-D one, whatever the payload size. Both elements here + are smaller, and the old wrapper overran each by 24 bytes. SuiteSparse's + generic reduce keeps a UDT accumulator in a stack array sized to the + element, so ``monoid.any`` overflowed it on every fold; depending on the build that clobbered a spilled pointer (segfault or SIGBUS) or silently produced a wrong answer. - Drive the compiled wrapper directly, over heap buffers with slack, so a - regression trips an assert instead of corrupting a stack frame. The two - sentinels must differ: the descriptor load/store is a byte-preserving copy - of the source element plus its trailing bytes, so if ``y``'s slack held the - same sentinel as ``z``'s guard, the overflow would rewrite ``z``'s guard - bytes with identical values and the check would be blind to it. + ``binary.any`` compiles through the generic ``_numba_func`` branch of + ``BinaryOp._compile_udt``, the same one a user's ``is_udt=True`` op takes, + so its wrapper stands in for both. Driving it over heap buffers with slack + makes a regression trip an assert instead of corrupting a stack frame. The + two sentinels must differ: the old load/store copied the source element + plus its trailing bytes, so if ``y``'s slack held ``z``'s guard value, the + overrun rewrote the guard with identical bytes and went unseen. """ import ctypes @@ -1999,10 +1955,7 @@ def test_udt_array_any_wrapper_stays_within_element(): from graphblas.core.operator.base import _get_udt_wrapper - # ``register_anonymous`` caches per np.dtype, so this may return the same - # DataType as other float64[4] tests, renamed. That is fine here: the - # wrapper below is compiled fresh and nothing asserts on cached JIT state. - udt = dtypes.register_anonymous(np.dtype((np.float64, (4,))), "_AnyOverflowArr") + udt = dtypes.register_anonymous(np.dtype((np.float64, shape)), f"_ArrWrapPin{len(shape)}D") # Mirror the generic ``_numba_func`` branch of ``BinaryOp._compile_udt``. numba_func = binary.any._numba_func @@ -2016,18 +1969,18 @@ def test_udt_array_any_wrapper_stays_within_element(): call = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p)(cfunc.address) itemsize = udt.np_type.itemsize - slack = 128 # the descriptor overran by 24 bytes; leave generous headroom + slack = 128 # the old wrapper overran by 24 bytes; leave generous headroom z = np.full(itemsize + slack, 0xAB, dtype=np.uint8) x = np.full(itemsize + slack, 0xCD, dtype=np.uint8) y = np.full(itemsize + slack, 0xCD, dtype=np.uint8) - xvals = np.array([1.0, 2.0, 3.0, 4.0]) - yvals = np.array([10.0, 20.0, 30.0, 40.0]) - x[:itemsize] = xvals.view(np.uint8) - y[:itemsize] = yvals.view(np.uint8) + xvals = np.arange(1.0, itemsize // 8 + 1).reshape(shape) + yvals = 10 * xvals + x[:itemsize] = xvals.ravel().view(np.uint8) + y[:itemsize] = yvals.ravel().view(np.uint8) call(z.ctypes.data, x.ctypes.data, y.ctypes.data) # ``any`` uses ``_second`` semantics, so the payload must be ``y``'s. - np.testing.assert_array_equal(z[:itemsize].view(np.float64), yvals) + np.testing.assert_array_equal(z[:itemsize].view(np.float64).reshape(shape), yvals) overrun = np.flatnonzero(z[itemsize:] != 0xAB) assert overrun.size == 0, f"wrote {overrun.size} bytes past the element at offsets {overrun}" @@ -2035,7 +1988,7 @@ def test_udt_array_any_wrapper_stays_within_element(): # overflowed. Kept after the byte-level checks so a regression fails the # assert above instead of reaching code that may crash the process. v = Vector(udt, size=3) - rows = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]] + rows = [xvals, yvals, xvals + yvals] for i, row in enumerate(rows): v[i] = row res = v.reduce(monoid.any).new() @@ -2124,6 +2077,64 @@ def _combine(x, y): # pragma: no cover (numba) np.testing.assert_array_equal(result[1].new().value["vec"], [8.0, 10.0, 12.0]) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_returning_operand_keeps_its_type(): + """A UDF that returns an array-UDT operand as-is produces that operand's UDT. + + Numba collapses a layered ``FP64[5][2]`` to the flat ``FP64[2, 5]``'s + ``nestedarray(float64, (2, 5))``, and ``lookup_dtype`` keys on the Numba + type, so it named whichever of the two was registered last. ``return y`` + over the flat UDT resolved to the layered one, which SuiteSparse rejects + as a domain mismatch in a monoid or on assignment into a flat vector. + """ + # Flat first, so the registry's last word on the shared Numba type is the + # layered UDT: the order that resolved the flat op to the wrong type. + flat = dtypes.register_anonymous(np.dtype((np.float64, (2, 5))), "_RetOperandFlat") + layered = dtypes.register_anonymous( + np.dtype((np.dtype((np.float64, (5,))), (2,))), "_RetOperandLayered" + ) + assert flat.numba_type == layered.numba_type + + def _second(x, y): # pragma: no cover (numba) + return y + + op = BinaryOp.register_anonymous(_second, "_ret_operand_second", is_udt=True) + assert op[flat].return_type is flat + assert op[layered].return_type is layered + # Either operand could be the one returned, so refuse to guess. + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op[flat, layered] + + v = Vector(flat, size=2) + v[0] = np.arange(10.0).reshape(2, 5) + v[1] = np.arange(10.0, 20.0).reshape(2, 5) + w = Vector(flat, size=2) + w << op(v & v) + np.testing.assert_array_equal(w[1].new().value, np.arange(10.0, 20.0).reshape(2, 5)) + # ``binary.any`` resolves its return type the same way. + res = v.reduce(monoid.any).new().value + assert any(np.array_equal(res, v[i].new().value) for i in range(2)) + + # A record's array field returned as-is also knows its extents, so it is + # not matched to an operand of the same rank but a different length. The + # field's 13-long shape must stay unique to this test: registering an + # ``FP64[13]`` UDT anywhere in the session gives ``lookup_dtype`` an answer + # for the returned field, and this resolves to it instead of raising. + rec = dtypes.register_anonymous( + np.dtype([("ro_n", np.int64), ("ro_vec", np.float64, (13,))], align=True), + "_RetOperandRec", + ) + udt14 = dtypes.register_anonymous(np.dtype((np.float64, (14,))), "_RetOperand14") + + def _field(x, y): # pragma: no cover (numba) + return x["ro_vec"] + + op2 = BinaryOp.register_anonymous(_field, "_ret_operand_field", is_udt=True) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + op2[rec, udt14] + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_auto_semiring(): From 4e7ec950bf132219395cadf59b7c63fb6f19296c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 01:10:30 -0500 Subject: [PATCH 4/7] Use unique dtype to avoid clobbering types in lookup_dtype --- graphblas/tests/test_op.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index d000e3f1d..4cad56723 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2118,11 +2118,11 @@ def _second(x, y): # pragma: no cover (numba) # A record's array field returned as-is also knows its extents, so it is # not matched to an operand of the same rank but a different length. The - # field's 13-long shape must stay unique to this test: registering an - # ``FP64[13]`` UDT anywhere in the session gives ``lookup_dtype`` an answer + # field's 19-long shape must stay unique to this test: registering an + # ``FP64[19]`` UDT anywhere in the session gives ``lookup_dtype`` an answer # for the returned field, and this resolves to it instead of raising. rec = dtypes.register_anonymous( - np.dtype([("ro_n", np.int64), ("ro_vec", np.float64, (13,))], align=True), + np.dtype([("ro_n", np.int64), ("ro_vec", np.float64, (19,))], align=True), "_RetOperandRec", ) udt14 = dtypes.register_anonymous(np.dtype((np.float64, (14,))), "_RetOperand14") From 4026f2a9aa38c64b158e7fb15c01ea964d98f42c Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 11:07:33 -0500 Subject: [PATCH 5/7] Flatten nested subarray dtypes so each UDT layout has one DataType numpy keeps np.dtype((FP64[3], (2,))) distinct from FP64[2, 3], but nothing else does: same bytes, same expanded array, same Numba type. As separate DataTypes the nested one broke auto-lifted ops, monoids, and JIT, could not mix with its flat twin, and made lookup_dtype's Numba-type keys resolve to whichever was registered last. - register_anonymous and dtypes.ss.register_new flatten nested subarrays, including inside records - Drop UDT Numba types as registry keys; lookup_dtype converts a Record or NestedArray to its numpy dtype instead - Simplify _resolve_udt_return_type: known extents resolve directly, and only an array built in the UDF is matched to an operand by rank - View, rather than cast, values whose dtype differs only in nesting; numpy's cast between the two scrambles values - Fix stale field-type rules in the UDT guide and note that records with array fields skip the JIT --- docs/user_guide/udt.rst | 12 ++- graphblas/core/dtypes.py | 70 +++++++++++++++- graphblas/core/operator/base.py | 91 ++++++-------------- graphblas/core/scalar.py | 11 +-- graphblas/core/ss/dtypes.py | 5 +- graphblas/core/utils.py | 11 ++- graphblas/tests/test_dtype.py | 109 +++++++++++++++++++++++- graphblas/tests/test_op.py | 142 +++++++++++++++++++++++++++----- graphblas/tests/test_scalar.py | 15 ++-- 9 files changed, 352 insertions(+), 114 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 67ec2cc5d..d43331257 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -32,7 +32,9 @@ A UDT is any ``numpy.dtype`` you register with python-graphblas. There are three point3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Point3") Multi-dimensional shapes work too (``np.dtype((np.float64, (2, 4)))``); the -layout is flattened row-major in C. +layout is flattened row-major in C. An array of arrays has the same layout, so +it registers as the same UDT: ``np.dtype((point3.np_type, (2,)))`` is +``FP64[2, 3]``. Record fields that nest arrays are flattened the same way. **Dataclass UDTs** are record UDTs derived from a ``@dataclass``: @@ -68,7 +70,11 @@ Field type rules: - Numeric scalar types (``int``, ``float``, ``bool``, ``complex``, the corresponding numpy scalar types) are supported. -- Strings, Python objects, and nested UDTs are not. +- Fixed-shape arrays (``("pos", np.float64, (3,))``) and records (a nested + struct) are supported as fields, and built-in operators lift through them. + In the dict and dataclass forms, write an array field as ``"FP64[3]"``; the + dict form also takes an existing UDT, such as ``{"id": int, "pos": point3}``. +- Strings and Python objects are not supported. - Dataclass annotation strings (e.g., ``"int"``) resolve through ``lookup_dtype``. @@ -236,6 +242,8 @@ JIT is skipped when: falls through to the Numba cfunc path. - A field type isn't in the numpy-to-C map (rare; the standard numeric scalar types all map). +- A record has an array-valued field. Built-in operators still lift to it + through the Numba cfunc. - The numpy layout doesn't match what a C compiler would produce. The most common case is a packed record with mixed-width fields (e.g., ``np.dtype([("a", int32), ("b", float64)])`` without ``align=True``). diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 8eb09725b..73b905306 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -10,6 +10,7 @@ if _has_numba: import numba + from numba.np.numpy_support import as_dtype as _numba_as_dtype # Default assumption unless FC32/FC64 are found in lib _supports_complex = hasattr(lib, "GrB_FC64") or hasattr(lib, "GxB_FC64") @@ -200,6 +201,10 @@ def register_anonymous(dtype, name=None): name; the SuiteSparse-side C name set at first registration does not change. See ``DataType.jit_c_name``. + Nested subarrays are flattened first, so ``np.dtype((np.dtype((np.float64, + (3,))), (2,)))``, two ``FP64[3]`` elements, registers as ``FP64[2, 3]`` and + returns the same ``DataType``. This applies inside records too. + Parameters ---------- dtype : np.dtype | dict | str | @dataclass @@ -259,6 +264,7 @@ def register_anonymous(dtype, name=None): dtype = np.dtype((base_dtype.np_type, shape)) else: raise + dtype = _flatten_subarrays(dtype) if dtype in _registry: # Always use the same object, but use the latest name. The # Python-side ``rv.name`` updates here, but the SuiteSparse-side @@ -318,9 +324,6 @@ def register_anonymous(dtype, name=None): rv = DataType(name, gb_obj, None, f"uint8_t[{dtype.itemsize}]", numba_type, dtype) _registry[gb_obj] = rv _registry[dtype] = rv - if _has_numba: - _registry[numba_type] = rv - _registry[numba_type.name] = rv # Set JIT C type definition so SuiteSparse can JIT-compile kernels for this UDT. _set_udt_jit_c_definition(rv) return rv @@ -532,6 +535,12 @@ def lookup_dtype(key, value=None): return _registry[key] except (KeyError, TypeError): pass + if _has_numba and isinstance(key, (numba.types.Record, numba.types.NestedArray)): + # A UDT's Numba type is not a registry key. Convert it to its numpy + # dtype, which is already flat, and look that up, so a layout no UDT + # has registered yet (e.g., a record's array field returned by a UDF) + # resolves the same way regardless of what else is registered. + return lookup_dtype(_numba_as_dtype(key)) if value is not None and hasattr(value, "dtype") and value.dtype in _registry: return _registry[value.dtype] # np.dtype(x) accepts some weird values; we may want to guard against some @@ -585,6 +594,61 @@ def _default_name(dtype): return repr(dtype) +def _flatten_subarrays(np_type): + """Collapse each run of nested subarray dtypes in ``np_type`` into one subarray. + + numpy keeps ``np.dtype((np.dtype((np.float64, (3,))), (2,)))`` layered, as + two ``FP64[3]`` elements, and compares it unequal to the flat + ``np.dtype((np.float64, (2, 3)))``. Nothing else tells them apart: the bytes + are the same, an array of either expands to a ``(..., 2, 3)`` float64 array, + and Numba types both as ``nestedarray(float64, (2, 3))``. Registering the flat + form gives each layout one DataType. Records are rebuilt only when a field + changes, keeping their offsets, titles, itemsize, and alignment flag. + """ + if np_type.subdtype is not None: + base, shape = np_type.subdtype + flat_base = _flatten_subarrays(base) + if flat_base.subdtype is not None: + flat_base, inner_shape = flat_base.subdtype + shape += inner_shape + elif flat_base is base: + return np_type + return np.dtype((flat_base, shape)) + if np_type.names is not None: + fields = [np_type.fields[name] for name in np_type.names] + formats = [_flatten_subarrays(field[0]) for field in fields] + if all(new is field[0] for new, field in zip(formats, fields, strict=True)): + return np_type + spec = { + "names": list(np_type.names), + "formats": formats, + "offsets": [field[1] for field in fields], + "itemsize": np_type.itemsize, + } + if any(len(field) > 2 for field in fields): + spec["titles"] = [field[2] if len(field) > 2 else None for field in fields] + return np.dtype(spec, align=np_type.isalignedstruct) + return np_type + + +def _view_if_same_layout(values, np_type): + """Return ``values`` viewed as ``np_type`` if the two differ only in subarray nesting. + + A record whose field nests subarrays has a different dtype from its + registered, flattened form, and numpy casts between the two wrongly in + either direction: ``astype``, ``np.array(..., dtype=)``, and item assignment + all scramble the field's values. The bytes already match, so a view is exact. + """ + if ( + isinstance(values, (np.ndarray, np.void)) + and values.dtype != np_type + and values.dtype.itemsize == np_type.itemsize + and _flatten_subarrays(values.dtype) == np_type + ): + return values.view(np_type) + return values + + def _dtype_to_string(dtype): """Convert a dtype to a string that can be safely evaluated to recreate the dtype. diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index bac82df93..28c2291f8 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -267,51 +267,15 @@ def _summarize_numba_typing_error(exc): return line return "Numba could not compile the function for these input types" - def _array_udt_operands(dtypes): - """Return the distinct array-UDT operands in ``dtypes``, order preserved. - - ``x`` and ``y`` are usually the same UDT, hence the dedup. It is by - identity, not by shape: a flat ``FP64[2, 3]`` and a layered - ``FP64[3][2]`` are separate DataTypes with separate GrB_Type handles, - yet Numba collapses both to the same ``nestedarray(float64, (2, 3))``. - """ - operands = [] - for d in dtypes: - if ( - d._is_udt - and d.np_type.subdtype is not None - and not any(d is seen for seen in operands) - ): - operands.append(d) - return operands - def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. When a UDF returns a tuple, Numba infers ``Tuple(...)`` rather than a Record type. Match by field count, preferring a candidate whose field - types align with the Tuple's element types. An array return is matched - to an array-UDT operand by exact Numba type when the extents survived - (``NestedArray``), and by element type and rank when they did not. + types align with the Tuple's element types. An array the UDF builds + (``x + y``) is typed without its extents, so it takes the type of the + array-UDT operand with the same element type and rank. """ - if isinstance(numba_ret_type, numba.core.types.NestedArray): - # An array-UDT operand returned as-is keeps its own Numba type. - # Match it to that operand before ``lookup_dtype``, which keys on - # the Numba type and so answers with whichever UDT that shares - # it was registered last: a flat ``FP64[2, 3]`` and a layered - # ``FP64[3][2]`` both map to ``nestedarray(float64, (2, 3))``. - # Naming the other one gives the op an output type that a monoid, - # or an assignment back into the operand's vector, rejects as a - # domain mismatch. - same = [d for d in _array_udt_operands(dtypes) if d.numba_type == numba_ret_type] - if len(same) > 1: - raise UdfParseError( - f"UDT UDF returned {numba_ret_type!r}, which matches more than one " - f"input array UDT ({', '.join(str(d) for d in same)}). " - f"Make the operands the same type." - ) - if same: - return same[0] try: return lookup_dtype(numba_ret_type) except (ValueError, TypeError): @@ -368,35 +332,40 @@ def _leaves(d): f"scalar; tuple returns are only matched to record UDTs." ) elif isinstance(numba_ret_type, numba.core.types.Array): - # A UDF over an array UDT may build its result (``x + y``) instead - # of returning an operand. Numba types that as a plain Array, which - # ``lookup_dtype`` doesn't recognize, so match it back to an array - # UDT input by base element type and dimensionality. - array_inputs = _array_udt_operands(dtypes) + # Numba types an array the UDF builds as a plain ``Array``: element + # type and rank, but no extents, which only an operand can supply. + # (An operand returned as-is keeps its ``NestedArray`` type, extents + # and all, and ``lookup_dtype`` resolved it above.) + array_inputs = list( + dict.fromkeys( # dedup: ``x`` and ``y`` are often the same UDT + d for d in dtypes if d._is_udt and d.np_type.subdtype is not None + ) + ) if isinstance(numba_ret_type, numba.core.types.NestedArray): - # Its extents are known and matched no operand above, so a - # same-rank operand would be the wrong size (e.g., returning a - # record's 3-long array field over an ``FP64[4]`` operand). - candidates = [] + # Extents survived, so ``lookup_dtype`` above named a UDT for + # them unless SuiteSparse refused to create it: builds without + # variable-length arrays cap a UDT at 128 bytes before SS 9.0 + # and 1024 from then on. Matching those extents to a same-rank + # operand would hand SuiteSparse an element of the wrong size. + shaped_like = [] else: - # An Array type carries ``ndim`` but not its extents, so - # operands that differ only in length are indistinguishable - # here. Guessing would hand SuiteSparse an element of the - # wrong size, so say so. - candidates = [ + shaped_like = [ d for d in array_inputs if d.numba_type.dtype == numba_ret_type.dtype - and len(d.numba_type.shape) == numba_ret_type.ndim + and d.numba_type.ndim == numba_ret_type.ndim ] - if len(candidates) > 1: + if len(shaped_like) == 1: + return shaped_like[0] + if shaped_like: + # Operands that differ only in length are indistinguishable + # here. Guessing would hand SuiteSparse an element of the + # wrong size, so say so. raise UdfParseError( f"UDT UDF returned {numba_ret_type!r}, which matches more than one " - f"input array UDT ({', '.join(str(d) for d in candidates)}). " + f"input array UDT ({', '.join(str(d) for d in shaped_like)}). " f"Return one of the operands, or make the operands the same type." ) - if candidates: - return candidates[0] # An array UDT went in and an array came out, but not one that # fits: name the mismatch rather than fall through to the generic # "unsupported type", whose advice the user already followed. @@ -426,12 +395,6 @@ def _array_udt_view(dtype): moves the descriptor rather than the element payload, corrupting whatever follows it (and overrunning the element outright once the descriptor is the wider of the two). - - Read from ``numba_type`` rather than ``np_type.subdtype`` so this - agrees with the type the UDF was compiled against: numpy keeps nested - subarray dtypes layered, e.g. ``FP64[5]`` inside ``[6]`` stays - ``(dtype((' subarray_after: dtype = lookup_dtype(np.dtype((dtype.np_type, array.shape[subarray_after:]))) if ownable and (not array.flags.owndata or not array.flags.writeable): diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 3a5026b6a..f8c53f423 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -10,8 +10,8 @@ from graphblas import core, dtypes from graphblas.core import _supports_udfs as supports_udfs # noqa: F401 from graphblas.core import lib -from graphblas.core.operator.udt_utils import _has_jit_set # noqa: F401 -from graphblas.core.utils import _NP2 +from graphblas.core.operator.udt_utils import _has_jit_set +from graphblas.core.utils import _NP2, values_to_numpy_buffer from graphblas.dtypes import lookup_dtype suitesparse = gb.backend == "suitesparse" @@ -215,10 +215,113 @@ def test_default_names(): ) assert _default_name(np.dtype("(29,)uint8")) == "UINT8[29]" assert _default_name(np.dtype("(3,4)bool")) == "BOOL[3, 4]" - assert _default_name(np.dtype((np.dtype("(5,)float64"), (6,)))) == "FP64[5][6]" + assert _default_name(np.dtype((np.dtype([("x", np.int32)]), (6,)))) == "{'x': INT32}[6]" assert _default_name(np.dtype("S5")) == "dtype('S5')" +def test_nested_subarrays_register_flat(): + """An array-of-arrays dtype registers as the flat array UDT with its layout. + + numpy keeps ``np.dtype((np.dtype((np.float32, (7,))), (3,)))`` layered and + compares it unequal to ``np.dtype((np.float32, (3, 7)))``, but nothing else + tells the two apart: same bytes, same expanded ``(..., 3, 7)`` array, same + Numba type. As two DataTypes they had separate GrB_Types, which SuiteSparse + refused to mix, and the layered one broke code that reads ``subdtype`` as + ``(scalar, shape)``. + """ + flat = np.dtype((np.float32, (3, 7))) + layered = np.dtype((np.dtype((np.float32, (7,))), (3,))) + assert layered != flat + udt = lookup_dtype(layered) + assert udt.np_type == flat + assert udt.name == "FP32[3, 7]" + assert lookup_dtype(flat) is udt + assert dtypes.register_anonymous(layered) is udt + assert udt == layered + assert pickle.loads(pickle.dumps(udt)) is udt + if _has_jit_set: + assert udt.jit_c_definition is not None + + deeper = np.dtype((np.dtype((np.dtype((np.float32, (7,))), (2,))), (3,))) + assert lookup_dtype(deeper).np_type == np.dtype((np.float32, (3, 2, 7))) + + +def test_nested_subarrays_register_flat_in_records(): + """Record fields flatten too, keeping offsets, itemsize, alignment, and titles.""" + point = np.dtype((np.float32, (5,))) + layered = np.dtype([("fl_id", np.int16), ("fl_tri", point, (3,))], align=True) + flat = np.dtype([("fl_id", np.int16), ("fl_tri", np.float32, (3, 5))], align=True) + udt = lookup_dtype(layered) + assert udt.np_type == flat + assert udt.np_type.isalignedstruct + assert lookup_dtype(flat) is udt + + packed = np.dtype([("fl_a", np.int8), ("fl_b", point, (2,))]) + assert lookup_dtype(packed).np_type == np.dtype( + [("fl_a", np.int8), ("fl_b", np.float32, (2, 5))] + ) + + titled = np.dtype({"names": ["fl_t"], "formats": [(point, (4,))], "titles": ["Tee"]}) + assert lookup_dtype(titled).np_type == np.dtype( + {"names": ["fl_t"], "formats": [(np.float32, (4, 5))], "titles": ["Tee"]} + ) + + # An array of records whose field nests subarrays. + assert lookup_dtype(np.dtype((layered, (2,)))).np_type == np.dtype((flat, (2,))) + + +def test_nested_subarray_record_values_load_intact(): + """Values whose record dtype nests subarrays load into the flattened UDT intact. + + numpy casts between a field of two ``FP64[3]`` elements and an + ``FP64[2, 3]`` field wrongly in both directions: ``astype``, + ``np.array(..., dtype=)``, and item assignment all scramble the values. The + bytes already match, so ingest views such values instead of casting them. + """ + point = np.dtype((np.float64, (3,))) + layered = np.dtype([("li_k", np.int64), ("li_pts", point, (2,))], align=True) + udt = lookup_dtype(layered) + vals = np.zeros(2, layered) + vals["li_k"] = [1, 2] + vals["li_pts"] = np.arange(12.0).reshape(2, 2, 3) + + # Inferring the dtype views the values too, so the pair agrees. + arr, dtype = values_to_numpy_buffer(vals) + assert dtype is udt + assert arr.dtype == udt.np_type + np.testing.assert_array_equal(arr["li_pts"], vals["li_pts"]) + + v = gb.Vector.from_coo([0, 1], vals, dtype=udt) + np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) + v = gb.Vector.from_coo([0, 1], vals) + assert v.dtype is udt + np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) + v[0] = vals[1] + np.testing.assert_array_equal(v[0].new().value["li_pts"], vals["li_pts"][1]) + s = gb.Scalar.from_value(vals[0], dtype=udt) + np.testing.assert_array_equal(s.value["li_pts"], vals["li_pts"][0]) + + +@pytest.mark.skipif("not supports_udfs") +def test_lookup_numba_type_goes_through_numpy_dtype(): + """A UDT's Numba type resolves through its numpy dtype, not a registry key. + + Keyed directly, a layered and a flat array dtype, which share a Numba type, + resolved to whichever registered last, and a layout nobody had registered + yet did not resolve at all. + """ + import numba + + nb_type = numba.typeof(np.dtype((np.int32, (2, 9)))).dtype + udt = lookup_dtype(nb_type) + assert udt.np_type == np.dtype((np.int32, (2, 9))) + assert lookup_dtype(np.dtype((np.dtype((np.int32, (9,))), (2,)))) is udt + assert lookup_dtype(nb_type) is udt + + rec = dtypes.register_anonymous(np.dtype([("nb_q", np.int32), ("nb_r", np.float32)])) + assert lookup_dtype(numba.typeof(rec.np_type).dtype) is rec + + def test_record_dtype_from_dict(): dtype = dtypes.lookup_dtype({"x": int, "y": float}) assert dtype.name == "{'x': INT64, 'y': FP64}" diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 4cad56723..82aa7a2f8 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2082,29 +2082,21 @@ def _combine(x, y): # pragma: no cover (numba) def test_udt_array_udf_returning_operand_keeps_its_type(): """A UDF that returns an array-UDT operand as-is produces that operand's UDT. - Numba collapses a layered ``FP64[5][2]`` to the flat ``FP64[2, 5]``'s - ``nestedarray(float64, (2, 5))``, and ``lookup_dtype`` keys on the Numba - type, so it named whichever of the two was registered last. ``return y`` - over the flat UDT resolved to the layered one, which SuiteSparse rejects - as a domain mismatch in a monoid or on assignment into a flat vector. + A layered ``FP64[5][2]`` and the flat ``FP64[2, 5]`` share the Numba type + ``nestedarray(float64, (2, 5))``, and ``lookup_dtype`` used to key UDTs on + it, so ``return y`` named whichever of the two had registered last. + SuiteSparse rejects the other one as a domain mismatch in a monoid or on + assignment into the operand's vector. The two now register as one UDT, and + a Numba type resolves through its numpy dtype, not registration order. """ - # Flat first, so the registry's last word on the shared Numba type is the - # layered UDT: the order that resolved the flat op to the wrong type. flat = dtypes.register_anonymous(np.dtype((np.float64, (2, 5))), "_RetOperandFlat") - layered = dtypes.register_anonymous( - np.dtype((np.dtype((np.float64, (5,))), (2,))), "_RetOperandLayered" - ) - assert flat.numba_type == layered.numba_type + assert dtypes.register_anonymous(np.dtype((np.dtype((np.float64, (5,))), (2,)))) is flat def _second(x, y): # pragma: no cover (numba) return y op = BinaryOp.register_anonymous(_second, "_ret_operand_second", is_udt=True) assert op[flat].return_type is flat - assert op[layered].return_type is layered - # Either operand could be the one returned, so refuse to guess. - with pytest.raises(UdfParseError, match="matches more than one input array UDT"): - op[flat, layered] v = Vector(flat, size=2) v[0] = np.arange(10.0).reshape(2, 5) @@ -2116,11 +2108,8 @@ def _second(x, y): # pragma: no cover (numba) res = v.reduce(monoid.any).new().value assert any(np.array_equal(res, v[i].new().value) for i in range(2)) - # A record's array field returned as-is also knows its extents, so it is - # not matched to an operand of the same rank but a different length. The - # field's 19-long shape must stay unique to this test: registering an - # ``FP64[19]`` UDT anywhere in the session gives ``lookup_dtype`` an answer - # for the returned field, and this resolves to it instead of raising. + # A record's array field returned as-is carries its extents too, so it + # resolves to the array UDT of that shape, whether or not one existed. rec = dtypes.register_anonymous( np.dtype([("ro_n", np.int64), ("ro_vec", np.float64, (19,))], align=True), "_RetOperandRec", @@ -2131,8 +2120,119 @@ def _field(x, y): # pragma: no cover (numba) return x["ro_vec"] op2 = BinaryOp.register_anonymous(_field, "_ret_operand_field", is_udt=True) + assert op2[rec, udt14].return_type is dtypes.lookup_dtype(np.dtype((np.float64, (19,)))) + r = Vector(rec, size=1) + r[0] = (1, np.arange(19.0)) + u = Vector(udt14, size=1) + u[0] = np.zeros(14) + np.testing.assert_array_equal(r.ewise_mult(u, op2).new()[0].new().value, np.arange(19.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_building_array_takes_operand_shape(): + """An array the UDF builds takes the type of the one operand it can match. + + Numba types ``x + 1.0`` as a plain ``Array`` with an element type and a + rank but no extents, so only an operand can say how long the result is. + Two same-rank candidates, or none, is an error rather than a guess. + """ + udt15 = dtypes.register_anonymous(np.dtype((np.float64, (15,))), "_BuildArr15") + udt16 = dtypes.register_anonymous(np.dtype((np.float64, (16,))), "_BuildArr16") + + def _shift(x, y): # pragma: no cover (numba) + return x + 1.0 + + op = BinaryOp.register_anonymous(_shift, "_build_arr_shift", is_udt=True) + assert op[udt15].return_type is udt15 + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op[udt15, udt16] + + def _fold(x, y): # pragma: no cover (numba) + return x.reshape(3, 5) + + op2 = BinaryOp.register_anonymous(_fold, "_build_arr_fold", is_udt=True) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + op2[udt15] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): + """A return that kept its extents is not matched to a same-rank operand. + + ``lookup_dtype`` names a UDT for known extents, registering one if no UDT + has that layout yet, so the rank matcher normally sees only an ``Array`` + the UDF built, which carries no extents. Registration can fail, though: + SuiteSparse builds without variable-length arrays reject a UDT over 128 + bytes before 9.0 and over 1024 from then on. Matching known extents to a + same-rank operand would then hand SuiteSparse an element of the wrong + size, so refuse instead. Simulated here, since this build registers the + type. + """ + import numba + + from graphblas.core.operator import base as _base + + udt21 = dtypes.register_anonymous(np.dtype((np.float64, (21,))), "_RankGuard21") + ret = numba.typeof(np.dtype((np.float64, (23,)))).dtype + # Same element type and rank, different length: what the rank matcher sees. + assert ret.dtype == udt21.numba_type.dtype + assert ret.ndim == udt21.numba_type.ndim + + real_lookup = _base.lookup_dtype + + def refuse_nested(key, value=None): + if isinstance(key, numba.core.types.NestedArray): + raise ValueError("simulated: SuiteSparse refused to register this UDT") + return real_lookup(key, value) + + monkeypatch.setattr(_base, "lookup_dtype", refuse_nested) with pytest.raises(UdfParseError, match="matches no input array UDT"): - op2[rec, udt14] + _base._resolve_udt_return_type(ret, udt21) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_declared_with_nested_subarrays_gets_builtin_ops(): + """Built-in ops work on UDTs declared with nested subarray dtypes. + + numpy reports such a dtype's ``subdtype`` as ``(inner array dtype, outer + shape)``, which the auto-lift codegen took for ``(scalar, shape)``: + ``plus``, ``eq``, and ``ainv`` failed in Numba lowering, ``monoid.plus`` and + ``agg.sum`` could not build an identity, and a UDF returning a record + operand could come back as its flat-field twin. Registration flattens these + dtypes now, so they take the ordinary path. + """ + udt = dtypes.register_anonymous(np.dtype((np.dtype((np.int64, (3,))), (4,))), "_NestedDeclArr") + assert udt.np_type == np.dtype((np.int64, (4, 3))) + vals = np.arange(12).reshape(4, 3) + v = Vector(udt, size=2) + v[0] = vals + v[1] = vals * 10 + np.testing.assert_array_equal(binary.plus(v & v).new()[1].new().value, vals * 20) + assert binary.eq(v & v).new().to_coo()[1].all() + np.testing.assert_array_equal(unary.ainv(v).new()[0].new().value, -vals) + np.testing.assert_array_equal(v.reduce(monoid.plus).new().value, vals * 11) + np.testing.assert_array_equal(v.reduce(agg.sum).new().value, vals * 11) + + point = np.dtype((np.float64, (3,))) + rec = dtypes.register_anonymous( + np.dtype([("nd_n", np.int64), ("nd_tri", point, (2,))], align=True), "_NestedDeclRec" + ) + flat_twin = np.dtype([("nd_n", np.int64), ("nd_tri", np.float64, (2, 3))], align=True) + assert dtypes.register_anonymous(flat_twin) is rec + r = Vector(rec, size=2) + r[0] = (1, np.ones((2, 3))) + r[1] = (2, np.full((2, 3), 5.0)) + res = r.reduce(monoid.plus).new().value + assert res["nd_n"] == 3 + np.testing.assert_array_equal(res["nd_tri"], np.full((2, 3), 6.0)) + + def _first(x, y): # pragma: no cover (numba) + return x + + op = BinaryOp.register_anonymous(_first, "_nested_decl_first", is_udt=True) + assert op[rec].return_type is rec @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 49e7221b4..b8c00ffec 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -259,18 +259,17 @@ def test_udt_scalar_padding_is_zeroed(): reason="SuiteSparse < 9 rejects a 240-byte UDT on builds without VLA support", ) def test_udt_scalar_nested_array_roundtrip(): - """A UDT whose element is itself an array dtype reads back at full extent. + """A UDT declared as an array of arrays reads back at full extent. - numpy keeps subarray dtypes layered rather than flattening them, so a - 6-long array of 5-long FP64 arrays reports ``subdtype`` as a 40-byte inner - dtype with shape ``(6,)``. Viewing all 240 bytes as that inner dtype raised - ValueError: "Changing the dtype to a subarray type is only supported if the - total itemsize is unchanged". + numpy keeps subarray dtypes layered, so a 6-long array of 5-long FP64 + arrays reports ``subdtype`` as a 40-byte inner dtype with shape ``(6,)``. + Viewing all 240 bytes as that inner dtype raised ValueError: "Changing the + dtype to a subarray type is only supported if the total itemsize is + unchanged". Registration now flattens it to ``FP64[6, 5]``. """ inner = np.dtype((np.float64, (5,))) udt = dtypes.register_anonymous(np.dtype((inner, (6,))), "_NestedArrayUdt") - assert udt.np_type.itemsize == 240 - assert udt.np_type.subdtype[0].subdtype is not None # genuinely nested + assert udt.np_type == np.dtype((np.float64, (6, 5))) val = np.arange(30, dtype=np.float64).reshape(6, 5) got = Scalar.from_value(val, dtype=udt).value From 16225fbfc269696c83dc4bd5dfbf87eee0433805 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 12:32:37 -0500 Subject: [PATCH 6/7] Build the nested-subarray Vector only where build can type dup_op test_nested_subarray_record_values_load_intact failed on every CI job that ran the suitesparse-vanilla backend without Numba: KeyError: "any does not work with ({'li_k': INT64, 'li_pts': FP64[2, 3]}, {'li_k': INT64, 'li_pts': FP64[2, 3]}) types" On that backend, Vector.build defaults dup_op to binary.any, and a UDT gets binary.any only by compiling it with Numba. The values_to_numpy_buffer and Scalar checks need neither, so they still run everywhere; only the Vector half is gated. --- graphblas/tests/test_dtype.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index f8c53f423..33b6d8254 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -8,7 +8,7 @@ import graphblas as gb from graphblas import core, dtypes -from graphblas.core import _supports_udfs as supports_udfs # noqa: F401 +from graphblas.core import _supports_udfs as supports_udfs from graphblas.core import lib from graphblas.core.operator.udt_utils import _has_jit_set from graphblas.core.utils import _NP2, values_to_numpy_buffer @@ -291,16 +291,20 @@ def test_nested_subarray_record_values_load_intact(): assert arr.dtype == udt.np_type np.testing.assert_array_equal(arr["li_pts"], vals["li_pts"]) - v = gb.Vector.from_coo([0, 1], vals, dtype=udt) - np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) - v = gb.Vector.from_coo([0, 1], vals) - assert v.dtype is udt - np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) - v[0] = vals[1] - np.testing.assert_array_equal(v[0].new().value["li_pts"], vals["li_pts"][1]) s = gb.Scalar.from_value(vals[0], dtype=udt) np.testing.assert_array_equal(s.value["li_pts"], vals["li_pts"][0]) + # Off SuiteSparse, ``build`` uses ``binary.any`` as its ``dup_op``, which a + # UDT gets only by compiling it with Numba. + if suitesparse or supports_udfs: + v = gb.Vector.from_coo([0, 1], vals, dtype=udt) + np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) + v = gb.Vector.from_coo([0, 1], vals) + assert v.dtype is udt + np.testing.assert_array_equal(v[1].new().value["li_pts"], vals["li_pts"][1]) + v[0] = vals[1] + np.testing.assert_array_equal(v[0].new().value["li_pts"], vals["li_pts"][1]) + @pytest.mark.skipif("not supports_udfs") def test_lookup_numba_type_goes_through_numpy_dtype(): From 73580cf3e2315ddf699f34949af215d72b5ed8ba Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 12:32:49 -0500 Subject: [PATCH 7/7] Keep the array-UDT test types under SuiteSparse 8's 128-byte cap SuiteSparse built without variable-length arrays (MSVC, so the Windows builds) rejects a UDT larger than 128 bytes with GrB_INVALID_VALUE before 9.0. The Windows job that drew psg 8.0.2 failed two tests on it: test_nested_subarrays_register_flat (FP32[3, 2, 7], 168 bytes) and test_udt_array_udf_known_extents_never_match_by_rank (FP64[21], 168 bytes). test_udt_array_udf_returning_operand_keeps_its_type registers 152- and 160-byte types as well; it is marked slow, and that job did not run slow tests. None of the three is about size, so keep their shapes and shrink the element type (INT16 or FP32) rather than skip them on SuiteSparse < 9. --- graphblas/tests/test_dtype.py | 6 ++++-- graphblas/tests/test_op.py | 13 ++++++++----- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/graphblas/tests/test_dtype.py b/graphblas/tests/test_dtype.py index 33b6d8254..2e3e17d87 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -242,8 +242,10 @@ def test_nested_subarrays_register_flat(): if _has_jit_set: assert udt.jit_c_definition is not None - deeper = np.dtype((np.dtype((np.dtype((np.float32, (7,))), (2,))), (3,))) - assert lookup_dtype(deeper).np_type == np.dtype((np.float32, (3, 2, 7))) + # INT16 keeps this under 128 bytes, the most SuiteSparse < 9 accepts for a + # UDT on builds without variable-length arrays (Windows). + deeper = np.dtype((np.dtype((np.dtype((np.int16, (7,))), (2,))), (3,))) + assert lookup_dtype(deeper).np_type == np.dtype((np.int16, (3, 2, 7))) def test_nested_subarrays_register_flat_in_records(): diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 82aa7a2f8..9bc3f2cdc 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2110,17 +2110,19 @@ def _second(x, y): # pragma: no cover (numba) # A record's array field returned as-is carries its extents too, so it # resolves to the array UDT of that shape, whether or not one existed. + # FP32 keeps each UDT under 128 bytes, the most SuiteSparse < 9 accepts on + # builds without variable-length arrays (Windows). rec = dtypes.register_anonymous( - np.dtype([("ro_n", np.int64), ("ro_vec", np.float64, (19,))], align=True), + np.dtype([("ro_n", np.int64), ("ro_vec", np.float32, (19,))], align=True), "_RetOperandRec", ) - udt14 = dtypes.register_anonymous(np.dtype((np.float64, (14,))), "_RetOperand14") + udt14 = dtypes.register_anonymous(np.dtype((np.float32, (14,))), "_RetOperand14") def _field(x, y): # pragma: no cover (numba) return x["ro_vec"] op2 = BinaryOp.register_anonymous(_field, "_ret_operand_field", is_udt=True) - assert op2[rec, udt14].return_type is dtypes.lookup_dtype(np.dtype((np.float64, (19,)))) + assert op2[rec, udt14].return_type is dtypes.lookup_dtype(np.dtype((np.float32, (19,)))) r = Vector(rec, size=1) r[0] = (1, np.arange(19.0)) u = Vector(udt14, size=1) @@ -2173,8 +2175,9 @@ def test_udt_array_udf_known_extents_never_match_by_rank(monkeypatch): from graphblas.core.operator import base as _base - udt21 = dtypes.register_anonymous(np.dtype((np.float64, (21,))), "_RankGuard21") - ret = numba.typeof(np.dtype((np.float64, (23,)))).dtype + # FP32 keeps this UDT under that 128-byte cap, so it registers everywhere. + udt21 = dtypes.register_anonymous(np.dtype((np.float32, (21,))), "_RankGuard21") + ret = numba.typeof(np.dtype((np.float32, (23,)))).dtype # Same element type and rank, different length: what the rank matcher sees. assert ret.dtype == udt21.numba_type.dtype assert ret.ndim == udt21.numba_type.ndim