{{ message }}
ENH: fast path for full contiguous reductions - #31274
Merged
Merged
Conversation
Convert the function to use a single ``cleanup:`` exit point that decrefs the descriptors and returns ``result``. No behaviour change. This makes the function easy to extend with extra paths (e.g. fast paths for special inputs) without duplicating the cleanup code.
For full reductions (``axis=None``) on contiguous, aligned, non-object arrays where the matching dtype matches the input dtype and the operation has an identity, call the strided reduce loop directly on the input buffer. This bypasses ``NpyIter`` and ``PyUFunc_ReduceWrapper`` entirely. Speeds up ``np.sum``, ``np.prod``, ``np.any``, ``np.all`` and similar reductions on small/medium arrays by ~1.3x; ``np.array_equal`` and ``np.allclose`` benefit indirectly because they call ``.all()`` internally on a comparison's contiguous boolean output. The reduction loop writes into the result array's data buffer directly, so no extra ``memcpy`` is needed. Any unmet condition or runtime failure within the fast path falls through to the existing slow path.
ngoldbaum
added
56 - Needs Release Note.
Member
|
Do the existing benchmarks in the benchmark suite cover the cases in the benchmark you included in the description? This needs a release note to be mergeable. @mhvk would you mind taking a look at the C implementation for the new fast path? Please let me know if you'd prefer I don't ping you like this on AI-generated PRs. |
Contributor
Author
Yes, they are covered by Details
I will add one ( maybe it can be merged later with some of the other performance related PRs). |
mhvk
reviewed
Apr 22, 2026
mhvk
left a comment
mhvk
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the ping. This looks good and the performance increase is quite impressive. My main comment is to try to make it more like we do in regular calls, using a separate try_fast function, for readability.
| && PyArray_DESCR(arr) == descrs[1] | ||
| && ufuncimpl->get_reduction_initial != NULL | ||
| && !PyDataType_REFCHK(descrs[0])) { | ||
| npy_intp count = PyArray_SIZE(arr); |
Contributor
There was a problem hiding this comment.
How about making this a separate function, so the flow is less interrupted? It would then be analogous to try_trivial_single_output_loop in this same file, and can similarly have a nice comment above it. (Of course, the compiler will just put it inline, so this is really for human consumption only.) Maybe only keep in the ->get_reduction_initial part here, as that is part of the if statement at some level.
Move the inline fast path inside ``PyUFunc_Reduce`` into a static helper ``try_reduce_contiguous`` that returns 1/0/-1 (success / not applicable / hard error). The caller now reduces to a single conditional call. The condition ``PyArray_ISCARRAY_RO(arr)`` is widened to also accept ``PyArray_ISFARRAY_RO(arr)``: for full reductions the loop walks all elements in memory order, which is correct for any contiguous layout. For multi-axis reductions, this makes order-of-iteration vary (it matches storage order rather than NpyIter's order), so an explicit ``NPY_METH_IS_REORDERABLE`` check is added. This also closes a small latent gap where the previous fast path silently succeeded on non-reorderable ops with multi-D C-contiguous input, while the slow path raises ``ValueError``.
ngoldbaum
reviewed
Apr 23, 2026
Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>
seberg
removed
the
56 - Needs Release Note.
seberg
approved these changes
Apr 28, 2026
seberg
left a comment
seberg
left a comment
Member
There was a problem hiding this comment.
Cool, LGTM. One small comment to refine, but doesn't even matter, I was first surprised it's just 1.3x, but any/all being 1.9x makes sense then ;).
It might make sense to also deal with the no-identity version, but probably not in this PR.
| if (!has_initial) { | ||
| /* No identity available -- fall back to the slow path. */ | ||
| Py_DECREF(result); | ||
| return 0; |
Member
There was a problem hiding this comment.
Too bad we create/delete an array here, but OK. (I was wondering how annoying it is to just handle this, but maybe not, it does require copying the first element from the array over, which might be a bit ugly.)
Contributor
Author
There was a problem hiding this comment.
I added cases without an identify to the PR. Added code is limited, and it avoids creating an array for nothing.
| /* Allocate the 0-d result first so the loop can write into it. */ | ||
| Py_INCREF(descrs[0]); | ||
| PyArrayObject *result = (PyArrayObject *)PyArray_NewFromDescr( | ||
| &PyArray_Type, descrs[0], 0, NULL, NULL, NULL, 0, NULL); |
Member
There was a problem hiding this comment.
This path is so close to being able to just return a scalar directly :) (if it worked out with the array-wrap)... But let's not think about it here.
| PyArrayMethodObject *ufuncimpl = context->method; | ||
| if (!(out == NULL && wheremask == NULL && initial == NULL && keepdims == 0 | ||
| && naxes == ndim | ||
| && (PyArray_ISCARRAY_RO(arr) || PyArray_ISFARRAY_RO(arr)) |
Member
There was a problem hiding this comment.
We could actually use PyArray_TRIVIALLY_ITERABLE I think. Then use the actual strides below (because that lets 1-D arrays always pass).
(Doesn't include ALIGNED check though.)
seberg
reviewed
Apr 28, 2026
| res = -1; | ||
| } | ||
| if (res == 0 && needs_fperr) { | ||
| res = _check_ufunc_fperr(errormask, ufunc_name); |
Member
There was a problem hiding this comment.
Suggested change
| res = _check_ufunc_fperr(errormask, ufunc_name); | |
| res = _check_ufunc_fperr(errormask, "reduce"); |
Weird that the new changes now hit path that actually tests this...
Other changes look good since you have the REFCHK in there, I think. We could possibly broaden it up eventually, but really no need here.
EDIT: Ah, because these tests are functions without an identity.
EDIT2: Commit, since I think it might be ready then, but feel free to force push it away!
Member
MaanasArora
pushed a commit
to MaanasArora/numpy
that referenced
this pull request
May 7, 2026
Adds a fast path in PyUFunc_Reduce for the common case of a full reduction (axis=None) over a trivially-iterable, aligned, non-object/reference input where the matching dtype equals the input dtype. Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com> Co-authored-by: Sebastian Berg <sebastianb@nvidia.com>
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
4 participants
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.
PR summary
Adds a fast path in
PyUFunc_Reducefor the common case of a full reduction (axis=None) over a contiguous, aligned, non-object input where the matching dtype equals the input dtype and the operation has an identity value (e.g.np.sum,np.prod,np.any,np.all).(EDIT: Seberg, later broadened to include any 1-D arrays and no identity value.)
Benchmark results
np.sum(array_4)np.sum(array_20)np.sum(array_100)np.prod(array_20)np.sum(10x10)np.sum(2x2)np.any(bool_20)np.all(bool_20)np.array_equal(a,b)np.allclose(a,b)np.max(array_20)np.sum(10x10,axis=0)np.sin(array_20)np.array_equalandnp.allclosebenefit indirectly: both internally call.all()on a contiguous boolean array.Benchmark script
AI Disclosure
Claude code was used to identify performance bottlenecks for the reductions. Improvement of the general case is also possible, but requires many changes for a much smaller gain. The fast path was written by Claude and manually refined.