Skip to content

Learn built-in UDF return types by typing alone, defer the compiles - #612

Open
eriknw wants to merge 1 commit into
31-numba-cache-builtin-udfsfrom
32-udf-typing-only-types
Open

Learn built-in UDF return types by typing alone, defer the compiles#612
eriknw wants to merge 1 commit into
31-numba-cache-builtin-udfsfrom
32-udf-typing-only-types

Conversation

@eriknw

@eriknw eriknw commented Aug 4, 2026

Copy link
Copy Markdown
Member

First touch of a built-in UDF binop compiled it for all 11 to 13 sample
dtypes, and threw nearly all of that away. The compiles existed only to
read back a return type per dtype for op.types; a program that then used
one dtype paid for the rest.

Get the return types from numba's typing pipeline alone (run_frontend
plus type_inference_stage, no lowering), apply the existing downcast
heuristic over the full result set, and defer each dtype's lowering,
cfunc wrapper, and GrB_BinaryOp_new to the first op[dtype] access via
_build_deferred/_materialize_deferred on OpBase. op.types is unchanged:
dumped for all five ops before and after and diffed, identical.

Numba exposes no public way to ask for a return type without also
compiling for it, so _infer_ret_types_typing_only reaches into
numba.core. We support numba back to 0.57 and this was verified against
0.66.0 only, so the helper returns None on a failed import or any
surprise from the call and _build falls back to the eager loop. Verified
by forcing _HAS_TYPING_ONLY = False: the suite passes with only
test_builtin_udf_types_precede_compilation failing, which is the intended
signal. A numba that moves these internals turns into a red test rather
than a silent return to the old cost.

Only the built-in UDFs take this path, keyed off the same flag that marks
them cacheable. Typing is more permissive than lowering, so a user
function whose dtype types but fails to lower would be wrongly advertised
in .types; register_new and register_anonymous stay eager and keep
validating lowerability per dtype. Semiring construction iterates
_typed_ops directly, so it materializes a deferred multiplier first.

Measured here, fresh process, warm numba cache, medians of 3, load
average 4 to 5 on a shared machine so these are indicative and the
before-numbers are the noisier half. First-touch floordiv: 474 -> 248 ms.
First-touch of all five: 753 -> 289 ms. The deferred per-dtype build does
not vanish, it moves to first use of that dtype.

This does not speed up a program that ends up using every dtype; it stops
programs that use one or two from paying for all of them. It does not
touch UDT auto-lift, which was already compiled on demand.

_finalize_typed_binaryop keeps error_model="numpy" on the cfunc, as the
call site it was extracted from had it and every other cfunc in the
package sets it. No test here distinguishes it: with it removed from the
cfunc, from the njit, or from both, and with the numba cache cleared,
integer floordiv by zero still returns numpy's answer.

Deferral surfaced one more direct consumer of _typed_ops: the typed
commutes_to properties asked whether the PARTNER op had already compiled
this dtype, so in a fresh process floordiv[INT64].commutes_to answered
None until rfloordiv happened to be built, and the answer depended on
access order. The membership test now consults .types, whose keyset is
identical to _typed_ops once materialized, letting getitem drive the
deferred build. test_deferred_commutes_to pins it in a subprocess, where
access order is controlled.


Stack created with GitHub Stacks CLIGive Feedback 💬

@eriknw
eriknw marked this pull request as ready for review August 4, 2026 16:07
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch 2 times, most recently from 7e70d25 to bc9e14e Compare August 5, 2026 00:06
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch 2 times, most recently from c4efaf4 to e17f97f Compare August 5, 2026 17:44
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch from e17f97f to a020a2c Compare August 5, 2026 18:03
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch 2 times, most recently from 6fbd969 to ace2c7a Compare August 6, 2026 07:59
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch from ace2c7a to c489e57 Compare August 6, 2026 15:39
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch from c489e57 to 04b6930 Compare August 6, 2026 15:41
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch 2 times, most recently from a9184d0 to f465e6d Compare August 6, 2026 20:41
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch from f465e6d to 31c66f3 Compare August 7, 2026 02:49
First touch of a built-in UDF binop compiled it for all 11 to 13 sample
dtypes, and threw nearly all of that away. The compiles existed only to
read back a return type per dtype for op.types; a program that then used
one dtype paid for the rest.

Get the return types from numba's typing pipeline alone (run_frontend
plus type_inference_stage, no lowering), apply the existing downcast
heuristic over the full result set, and defer each dtype's lowering,
cfunc wrapper, and GrB_BinaryOp_new to the first op[dtype] access via
_build_deferred/_materialize_deferred on OpBase. op.types is unchanged:
dumped for all five ops before and after and diffed, identical.

Numba exposes no public way to ask for a return type without also
compiling for it, so _infer_ret_types_typing_only reaches into
numba.core. We support numba back to 0.57 and this was verified against
0.66.0 only, so the helper returns None on a failed import or any
surprise from the call and _build falls back to the eager loop. Verified
by forcing _HAS_TYPING_ONLY = False: the suite passes with only
test_builtin_udf_types_precede_compilation failing, which is the intended
signal. A numba that moves these internals turns into a red test rather
than a silent return to the old cost.

Only the built-in UDFs take this path, keyed off the same flag that marks
them cacheable. Typing is more permissive than lowering, so a user
function whose dtype types but fails to lower would be wrongly advertised
in .types; register_new and register_anonymous stay eager and keep
validating lowerability per dtype. Semiring construction iterates
_typed_ops directly, so it materializes a deferred multiplier first.

Measured here, fresh process, warm numba cache, medians of 3, load
average 4 to 5 on a shared machine so these are indicative and the
before-numbers are the noisier half. First-touch floordiv: 474 -> 248 ms.
First-touch of all five: 753 -> 289 ms. The deferred per-dtype build does
not vanish, it moves to first use of that dtype.

This does not speed up a program that ends up using every dtype; it stops
programs that use one or two from paying for all of them. It does not
touch UDT auto-lift, which was already compiled on demand.

_finalize_typed_binaryop keeps error_model="numpy" on the cfunc, as the
call site it was extracted from had it and every other cfunc in the
package sets it. No test here distinguishes it: with it removed from the
cfunc, from the njit, or from both, and with the numba cache cleared,
integer floordiv by zero still returns numpy's answer.

Deferral surfaced one more direct consumer of _typed_ops: the typed
commutes_to properties asked whether the PARTNER op had already compiled
this dtype, so in a fresh process floordiv[INT64].commutes_to answered
None until rfloordiv happened to be built, and the answer depended on
access order. The membership test now consults .types, whose keyset is
identical to _typed_ops once materialized, letting __getitem__ drive the
deferred build. test_deferred_commutes_to pins it in a subprocess, where
access order is controlled.
@eriknw
eriknw force-pushed the 32-udf-typing-only-types branch from 31c66f3 to 3f97f50 Compare August 7, 2026 05:09
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