ENH: Refactor ufunc implementation to avoid intermediate tuples by eendebakpt · Pull Request #31009 · numpy/numpy

eendebakpt · 2026-03-13T22:05:39Z

PR summary

We replace the ufunc_full_args structure which contains tuples with a struct with plain C arrays. The input and output arguments are now allocated in the same way is the signature, operands, operand_DTypes and operation_descrs.

With this we can:

Avoid construction of intermediate tuples in the common case
Use borrowed references on the input arguments (for the common case)
In future PRs: use borrowed references on the operand_DTypes (and perhaps more cases)

On a FT build this improves performance:

x + x: Mean +- std dev: [main] 470 ns +- 1 ns -> [pr] 446 ns +- 4 ns: 1.06x faster
np.abs(x): Mean +- std dev: [main] 424 ns +- 9 ns -> [pr] 401 ns +- 8 ns: 1.06x faster
np.cos(x): Mean +- std dev: [main] 422 ns +- 13 ns -> [pr] 397 ns +- 10 ns: 1.06x faster
np.add.reduce(x): Mean +- std dev: [main] 1.04 us +- 0.02 us -> [pr] 1.01 us +- 0.03 us: 1.03x faster
np.add.accumulate(x): Mean +- std dev: [main] 560 ns +- 11 ns -> [pr] 548 ns +- 6 ns: 1.02x faster

Geometric mean: 1.05x faster

Script

import pyperf

runner = pyperf.Runner()

setup = """
import numpy as np
x = np.array([1., 2.])
y = np.array([-1., 12.])
"""

runner.timeit(name="x + x", stmt="x + y", setup=setup)
runner.timeit(name="np.abs(x)", stmt="np.abs(x)", setup=setup)
runner.timeit(name="np.cos(x)", stmt="np.cos(x)", setup=setup)
runner.timeit(name="np.add.reduce(x)", stmt="np.add.reduce(x)", setup=setup)
runner.timeit(name="np.add.accumulate(x)", stmt="np.add.accumulate(x)", setup=setup)

AI Disclosure

Claude code was used heavily in creating the PR.

seberg

As I had mentioned offline, I like this, but I think we need to make it pretty to be worthwhile.
So a few comments, maybe they make sense, or maybe I'll push some changes if I feel like it soon.

(Part of me wonders also if the in-args should just always be a copy.)

seberg · 2026-03-14T09:08:22Z

+{
+    _xdecref_pyobjects(full_args->out, full_args->nout);
+    full_args->nout = 0;
+}


This is also weird. If we need a helper to clean up full-args, it should just be a single helper that deals with everything.

seberg · 2026-03-14T09:13:37Z

+                    PyObject *item = PyTuple_GET_ITEM(out_obj, i);
+                    full_args_light->out[i] = Py_NewRef(item);
+                }
+                full_args_light->nout = nout;


Why are the output arguments owned suddenly?

The path where outer is true is a bit more complicated than the common path, this is now handled a bit better (but still, depending on outer the handling is slightly different)

seberg · 2026-03-14T09:16:20Z

