diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 18af2cf3e..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``. @@ -173,6 +179,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. @@ -216,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 287e525d5..28c2291f8 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -272,7 +272,9 @@ def _resolve_udt_return_type(numba_ret_type, *dtypes): 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 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. """ try: return lookup_dtype(numba_ret_type) @@ -329,17 +331,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): + # 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): + # 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: + shaped_like = [ + d + for d in array_inputs + if d.numba_type.dtype == numba_ret_type.dtype + and d.numba_type.ndim == numba_ret_type.ndim + ] + 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 shaped_like)}). " + f"Return one of the operands, or make the operands the same type." + ) + # 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. + if array_inputs: + 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: {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}. " 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). + """ + nested = dtype.numba_type + return nested.dtype, nested.shape + 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, 1)`` line - to add to the wrapper body (empty for non-record cases). + - ``setup_line`` is the optional ``var = numba.carray(var_ptr, n)`` line + 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. """ @@ -351,8 +416,12 @@ def _input_operand(dtype, var): f"{var}[0]", nt.CPointer(dtype.numba_type), ) - # Array UDT: pass the raw pointer. The UDF carrays it if it wants. - return "", f"{var}_ptr", nt.CPointer(dtype.numba_type) + base_numba, shape = _array_udt_view(dtype) + return ( + f" {var} = numba.carray({var}_ptr, {shape})\n", + var, + nt.CPointer(base_numba), + ) if dtype == BOOL: # Numba can't compile bool ptrs (numba/numba#5395); expose them # as int8 and cast on deref. @@ -364,9 +433,11 @@ def _output_handler(return_type, numba_ret_type): """Return ``(setup_line, ret_ptr_type, write_kind, write_info)``. ``write_kind`` is ``"record_fields"`` when the UDF returns a Tuple to be - unpacked into a record output, otherwise ``"direct"``. - ``write_info`` is the tuple-unpack field tuple for ``"record_fields"``, - or a ``(BL, BR, zname)`` 3-tuple for ``"direct"``. + unpacked into a record output, ``"array_elements"`` for an array UDT + 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 @@ -378,18 +449,21 @@ def _output_handler(return_type, numba_ret_type): and return_type._is_udt and return_type.np_type.names is not None ): - # ``write_info`` is the list of Python access paths to each leaf - # field (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for - # nested). The wrapper iterates these to write the flat tuple - # ``_result`` back leaf-by-leaf. + # ``write_info`` pairs each leaf field's Python access path + # (``"['a']"`` for flat, ``"['outer']['inner_a']"`` for nested) + # with whether that leaf is array-typed. The wrapper iterates + # these to write the flat tuple ``_result`` back leaf-by-leaf. from .udt_utils import _iter_record_leaves - leaf_python_paths = tuple(py for py, _c, _d in _iter_record_leaves(return_type.np_type)) + leaves = tuple( + (py, d.subdtype is not None) + for py, _c, d in _iter_record_leaves(return_type.np_type) + ) return ( " z = numba.carray(z_ptr, 1)\n", ret_ptr_type, "record_fields", - leaf_python_paths, + leaves, ) if return_type._is_udt: if return_type.np_type.subdtype is None: @@ -399,9 +473,13 @@ def _output_handler(return_type, numba_ret_type): "direct", ("", "", "z[0]"), ) - # Array UDT: write via z_ptr[0]. The UDF receives the carray and - # writes in place. - return "", ret_ptr_type, "direct", ("", "[0]", "z_ptr[0]") + base_numba, shape = _array_udt_view(return_type) + return ( + f" z = numba.carray(z_ptr, {shape})\n", + nt.CPointer(base_numba), + "array_elements", + None, + ) if return_type == BOOL: return "", ret_ptr_type, "direct", ("bool(", ")", "z_ptr[0]") return "", ret_ptr_type, "direct", ("", "", "z_ptr[0]") @@ -410,11 +488,29 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): """Assemble the Python source for a UDT cfunc wrapper. For record returns, the wrapper writes leaf-by-leaf; Numba does not - compile a nested-tuple assignment to an outer record field. + 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. 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 + buffer. """ + if zkind == "array_elements": + return ( + f"{signature_line}\n" + f"{body_setup}" + f" _result = {call_expr}\n" + f" z[:] = _result\n" + ) if zkind == "record_fields": + # Array-typed leaves slice-assign. Numba's record-field setitem + # copies the destination's extent regardless of the source's, so a + # short source is read past its end and a long one is silently + # truncated; ``[:]`` makes both a shape error instead. field_assigns = "".join( - f" z[0]{path} = _result[{i}]\n" for i, path in enumerate(zinfo) + f" z[0]{path}{'[:]' if is_array else ''} = _result[{i}]\n" + for i, (path, is_array) in enumerate(zinfo) ) return ( f"{signature_line}\n" diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index 3e2a60d45..b55a17f13 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -349,9 +349,11 @@ def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``. Only record-in-record nesting is recognised. Array-typed fields are - yielded as a single leaf (the codegen treats them opaquely; arithmetic - on an array-typed sub-field isn't supported by either the cfunc or the - JIT path today). + yielded as a single leaf, which the Numba path handles: the generated + expression operates on the whole sub-array and the wrapper slice-assigns + it back. The JIT path never sees such a record, because a subarray dtype + has no entry in ``NP_TO_C_TYPES`` and ``_udt_c_typedef`` bails out, so + the op keeps the cfunc for every call. """ for name in np_type.names: field_dtype = np_type.fields[name][0] diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 5b5e56299..9755acb3c 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -8,6 +8,7 @@ from ..exceptions import EmptyObject, check_status from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils from .base import BaseExpression, BaseType, call +from .dtypes import _view_if_same_layout from .expr import AmbiguousAssignOrExtract from .operator import get_typed_op from .utils import _Pointer, output_type, wrapdoc @@ -356,15 +357,7 @@ def value(self): rv = np.array(ffi.buffer(scalar.gb_obj[0 : np_type.itemsize])) if np_type.subdtype is None: return rv.view(np_type)[0] - # Collapse nested subarray dtypes before viewing. numpy keeps them - # layered, so a 6-long array of 5-long FP64 arrays has - # ``subdtype == (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..2e3e17d87 100644 --- a/graphblas/tests/test_dtype.py +++ b/graphblas/tests/test_dtype.py @@ -8,10 +8,10 @@ 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 # 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,119 @@ 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 + + # 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(): + """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"]) + + 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(): + """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 de1c01b88..9bc3f2cdc 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1926,6 +1926,318 @@ def test_udt_auto_monoid(): assert any_res in ((7, 8.0), (11, 12.0)) +@pytest.mark.skipif("not supports_udfs") +@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. + + ``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 + + import numba + + from graphblas.core.operator.base import _get_udt_wrapper + + 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 + 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 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.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).reshape(shape), 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 = [xvals, yvals, xvals + yvals] + 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(): + """An array-UDT UDF may build its result instead of returning an operand. + + The wrapper hands the UDF a numpy view of each element in the UDT's + declared shape, so ordinary array expressions work and the result is copied + back element-wise. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (8,))), "_ArrRetUDT") + v = Vector(udt, size=2) + v[0] = np.arange(8.0) + v[1] = np.arange(8.0, 16.0) + w = Vector(udt, size=2) + w[0] = np.full(8, 100.0) + w[1] = np.full(8, 200.0) + + def _add(x, y): + return x + y # pragma: no cover (numba) + + add_op = BinaryOp.register_anonymous(_add, "_arr_ret_add", is_udt=True) + result = add_op(v & w).new() + np.testing.assert_array_equal(result[0].new().value, np.arange(8.0) + 100.0) + np.testing.assert_array_equal(result[1].new().value, np.arange(8.0, 16.0) + 200.0) + + def _double(x): + return x * 2 # pragma: no cover (numba) + + double_op = UnaryOp.register_anonymous(_double, "_arr_ret_double", is_udt=True) + np.testing.assert_array_equal(double_op(v).new()[0].new().value, np.arange(8.0) * 2) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_multidim_array_keeps_shape_in_udf(): + """A multi-dimensional array UDT reaches the UDF in its declared shape. + + The wrapper builds the operand with ``numba.carray(ptr, shape)``, so 2-D + indexing and ``.shape`` work. Addressing it as a flat run of elements + would still be memory-safe but would silently drop that metadata. + """ + udt2d = dtypes.register_anonymous(np.dtype((np.float64, (2, 3))), "_ArrRet2D") + m = Vector(udt2d, size=1) + m[0] = np.arange(6.0).reshape(2, 3) + n = Vector(udt2d, size=1) + n[0] = np.full((2, 3), 10.0) + + def _add_corner(x, y): + # Fails to compile unless `x` really is 2-D with shape metadata. + return x + y[0, 0] + x.shape[1] # pragma: no cover (numba) + + add_2d = BinaryOp.register_anonymous(_add_corner, "_arr_ret_add_2d", is_udt=True) + np.testing.assert_array_equal( + add_2d(m & n).new()[0].new().value, np.arange(6.0).reshape(2, 3) + 10.0 + 3 + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_array_field_roundtrip(): + """A record UDT with an array field writes exactly that field's extent. + + Numba's record-field setitem copies the *destination* extent whatever the + source's length, so a short return used to read past the end of the source + array. The wrapper slice-assigns array leaves to make that a shape error. + """ + spec = np.dtype([("count", np.int64), ("vec", np.float64, (3,))], align=True) + udt = dtypes.register_anonymous(spec, "_RecArrField") + v = Vector(udt, size=2) + v[0] = (1, [1.0, 2.0, 3.0]) + v[1] = (2, [4.0, 5.0, 6.0]) + + def _combine(x, y): # pragma: no cover (numba) + return (x["count"] + y["count"], x["vec"] + y["vec"]) + + op = BinaryOp.register_anonymous(_combine, "_rec_arr_field", is_udt=True) + result = v.ewise_mult(v, op).new() + assert result[0].new().value["count"] == 2 + np.testing.assert_array_equal(result[0].new().value["vec"], [2.0, 4.0, 6.0]) + 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. + + 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 = dtypes.register_anonymous(np.dtype((np.float64, (2, 5))), "_RetOperandFlat") + 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 + + 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 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.float32, (19,))], align=True), + "_RetOperandRec", + ) + 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.float32, (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 + + # 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 + + 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"): + _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") @pytest.mark.slow def test_udt_auto_semiring(): 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