Tracing np.add, all the way down

The notes for this blog post have been sitting in my drafts folder for half a year now. I’ve done a little work on NumPy itself in the past year. Nothing notable, but enough to have to find my way around the source. That gave me the idea to write this, but then, as other obligations overshadowed my NumPy contributions, it just started rotting quietly. One or two NumPy releases later I finally picked it up again, retraced my steps, and here we are. Here’s the premise: np.add(a, b) might well be among the most executed lines of numerical Python in the world, and most of us have a working mental model that’s equivalent to “it adds the arrays, in C, quickly”. That model is correct, but there’s a lot of machinery between the Python call and the loop that does the adding, and I think it’s a fun machine to take apart. So today we’ll trace a single call, np.add(a, b) with two float64 arrays, from the Python entry point all the way down to the SIMD kernel, reading the actual NumPy source as we go. We’ll learn a lot, I hope! Everything below is pinned to NumPy 2.5.2, the current release as I write this, and all links point into that tag. The internals move around between versions1, so if you’re spelunking along at home, check out the matching tag. I’ll assume you’re at least somewhat comfortable reading C, but no NumPy internals knowledge is required, that’s what we’re here for. The map Before we dive in, here’s the treasure map, so you always know where we are: np.add(a, b) (Python) │ ▼ ufunc_generic_fastcall (C: parse arguments) │ ▼ __array_ufunc__ override check (may divert to other libraries) │ ▼ promotion & dispatch (find the float64 loop, cache it) │ ▼ trivial loop or NpyIter (iteration strategy) │ ▼ DOUBLE_add (the actual inner loop, SIMD) Each of these is a section below. Let’s start at the top. np.add is an object The first thing to know is that np.add is not a normal Python function. It’s an instance of numpy.ufunc, a C-defined type2: >>> type(np.add) <class 'numpy.ufunc'> >>> np.add.nin, np.add.nout (2, 1) >>> len(np.add.types) 22 >>> np.add.types[11:14] ['ee->e', 'ff->f', 'dd->d'] A ufunc is, at its core, a bundle of inner loops. One small C function per supported type signature, plus metadata about how many inputs and outputs there are. np.add ships 22 of them, though I should say that types only lists the classic ones: loops registered the modern way (more on that distinction later) live in an internal mapping on the ufunc that Python never sees. The one we’re chasing today is dd->d: double, double, to double. The whole rest of this post is about how NumPy gets from your call to that one entry, and what happens once it’s found. Other paths may vary, it’s a big piece of kit! Let’s dive into the cave. Into the C When Python sees np.add(a, b), it calls the ufunc object. The ufunc type implements the vectorcall protocol, so the call lands in ufunc_generic_vectorcall, which immediately forwards to the real workhorse, ufunc_generic_fastcall. That function is long, but it reads like a checklist, and it is the skeleton of the whole operation.

The signature is what you fix explicitly when you call np.add(a, b, dtype=...); in our call it’s empty. A “promoter” is a registered helper that handles cases where no loop matches directly by rewriting the requested types and letting dispatch run again. Confusingly, the everyday mixed case, np.add(int32_array, float64_array), doesn’t even use one: when dispatch comes up empty there, it falls back to the ufunc’s old type resolution machinery (PyUFunc_AdditionTypeResolver, in our case) to pick the common types, and then re-enters dispatch with those, landing on dd->d. And the cache in step 2 matters a lot! The full resolution only happens the first time you call a ufunc with a given combination of types. For an ordinary cacheable case like ours, every later call with the same types is a single hash lookup on the DType classes. It’s pure machinery, but it’s useful. To understand what we now have, we have to engage in some archaeology and word-slinging. What promote_and_get_ufuncimpl returns is a PyArrayMethodObject, the modern (post-NEP 43) representation of “one concrete implementation of an operation for concrete DTypes”. For float64 addition, though, the ArrayMethod is a thin wrapper around something much older. When the actual loop is needed, get_wrapped_legacy_ufunc_loop calls PyUFunc_DefaultLegacyInnerLoopSelector, which does exactly what one would write the first time around! It walks the ufunc’s types table, entry by entry, until it finds dd->d, and returns ufunc->functions[i], a plain C function pointer. The wrapper that adapts it to the modern interface is adorable: static int generic_wrapped_legacy_loop(PyArrayMethod_Context *NPY_UNUSED(context), char *const *data, const npy_intp *dimensions, const npy_intp *strides, NpyAuxData *auxdata) { legacy_array_method_auxdata *ldata = (legacy_array_method_auxdata *)auxdata; ldata->loop((char **)data, dimensions, strides, ldata->user_data); if (ldata->pyerr_check && PyErr_Occurred()) { return -1; } return 0; } That ldata->loop call is the “classic” ufunc inner-loop interface, PyUFuncGenericFunction, unchanged for decades. It’s an array of data pointers, an element count, a stride per operand, and an opaque payload. Every builtin numeric loop in NumPy still has this shape. Everything above it exists to line up memory so that calling it is correct4. We are deep in the cave now. It’s almost entirely dark. To iterate or not to iterate We have a loop, now we need to feed it.

No iterator is constructed at all. For the everyday np.add(a, b) of two well-behaved same-shape arrays, this is the path you’re on. Everything else goes through execute_ufunc_loop and NpyIter, NumPy’s general array iterator.