Learn built-in UDF return types by typing alone, defer the compiles - #612
Open
eriknw wants to merge 1 commit into
Open
Learn built-in UDF return types by typing alone, defer the compiles#612eriknw wants to merge 1 commit into
eriknw wants to merge 1 commit into
Conversation
eriknw
marked this pull request as ready for review
August 4, 2026 16:07
eriknw
force-pushed
the
32-udf-typing-only-types
branch
2 times, most recently
from
August 5, 2026 00:06
7e70d25 to
bc9e14e
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
2 times, most recently
from
August 5, 2026 17:44
c4efaf4 to
e17f97f
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
from
August 5, 2026 18:03
e17f97f to
a020a2c
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
2 times, most recently
from
August 6, 2026 07:59
6fbd969 to
ace2c7a
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
from
August 6, 2026 15:39
ace2c7a to
c489e57
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
from
August 6, 2026 15:41
c489e57 to
04b6930
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
2 times, most recently
from
August 6, 2026 20:41
a9184d0 to
f465e6d
Compare
eriknw
force-pushed
the
32-udf-typing-only-types
branch
from
August 7, 2026 02:49
f465e6d to
31c66f3
Compare
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
force-pushed
the
32-udf-typing-only-types
branch
from
August 7, 2026 05:09
31c66f3 to
3f97f50
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 CLI • Give Feedback 💬