Skip to content
Merged
17 changes: 12 additions & 5 deletions docs/user_guide/init.rst
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,9 @@ often several times slower (``Matrix.reduce_rowwise`` with ``lor`` on ``BOOL`` o
``plus`` on ``UINT64`` ran 2.7 to 8 times slower on 10 million entries). The JIT control,
``gb.ss.config["jit_c_control"]``, is ``"run"`` by default: run only kernels already
loaded in the process, and compile nothing. UDT operators, and types and operators
registered from C source, raise this default to ``"on"`` when they need a kernel (see
:doc:`udt`), but leave a control you set as it is. To opt in for everything, set the
registered from C source, raise this default to ``"on"`` the first time they need a
kernel, for the rest of the process and for every operation (see :doc:`udt`), but leave
a control you set as it is. To opt in for everything, set the
control to ``"on"`` (compile, load and run):

.. code-block:: bash
Expand All @@ -96,7 +97,11 @@ later to ``gb.ss.config`` is that same default, so to keep UDT work from compili
directory and compiler), and falls back to ``"run"``, with no error, where the JIT cannot
work: SuiteSparse built without the JIT (as in the PyPI wheels), SuiteSparse 7, no working
C compiler (in an activated conda environment, ``conda install c-compiler``), or a cache
directory it cannot write to. On Windows, SuiteSparse built with MSVC (as conda-forge
directory it cannot write to. That check is for a configured ``"on"`` only: an ``"on"``
written later to ``gb.ss.config`` is taken as it is, and with a C compiler that does not
work, the first operation that needs a new kernel raises ``JitError``, and SuiteSparse
lowers the control to ``"load"``. ``gb.ss.fix_jit_config()`` repairs the compiler setting
where it can and checks it. On Windows, SuiteSparse built with MSVC (as conda-forge
builds it) compiles JIT kernels only through cmake, and python-graphblas does not turn the
JIT on there yet, so a configured ``"on"`` becomes ``"run"``.

Expand All @@ -109,8 +114,10 @@ With the JIT on:
(`GraphBLAS issue #449 <https://github.com/DrTimothyAldenDavis/GraphBLAS/issues/449>`_).
With multiprocessing, dask, pytest-xdist and the like, set
``GRAPHBLAS_JIT_C_CONTROL=run`` in the environment the workers start in, so that none
of them compiles: UDT work raises only the default ``"run"``, not one you set. That
costs UDTs their compiled kernels: UDT operators call their Numba functions instead
of them compiles: UDT work raises only the default ``"run"``, not one you set. (At the
default, the first UDT operator that needs a kernel turns the JIT on for the whole
process, so such processes race too.) Setting ``"run"`` costs UDTs their compiled
kernels: UDT operators call their Numba functions instead
(typically 2 to 3 times slower for elementwise operations), and registering types or
operators from C source fails. To keep the kernels, give each process its own cache
instead, by setting the ``GRAPHBLAS_CACHE_PATH`` environment variable before graphblas
Expand Down
58 changes: 40 additions & 18 deletions docs/user_guide/udt.rst
Original file line number Diff line number Diff line change
Expand Up @@ -376,9 +376,11 @@ from Numba's typing pass instead of a 200-line traceback.
JIT and introspection
---------------------

When the UDT name and all field names are valid C identifiers (not C reserved
words), and the field types map to C primitives, SuiteSparse JIT-compiles a
dedicated C kernel for each ``op[udt]`` lookup. The kernel is cached on disk
When all field names are valid C identifiers (not C reserved words), and the
field types map to C primitives, SuiteSparse JIT-compiles a dedicated C kernel
for each ``op[udt]`` lookup. A lifted op on a UDT and a built-in type, or on two
UDT types that pair, compiles too, under a name such as
``plus_7_int64_t_gbudt_Edge``. The kernel is cached on disk
and reused across processes (keyed by content hash, so renaming a UDT doesn't
invalidate). JIT lets SuiteSparse inline the kernel into its eWise and reduce
templates, eliminating the per-element function-call overhead the Numba
Expand All @@ -389,11 +391,11 @@ Inspect what SuiteSparse is JIT-ing from Python:

.. code-block:: python

udt.jit_c_name # 'Edge', or None if not JIT-able
udt.jit_c_definition # 'typedef struct { ... } Edge ;'
udt.jit_c_name # 'gbudt_Edge', or None if not JIT-able
udt.jit_c_definition # 'typedef struct { ... } gbudt_Edge ;'

op = binary.plus[udt]
op.jit_c_name # 'plus_Edge'
op.jit_c_name # 'plus_gbudt_Edge'
op.jit_c_source # full C source SS will compile

monoid.plus[udt].jit_c_source # same kernel
Expand All @@ -404,10 +406,11 @@ Inspect what SuiteSparse is JIT-ing from Python:

JIT is skipped when:

- A field name (or the UDT name, if you supplied one) is a C reserved word
(``class``, ``return``, ...) or a stdlib macro/typedef pulled in by
``GraphBLAS.h`` (``NULL``, ``FILE``, ``M_PI``, ``complex``, ...). The op
falls through to the Numba cfunc path.
- A field name is a C reserved word (``class``, ``return``, ...), a macro or
typedef of the C headers the kernels include (``NULL``, ``FILE``, ``M_PI``,
``complex``, ``INT_MAX``, ``SIZE_MAX``, ``BUFSIZ``, ``SIGINT``, ...), or
starts with ``GB_``, ``GrB_`` or ``GxB_``, which GraphBLAS keeps for its own
names. The op 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
Expand All @@ -418,11 +421,25 @@ JIT is skipped when:
Pass ``align=True`` to ``np.dtype``, or use the dict / dataclass form,
which auto-aligns.
- The JIT compiler isn't usable after auto-fix (see below).

Anonymous UDTs (no ``name=`` argument) still take the JIT path: a synthetic
``_gbudt_NNN`` C name is minted so SuiteSparse always has a registerable
identifier. Pass ``name=`` only for readable JIT cache filenames and
introspection.
- An op on a UDT and another type is an array UDT that broadcasts to another
shape, a record with an array-valued field, or ``eq`` or ``ne``; or an op on
one UDT gives another type (``v / v`` on ``INT8[3]`` is ``FP64[3]``). These
keep the Numba cfunc without a warning.

SuiteSparse knows a UDT by ``gbudt_`` and its name, which cannot collide with
an identifier in the C code of a kernel, whatever the name. A name starting
with ``_gb``, which python-graphblas keeps for the types it makes itself, is
the exception: it is the C name as it is. A UDT without a name, or whose name
is not a C identifier or would make that C name too long to leave room for its
ops' names (over 118 characters), gets a synthetic ``_gbudt_NNN`` C name and
still takes the JIT path. Pass ``name=`` only for
readable JIT cache filenames and introspection. To name such a UDT in C source
of your own, as for ``gb.binary.ss.register_new``, use its ``jit_c_name``, not
its Python name, which SuiteSparse does not know. A type registered from C
source (``gb.dtypes.ss.register_new``) keeps the name you give it, so give it
118 characters or fewer: the built-in operators lifted to it are named with
that name after them (``floordiv_``, ``plus_``, ...), and a longer one leaves
them names SuiteSparse refuses, so they run without the JIT.

When the cause is the UDT itself (the first four cases above), the first time
auto-lift produces a non-JIT'd op for a given ``(op, dtype)`` pair in a
Expand All @@ -439,7 +456,9 @@ category or by message::
warnings.filterwarnings("ignore", message="UDT operator running without JIT")

In every skipped case the operation still works correctly via the Numba
function-pointer path; you only lose the JIT speedup.
function-pointer path; you only lose the JIT speedup. A field named like a
macro that only some platform's C headers define (beyond the C standard's)
is not skipped, and its kernel fails to compile; rename the field.

