ENH: support reduction loops in ufunc.accumulate - #32213

This is a sister PR to #32212 so depending on the merge order, a little bit of buffering is required between them to consider the review comments of the one to the other and rebase whilst fixing merge conflicts.

The branch is now on top of #32212 and I have converted it to draft. The review will continue in that one, and once that one has converged, all the findings from that review will be applied here and once its merged on main the diff will shorten significantly. Now the branch holds 2 PRs in the diff.

Okay this is now rebased on top of main since #32212 has been merged and the feedback from that PR has been considered here. This should be ready for a review pass now @ngoldbaum

I only found one issue during this review pass. Consider this IPython script:

import numpy as np

a = np.ones(200_000, dtype=object)
indices = np.arange(a.size)
out = np.empty_like(a)
rows = a.reshape(-1, 2)
%timeit np.add.reduceat(a, indices, out=out)
%timeit np.add.accumulate(rows, axis=1)

rows = np.ones((100_000, 2))
%timeit np.add.accumulate(rows, axis=1)

On 47f8d7a (just before #32212 was merged), this prints on my ARM macbook:

286 μs ± 1.34 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
872 μs ± 7.27 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
362 μs ± 9.35 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

And on this PR it prints:

556 μs ± 959 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1 ms ± 2.25 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
519 μs ± 609 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Object reductions show the slowdown more clearly for both operations but I'm only able to reproduce it for accumulate using float arrays.

I think to get the performance back, you need to restore the old specialized single-output loops with something like:

if (nout == 1) {
    /* Specialized single-output loop. */
}
else {
    /* Generalized multi-output loop. */
}

That said, I'm not sure how load-bearing these code paths are for performance or if it's worth hyper-optimizing.

I'll do an investigation. I did not see this when we had done reduce only and I also don't understand how the code changes could cause that. I mean for very small arrays just because we're handling an array of size 1 of things instead of just the thing now I would say it's fine but your array is kinda large here.

Oh I see what you did in the benchmark, you are calling reductions along small axes....yeah and the inner loop is called once per row and now each inner loop has more fixed setup overhead I guess which is comparable to the actual loop call. I originally thought you had found some O(N) regression but yeah...it's a small dimension problem. Got scared a bit for my code 🤣

Oh the main issue is not that we are dealing with size 1 arrays now to cover the multi output reductions....It's the refactoring to deduplicate code that introduced more functions and function calls and helpers on the hot path of calling the loop for these small dimension reductions....

It's the refactoring to deduplicate code that introduced more functions and function calls and helpers on the hot path of calling the loop for these small dimension reductions....

Hmm maybe we can give hints to the compiler to encourage inlining

It's the refactoring to deduplicate code that introduced more functions and function calls and helpers on the hot path of calling the loop for these small dimension reductions....

Hmm maybe we can give hints to the compiler to encourage inlining

Yeah I'm trying different things out accros macos and linux to see if I can do anything. There's NPY_INLINE but functions calls are not 100% of the regression. The hammer is as you said the if (nout == 1) {...} else {...} pattern which the branch predictor will get very quickly on the hot path so if my attempts fail, we got that.
I'm basically a bit scared of tryng to make the compiler do something since numpy has to support all these architectures and all these different compilers so even if I find something on my machines, it's hard to make a general statement I guess.

@ngoldbaum see cb5224f for my attempt.

Looks nice! Just a thought: why not specialize two-output reductions too?

Looks nice! Just a thought: why not specialize two-output reductions too?

Hmm that's interesting, do you wanna move this to a follow-up PR perhaps? Worth an investigation for sure but as long as we're not doing any bad regressions here I'd want to get accumulate out of the to-do list.

@ikrommyd I was thinking of something like this: https://gist.github.com/ngoldbaum/37a0d444df0f2b9b48822ea0b476476e

On my Mac using Clang, the patch makes the minmax call timed in this IPython script about 10% faster:

import numpy as np

a = np.ones((10_000, 15))
%timeit np.minmax(a, axis=1)

(103 vs 92.5 microseconds).

I added an nout=2 specialization too @ngoldbaum.
Had to rebase as well due to conflicts. Should be ready for a last check.

Thanks Nathan! Folded the patch in here. Will merge once CI is happy one last time. Cheers!