diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index 670cbc6e3..2f5bdcc01 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -188,3 +188,27 @@ time: # numba.njit has run. The compile happens on first lookup (``unary.heavy_op[int]``). + +Compilation caching +------------------- + +A handful of built-in operators are themselves UDFs (``binary.floordiv``, +``rfloordiv``, ``absfirst``, ``abssecond``, and ``rpow``). Because these are +plain module-level functions, python-graphblas compiles them with Numba's +``cache=True``, so the Numba half of the compile is paid once per machine +instead of once per Python process. Numba writes the cache (``.nbi`` index and +``.nbc`` data files) into the ``__pycache__`` directory next to the operator +module. A warm cache shaves roughly 20% off a single op's first-touch build and +about 45% when several are used; the rest is Numba/LLVM startup and per-process +object linking, which caching cannot remove. + +If that ``__pycache__`` directory is not writable (a read-only install, for +example), Numba silently falls back to its per-user cache directory +(``~/.cache/numba`` on Linux, ``~/Library/Caches/numba`` on macOS; set +``NUMBA_CACHE_DIR`` to override). No warning is emitted for the fallback. + +Operators you create with ``register_new`` or ``register_anonymous`` are never +passed ``cache=True``: Numba can only cache a function with a stable on-disk +source, which a lambda or an interactively defined function does not have. A +registered operator's compile is paid once per process, and ``lazy=True`` is +the tool for deferring it. diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index bdcfccc71..01c790acd 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -519,7 +519,7 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: @@ -532,7 +532,12 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): # setting it only on the wrapper leaves ``x // 0`` raising # ZeroDivisionError inside a cfunc, where Numba prints the traceback # and returns, handing GraphBLAS an element it never wrote. - binary_udf = numba.njit(func, error_model="numpy") + # + # ``cache`` is only True for the module-level built-in UDFs. Numba can + # only key a cache entry on a stable on-disk source, which lambdas and + # interactively defined functions do not have, so ops registered by + # users stay uncached. + binary_udf = numba.njit(func, error_model="numpy", cache=cache) new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) return_types = {} nt = numba.types @@ -730,7 +735,9 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals return cls._build(name, func, anonymous=True, is_udt=is_udt) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, _cache=False + ): """Register a new BinaryOp and save it to ``graphblas.binary`` namespace. Parameters @@ -797,13 +804,14 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal "func": func, "parameterized": parameterized, "is_udt": is_udt, + "_cache": _cache, }, ) elif parameterized: binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func, is_udt=is_udt) + binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache) setattr(module, funcname, binary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -856,13 +864,20 @@ def _initialize(cls): if _supports_udfs: # Add floordiv # cdiv truncates towards 0, while floordiv truncates towards -inf - BinaryOp.register_new("floordiv", _floordiv, lazy=True) # cast to integer - BinaryOp.register_new("rfloordiv", _rfloordiv, lazy=True) # cast to integer + # cache=True persists the numba compilation of these built-in UDFs to + # disk so it is paid once per machine, not once per process. It shaves + # ~20% off a single op's first-touch build and ~45% when several are + # used (the per-process object-link and LLVM-init cost is not cached). + # Only module-level built-ins get cache=True (see _build). + BinaryOp.register_new("floordiv", _floordiv, lazy=True, _cache=True) # cast to integer + BinaryOp.register_new( + "rfloordiv", _rfloordiv, lazy=True, _cache=True + ) # cast to integer # For aggregators - BinaryOp.register_new("absfirst", _absfirst, lazy=True) - BinaryOp.register_new("abssecond", _abssecond, lazy=True) - BinaryOp.register_new("rpow", _rpow, lazy=True) + BinaryOp.register_new("absfirst", _absfirst, lazy=True, _cache=True) + BinaryOp.register_new("abssecond", _abssecond, lazy=True, _cache=True) + BinaryOp.register_new("rpow", _rpow, lazy=True, _cache=True) # For algorithms binary._delayed["binom"] = (_register_binom, {}) # Lazy with custom creation diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 71543251d..4361b8dc2 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3808,3 +3808,22 @@ def test_initialize_does_not_build_lazy_udfs(): still_lazy = lines.get("still lazy", "").split() assert still_lazy == ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"], report + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udfs_are_disk_cached(): + # The built-in UDF binops are module-level functions, so numba can persist + # their compilation across processes. Anything a user registers cannot be: + # numba keys a cache entry on a stable on-disk source, which a lambda or an + # interactively defined function does not have. + from numba.core.caching import NullCache + + for name in ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"]: + numba_func = getattr(binary, name)._numba_func + assert not isinstance(numba_func._cache, NullCache), name + + def _uncached_probe(x, y): + return x + y + + user_op = BinaryOp.register_anonymous(_uncached_probe) + assert isinstance(user_op._numba_func._cache, NullCache)