+    }
+    else {
+        /* Borrowed references directly from vectorcall args */
+        full_args_light.in = args;


Does it matter with the borrowed vs. not. Most of the time, I suspect we'll touch the refcount anyway, so just copying them over always is maybe just as well?

seberg · 2026-03-14T10:15:37Z

-        Py_XDECREF(full_args.in);
-        Py_XDECREF(full_args.out);
+        _xdecref_pyobjects(in_args, full_args_light.nin);
+        ufunc_full_args_light_clear_out(&full_args_light);


Why owning reference(s) here and elsewhere?

seberg · 2026-03-14T10:22:22Z

@@ -4477,24 +4530,27 @@ ufunc_generic_fastcall(PyUFuncObject *ufunc,
        goto fail;


This fail goto looks like it should explode cleanup? Although, I am surprised we wouldn't have a single test excercising it.

seberg · 2026-03-14T10:27:07Z

 PyUFunc_CheckOverride(PyUFuncObject *ufunc, char *method,
-        PyObject *in_args, PyObject *out_args, PyObject *wheremask_obj,
+        PyObject *const *in_args, int nin,
+        PyObject *const *out_args, int nout,


Maybe we should just move the struct definition and pass it? (Maybe even merge with NpyUFuncContext, but may want to split out ufunc there then.)

mhvk

@eendebakpt - Is there a reason this is still a draft? It looks very good to me! As I've always rather disliked the needless creation of input and output tuples, I'm very happy to see that removed here...

mhvk · 2026-03-27T08:31:08Z

    }

-    if (out_args != NULL) {
+    if (out_args != NULL && nout > 0) {


Can nout be zero if out_args is defined? (If this is defensive programming, add a comment, or perhaps just do assert nout > 0 on the next line).

eendebakpt · 2026-03-27T10:56:30Z

@eendebakpt - Is there a reason this is still a draft? It looks very good to me! As I've always rather disliked the needless creation of input and output tuples, I'm very happy to see that removed here...

No real reason it was still draft. With these refactorings it is not always clear where to stop (doing more refactoring can be good, but also makes the diff larger), and I had not made my mind up entirely.

seberg

Asked claude to nitpick (it found a bug) and scrolled through and have two nitpicks on my own (claude some more, but even I thought they were overly nitpicky and/or unrelated).

Part of that "maybe NULL" and "may be owned" is annoying, but it seems fine.

Overall, I think I am getting happy and we could just go and put it in. It's not the prettiest, but the existing tuple isn't either really...

seberg · 2026-04-09T13:52:07Z

-            PyTuple_SET_ITEM(full_args.out, i-nin, tmp);
+            ufunc_output[i-nin] = Py_NewRef(tmp);
        }
+        nout_args = nout;


Claude was right on this one, you have to update it within the loop, so that goto fail does the right thing (or before the goto fail, but doubt that is nicer).

seberg · 2026-04-09T13:52:56Z

-    NPY_ALLOC_WORKSPACE(scratch_objs, void *, UFUNC_STACK_NARGS * 4 + 2, nop * 4 + 2);
+    NPY_ALLOC_WORKSPACE(scratch_objs, void *,
+                        UFUNC_STACK_NARGS * 4 + 2,
+                        nop * 4 + 2 + nop);


Might just write at as nop * 5, since it's just a + anyway?

seberg · 2026-04-09T14:18:10Z

-                goto fail;
+            else {
+                ufunc_output = args_scratch + nin;
+                int n = _set_full_args_out(nout, out_obj, ufunc_output);


I might suggesting moving that ufunc_output = NULL logic into the the function (and maybe rename it).
Just because this is duplicated... I admit it getting set to NULL may be a tad cryptic, but then it can just keep returning 0 and -1 and is concise here?

One thing is that nout_args is really just nout except when ufunc_output == NULL and some places even rely on that.

(The other nit is that the _set_full_args_out name is a bit outdated, but I don't really care.)

Refactor the ufunc call path to pass inputs and outputs as C arrays of PyObject* instead of building the intermediate `ufunc_full_args` in/out tuples. convert_ufunc_arguments, resolve_descriptors, npy_find_array_wrap, and the override machinery now take pointer/length pairs; scratch space comes from a single stack workspace. This removes two tuple allocations from every ufunc call. See numpy#31009. Co-Authored-By: Claude Fable 5 <[email protected]>

seberg · 2026-09-03T13:06:38Z

I guess this keeps getting conflicts :(, want to try to push it over? There is also some talk about whether we can get stable API working well enough and if not vital, this might actually be somewhat nice for that (as it removes some relatively hot-tuple building).

mhvk

@eendebakpt - looked back here again because of a comment. Two totally nitpicky comments if you still need to do a rebase anyway.

A question would be if it might not make more sense to keep in and out together in a single array (even if one has pointers to different parts), but, really, what you have is great and very clear, so the best thing is probably just to get it in!

mhvk · 2026-09-03T16:04:44Z

@@ -3754,22 +3759,23 @@ _set_full_args_out(int nout, PyObject *out_obj, ufunc_full_args *full_args)
            return -1;
        }
        if (tuple_all_none(out_obj)) {


In principle, there might be a bit more speed gain to be had by doing this checking for all arguments being None on the fly in the loop below.

In that way if all are None we then have an extra incref/decref pair for all arguments. For the common case (out not set) it does not matter either way.

Fair enough, though if a tuple is present, it is nearly guaranteed not to be all-None.

Squashed rebase of branch ufunc_tuple_elimination_v2 (numpygh-31009) onto current main. Pre-rebase tip was 64233ad. Conflict resolution: main gained NEP-50 style promotion for exact Python strings (numpygh-32040), which passes the original inputs down to `resolve_descriptors` for `NPY_ARRAY_WAS_PYTHON_STR` operands as well. That is now carried by the borrowed-reference `inputs` array instead of the removed `inputs_tup` tuple. Co-authored-by: Sebastian Berg <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

eendebakpt · 2026-09-03T18:24:35Z

A question would be if it might not make more sense to keep in and out together in a single array (even if one has pointers to different parts), but, really, what you have is great and very clear, so the best thing is probably just to get it in!
The two are allocated together in args_scratch already. Are do you mean something else?

seberg · 2026-09-03T18:33:34Z

Well, let me get this in before it diverges again (and probably annoying someone else with merge conflict :)).

seberg · 2026-09-03T18:34:33Z

The two are allocated together in args_scratch already. Are do you mean something else?

Yeah, I think the question is more whether we could just drag around a single array for in+out. Conceptually, I like that more, but I realize there were a few paths were you needed the distinction.

mhvk · 2026-09-03T18:35:09Z

The two are allocated together in args_scratch already. Are do you mean something else?

Ah, yes, I see now that for the regular ufunc call they indeed are! I think the reduction case was different, that's why this came up. Anyway, I like it how it is!

seberg marked this pull request as draft March 14, 2026 09:02

seberg reviewed Mar 14, 2026

View reviewed changes

eendebakpt force-pushed the ufunc_tuple_elimination_v2 branch from ff25bc2 to a03253a Compare March 18, 2026 21:29

tylerjereddy added the 01 - Enhancement label Mar 18, 2026

eendebakpt mentioned this pull request Mar 26, 2026

ENH: [RFC] Improve performance of numpy operations involving scalars #24345

Open

mhvk reviewed Mar 27, 2026

View reviewed changes

eendebakpt marked this pull request as ready for review March 27, 2026 10:55

eendebakpt changed the title ~~Draft: [ENH] Refactor ufunc implementation to avoid intermediate tuples~~ [ENH] Refactor ufunc implementation to avoid intermediate tuples Mar 27, 2026

seberg reviewed Apr 9, 2026

View reviewed changes

eendebakpt changed the title ~~[ENH] Refactor ufunc implementation to avoid intermediate tuples~~ ENH: Refactor ufunc implementation to avoid intermediate tuples Jul 14, 2026

mhvk approved these changes Sep 3, 2026

View reviewed changes

eendebakpt force-pushed the ufunc_tuple_elimination_v2 branch from 64233ad to 125a2a6 Compare September 3, 2026 17:42

seberg merged commit c2cbbf5 into numpy:main Sep 3, 2026
91 checks passed

eendebakpt deleted the ufunc_tuple_elimination_v2 branch September 3, 2026 19:27

		@@ -4477,24 +4530,27 @@ ufunc_generic_fastcall(PyUFuncObject *ufunc,
		goto fail;

Uh oh!

Conversation

eendebakpt commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR summary

AI Disclosure

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mhvk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eendebakpt commented Mar 27, 2026

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

seberg commented Sep 3, 2026

Uh oh!

mhvk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eendebakpt commented Sep 3, 2026

Uh oh!

seberg commented Sep 3, 2026

Uh oh!

Uh oh!

seberg commented Sep 3, 2026

Uh oh!

mhvk commented Sep 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

eendebakpt commented Mar 13, 2026 •

edited

Loading