JIT compiler auto-fix
~~~~~~~~~~~~~~~~~~~~~
Expand All @@ -461,9 +480,12 @@ later operations compute.
``jit_c_control`` starts at ``graphblas.config["jit_c_control"]``, ``'run'`` by
default (run cached kernels only; no compile, no load from disk); see
:ref:`suitesparse_jit`. That default ``'run'`` is raised to ``'on'`` (compile,
load, and run) by each of these, since each asks SuiteSparse to compile C source:
load, and run) by each of these, since each asks SuiteSparse to compile C source.
It stays ``'on'`` for the rest of the process, for every operation, built-in ones
included:

- a UDT operator that gets C source, such as the first ``binary.plus[udt]``;
- a UDT operator that gets C source, such as the first ``binary.plus[udt]`` or
``binary.plus[udt, "INT64"]``;
- ``gb.dtypes.ss.register_new``, which registers a type from a C typedef;
- ``gb.binary.ss.register_new`` and its ``unary``, ``indexunary``, ``select``,
and ``indexbinary`` siblings, which register an operator from a C definition.
Expand Down
25 changes: 13 additions & 12 deletions graphblas/core/dtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,14 +108,14 @@ def _deserialize(name, dtype, is_anonymous):
def _get_udt_c_type_info(datatype):
"""Return ``(c_name, c_typedef)`` for a UDT, or ``None`` if not expressible in C.

The returned ``c_name`` may differ from ``datatype.name``: if the user
didn't supply a name, or supplied one that isn't a valid C identifier,
``_udt_c_typedef`` synthesizes a fresh ``_gbudt_NNN`` name so the UDT
can still take the JIT path. Returns ``None`` only for the cases that
remain unrepresentable in C: an unsupported field type (object,
datetime, string), or a field name that collides with a C reserved
word or isn't an identifier. Supports record UDTs (including nested
records) and array UDTs (e.g., ``FP64[3]``).
The returned ``c_name`` is ``gbudt_`` and ``datatype.name``, or a fresh
``_gbudt_NNN`` name when the user didn't supply a name, or supplied one
that isn't a valid C identifier or is too long for SuiteSparse
(``udt_codegen._pick_c_type_name``), so the UDT can still take the JIT
path. Returns ``None`` only for the cases that remain unrepresentable in
C: an unsupported field type (object, datetime, string), or a field name
that collides with a C reserved word or isn't an identifier. Supports
record UDTs (including nested records) and array UDTs (e.g., ``FP64[3]``).
"""
from .operator.udt_codegen import _udt_c_typedef

