Skip to content

Adds IndexBinaryOp and fixes for JIT - #576

Merged
jim22k merged 3 commits into
python-graphblas:mainfrom
jim22k:udts_jit_and_ibos2
Jul 23, 2026
Merged

jim22k merged 3 commits into
python-graphblas:mainfrom
jim22k:udts_jit_and_ibos2

Conversation

@jim22k

@jim22k jim22k commented Jun 8, 2026

Copy link
Copy Markdown
Member

Summary

This commit adds full IndexBinaryOp UDT support and a suite of SuiteSparse JIT reliability fixes.


New Features

IndexBinaryOp on User-Defined Types (UDTs)

  • IndexBinaryOp._compile_udt now works for UDT operands, including UDT theta values. A new _get_udt_wrapper_indexbinary helper in base.py generates the Numba cfunc wrapper for the 4-index + theta signature (x, ix, jx, y, iy, jy, theta → z).
  • _BoundIndexBinaryOp (the result of ibo[dtype](theta)) gained pickling support via __reduce__, proper type markers (_is_udt, commutes_to, etc.), and stores _theta for round-trip serialization.
  • ParameterizedIndexBinaryOp pickle/deserialize now round-trips the is_udt flag.

JIT C Code Generation for UDTs (udt_utils.py — new file, ~979 lines)

  • Generates C struct typedefs from numpy dtypes (record and array UDTs), enabling SuiteSparse to JIT-compile optimized kernels instead of falling back to Numba cfuncs.
  • Handles nested record UDTs, array UDTs, C-reserved name collisions (synthesizes _gbudt_NNN names automatically), and complex-type field restrictions.
  • Provides _compile_codegen — a safe exec-based helper with AST pre-validation and linecache registration for readable tracebacks.

JIT Config Auto-Repair (jit_config.py — new file)

  • fix_jit_config(): replaces conda-build-baked compiler paths with one from $CONDA_PREFIX/bin/ or sysconfig, strips stale -isysroot / -fdebug-prefix-map flags, and optionally probes the result with a trivial JIT compile.
  • _auto_fix_jit_at_import(): called automatically at gb.ss import time to silently repair common mis-configurations without user intervention.
  • jit_compiler_is_usable(): lightweight check (file existence only) for use in warnings.

DataType JIT Introspection

  • New jit_c_name and jit_c_definition properties expose the C type name and struct typedef that SuiteSparse actually registered — stable across Python-side renames.
  • _set_udt_jit_c_definition() sets GxB_JIT_C_NAME + GxB_JIT_C_DEFINITION on the GrB_Type at registration time.

register_new / register_anonymous Ergonomics

  • Both now accept a bare @dataclass class or instance as the sole argument, inferring name from the class name.

Fixes & Improvements

UDT Wrapper Refactor (base.py)

  • _get_udt_wrapper was a large monolithic function; it's been decomposed into _input_operand, _output_handler, and _compose_wrapper_body, handling record UDTs (tuple-unpack returns), array UDTs, and BOOL-as-INT8 coercion correctly in all combinations.
  • _resolve_udt_return_type resolves Numba BaseTuple return types back to the matching record UDT by field arity and type alignment, with actionable UdfParseError messages when the match fails.

Better UDF Error Messages

  • UdfParseError docstring updated to document wrapping of TypingError, LoweringError, UnsupportedError.
  • _compile_udf_for_udt catches the full NumbaError hierarchy and surfaces the actionable diagnostic line (via _summarize_numba_typing_error) rather than a raw Numba traceback.
  • New NoJITWarning warning class emitted once per process when an auto-lifted UDT op falls back to the cfunc path (JIT compiler unusable, jit_c_control off, or UDT not C-representable).

Scalar / UDT Thunk Binding

  • TypedBuiltinIndexBinaryOp.__call__ stores theta_value before wrapping in a Scalar, fixing a round-trip issue for UDT thunks where Scalar.from_value couldn't infer the dtype from a raw numpy array.

Tests & Docs

  • test_indexbinary.py expanded with UDT cases, pickling, and semiring round-trips.
  • test_op.py added ~1,200 lines covering UDT operator compilation, JIT codegen, and error paths.
  • test_ssjit.py expanded with JIT auto-fix, NoJITWarning, and C typedef generation tests.
  • test_dtype.py and test_pickle.py extended for new jit_c_name/jit_c_definition and register_new(dataclass) behaviour.
  • New user guide pages: udt.rst (UDT reference) and updated udf.rst and operators.rst.

@jim22k
jim22k requested a review from eriknw June 8, 2026 14:48
@jim22k
jim22k merged commit 6f1eb02 into python-graphblas:main Jul 23, 2026
8 of 16 checks passed
eriknw added a commit that referenced this pull request Nov 3, 2026
…ier collides with

