Skip to content

Adds IndexBinaryOp and fixes for JIT#576

Open
jim22k wants to merge 1 commit into
python-graphblas:mainfrom
jim22k:udts_jit_and_ibos2
Open

Adds IndexBinaryOp and fixes for JIT#576
jim22k wants to merge 1 commit 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
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