Expand Down Expand Up @@ -159,7 +159,8 @@ def _set_udt_jit_c_definition(datatype):
# A name already given to another layout (two UDTs registered under one
# name) gets a fresh one: SuiteSparse would share their JIT kernels.
while _jit_c_type_defs.setdefault(c_name, typedef) != typedef:
info = _udt_c_typedef(f"_gbudt_{next(_synthetic_udt_counter)}", datatype.np_type)
c_name = f"_gbudt_{next(_synthetic_udt_counter)}"
info = _udt_c_typedef(c_name, datatype.np_type, is_c_name=True)
c_name, typedef = info
lib.GrB_Type_set_String(datatype._carg, ffi.new("char[]", c_name.encode()), lib.GxB_JIT_C_NAME)
lib.GrB_Type_set_String(
Expand Down Expand Up @@ -276,9 +277,9 @@ def register_anonymous(dtype, name=None):
string, and nested-UDT fields are not.
name : str, optional
A human-readable name for this UDT. Used in error messages and
``repr``. Must be a valid C identifier (and not a C reserved word)
to be usable on the SuiteSparse JIT path; otherwise JIT is skipped
silently and ops fall back to the Numba cfunc path.
``repr``. Any name takes the SuiteSparse JIT path: its C name is
``gbudt_`` and the name, or a synthetic one when that is not a C
identifier or is too long.

Returns
-------
Expand Down
Loading
Loading