SuiteSparse's JIT puts a UDT's C name in each kernel beside its own
identifiers, and that name was the UDT's Python name. A UDT named like
one of them did not compile: SuiteSparse raised JitError (with no
message) for that op and set jit_c_control to 'load', so the process
compiled no more kernels, although the code comment and udt.rst say such
a UDT falls back to the cfunc. Main did the same. Names that failed, each
in a fresh process: z and x (parameters of our op definitions), A, B, C,
M, p, i, j, k, Ax, Cx, cnz, tid and result (locals of SuiteSparse's
templates; which kernels fail depends on the name), div, abs, exp,
round, free, index and printf (declared by the C headers), and GB_void
and GrB_Info (SuiteSparse's). A list of reserved names cannot cover the
template locals.

Now the C name of a named UDT is gbudt_ and its name (_pick_c_type_name),
so every one of those compiles and computes; a C keyword as a name
(class) now takes the JIT too, where it got a synthetic name. A name
that is not a C identifier, or would make a C name of 128 characters or
more, gets a synthetic _gbudt_NNN name: SuiteSparse refuses such a long
name and keeps the type with an empty one, and its first kernel (A.T.new()
say) raised JitError, on main too for a UDT name of 128 characters or
more. jit_c_name is unreleased (it came with #576, after 2025.2.0), so
changing it breaks no release.

udt.rst's example and the tests that pin C names say gbudt_. The test
of an op whose C name SuiteSparse refuses now uses a 119-character UDT
name: its C name (125 characters) fits, and plus_ and that does not. New
tests: UDTs named z, A, p, i, div and result take the JIT for ewise_mult,
reduce and mxv; UDT names of 121, 122 and 200 characters get C names
that fit, and transpose under the JIT.
eriknw added a commit that referenced this pull request Nov 13, 2026
* Skip the JIT for a UDT field named like a standard C macro

A record field named like an object-like macro of the C headers the JIT
kernels include, such as INT_MAX, SIZE_MAX, BUFSIZ, EXIT_FAILURE or
PRId64, is expanded by the preprocessor inside the struct, so the kernel
did not compile. SuiteSparse 9.4 and later raised JitError (with no
message) for that op and set jit_c_control to 'load', so the process
compiled no more kernels. Main did the same. The comment beside
_C_RESERVED said SuiteSparse swallowed the failure silently.

_C_RESERVED now also has the object-like macros the C standard defines in
<limits.h>, <stdint.h>, <inttypes.h>, <stdio.h>, <stdlib.h>, <errno.h>,
<math.h> and <stdatomic.h>, so such a field skips the JIT: the op runs
through the Numba cfunc with a NoJITWarning, as udt.rst says. Macros a
platform adds (POSIX errno codes and the like) are not listed and still
fail to compile.

* Give each UDT the C name gbudt_ and its name, which no kernel identifier collides with

SuiteSparse's JIT puts a UDT's C name in each kernel beside its own
identifiers, and that name was the UDT's Python name. A UDT named like
one of them did not compile: SuiteSparse raised JitError (with no
message) for that op and set jit_c_control to 'load', so the process
compiled no more kernels, although the code comment and udt.rst say such
a UDT falls back to the cfunc. Main did the same. Names that failed, each
in a fresh process: z and x (parameters of our op definitions), A, B, C,
M, p, i, j, k, Ax, Cx, cnz, tid and result (locals of SuiteSparse's
templates; which kernels fail depends on the name), div, abs, exp,
round, free, index and printf (declared by the C headers), and GB_void
and GrB_Info (SuiteSparse's). A list of reserved names cannot cover the
template locals.

Now the C name of a named UDT is gbudt_ and its name (_pick_c_type_name),
so every one of those compiles and computes; a C keyword as a name
(class) now takes the JIT too, where it got a synthetic name. A name
that is not a C identifier, or would make a C name of 128 characters or
more, gets a synthetic _gbudt_NNN name: SuiteSparse refuses such a long
name and keeps the type with an empty one, and its first kernel (A.T.new()
say) raised JitError, on main too for a UDT name of 128 characters or
more. jit_c_name is unreleased (it came with #576, after 2025.2.0), so
changing it breaks no release.

udt.rst's example and the tests that pin C names say gbudt_. The test
of an op whose C name SuiteSparse refuses now uses a 119-character UDT
name: its C name (125 characters) fits, and plus_ and that does not. New
tests: UDTs named z, A, p, i, div and result take the JIT for ewise_mult,
reduce and mxv; UDT names of 121, 122 and 200 characters get C names
that fit, and transpose under the JIT.

* Say in the JIT docs how mixed UDT ops compile and how long the JIT stays on

udt.rst said SuiteSparse compiles a kernel for each op[udt]. A lifted op
on a UDT and a built-in type, or on two UDT types that pair, compiles
too, and these mixed cases keep the Numba cfunc without a warning: an
array UDT that broadcasts to another shape, a record with an array field,
eq and ne on two types, and an op on one UDT that gives another type
(v / v on INT8[3]). The JIT section now says so, with the mixed C name.

UDT work raises the default 'run' to 'on' the first time it needs a
kernel, for the rest of the process and for every operation, built-in
ones included; a mixed op does too now that it has C source. init.rst
and udt.rst read as if the raise were per op. The concurrency bullet now
says such processes race too.

init.rst says python-graphblas checks the C compiler under "on" and
falls back to "run". Only a configured "on" is checked: an "on" written
later to gb.ss.config is taken as it is, and with a compiler that does
not work the first new kernel raises JitError and SuiteSparse lowers
the control to "load" (checked with /usr/bin/false as the compiler). It
says so, and that gb.ss.fix_jit_config() checks.

* Keep a name starting with _gb as a UDT's C name, unprefixed

python-graphblas names the types it makes itself with _gb (the
synthetic _gbudt_NNN and _gbnest_NNN, and the tensor engine's
_gbtensor_fp64xR), and no kernel identifier starts with it. Only
_gbudt_ and _gbnest_ were kept as they are, so the tensor engine's type
became gbudt__gbtensor_fp64x16, which its test of the name it claims
caught. Keep every _gb name.

* Name a C-source type by its own C name in the ops lifted to it

The prefix that keeps a numpy UDT's C name clear of kernel identifiers
was put on the name of every type an op definition names, and a type
registered from C source (dtypes.ss.register_new) has no jit_c_name, so
its Python name stood in, and got the prefix too: binary.plus[t] on a
type cpt_t defined plus_gbudt_cpt_t on a type gbudt_cpt_t that
SuiteSparse does not know. The kernel failed to compile, JitError, and
the JIT was 'load' for the rest of the process. Before the prefix it was
plus_cpt_t, and right.

Now a C-source type's jit_c_name is the name it was registered under
(and jit_c_definition its definition), and the op code generators name a
type by its jit_c_name as it is (_dtype_c_typedef, is_c_name=True), so
only a numpy UDT's Python name goes through _pick_c_type_name. That
function no longer needs to map a pinned gbudt_ name to itself, so a
user's name gbudt_Foo is now gbudt_gbudt_Foo and cannot collide with one
made here (rBj N1). A name starting with _gb, which python-graphblas
keeps for the types it makes itself (_gbudt_NNN, _gbnest_NNN, the tensor
engine's _gbtensor_...), is still kept as it is.

* Give a UDT a synthetic C name when its ops' names would not fit

A numpy UDT kept the C name gbudt_ and its name while that was under
SuiteSparse's 128 characters. But its one-type ops are named {op}_{type},
so a name of 113 to 121 characters left a C name its ops did not fit
beside: SuiteSparse refused floordiv_ or plus_ and it, the op kept its
cfunc, and the NoJITWarning said the UDT was not expressible as a C
struct, blaming its fields.

Now the C name may be at most _max_type_c_name_len() characters, which
leaves room for the longest prefix of an op built from it
(GxB_MAX_NAME_LEN - 1 - len("floordiv_") = 118, from the lifted ops, eq
and ne), and a longer name gets a synthetic _gbudt_NNN, so every op on
the UDT keeps its JIT. Mixed ops shorten their own names already.

The test of a refused op name now registers a type from C source, whose
125-character name is the user's, so plus_ and eq_ and it still do not
fit. The long-name test checks names of 112 (C name 118, kept), 113, 119
and 200 characters: floordiv, plus, eq and ainv all have C source, and
transpose and plus run under the JIT.

* Say that C source names a numpy UDT by its jit_c_name

A UDT from numpy is known to SuiteSparse as gbudt_ and its name, so C
source passed to gb.binary.ss.register_new and its siblings must name it
so: written with its Python name, the operator failed to compile, raised
JitError and left the JIT at 'load' for the rest of the process (rBj
L1). The five register_new docstrings and udt.rst now say to use
jit_c_name, and that a type registered from C source keeps its name.
udt.rst also gives the new threshold for a synthetic C name (over 118
characters) and the exception for a name starting with _gb, which is
kept as it is.

* Skip the JIT for a field named like a signal macro or with GraphBLAS's prefix

A record field named like one of <signal.h>'s nine C-standard macros
(SIGINT, SIG_DFL, ...), which the C library headers bring into the
kernels on macOS, or like one of GraphBLAS's own macros (GrB_INDEX_MAX,
GxB_IMPLEMENTATION_MAJOR) was expanded by the preprocessor inside the
struct, so the kernel failed to compile: JitError, then 'load' for the
rest of the process (rBj N2). udt.rst said only macros beyond the C
standard's still failed.

Now _C_RESERVED has the nine, and a field that starts with GB_, GrB_ or
GxB_ skips the JIT too: GraphBLAS keeps those prefixes for its own
names, and each version adds more of them. The op runs through the
Numba cfunc with a NoJITWarning, as for the other macros.

* Say how long a C-source type's name may be, and that gbudt_ is usual

The five ss.register_new docstrings said a numpy UDT's C name is gbudt_
and its name; a synthetic _gbudt_NNN is the exception, so they now say
usually. A type registered from C source keeps the name given, and with
119 to 127 characters the built-in ops lifted to it get C names
SuiteSparse refuses, so they lose the JIT, with a NoJITWarning that
blames the fields. udt.rst now says to keep such a name to 118
characters (rBj round 2).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant