Het volgen van np.add, tot op de bodem

De premisse is als volgt: np.add(a, b) is waarschijnlijk een van de meest uitgevoerde regels numeriek Python ter wereld. De meeste van ons hebben een mentaal model dat gelijkstaat aan: "het telt de arrays op, in C, en dat gaat snel". Dat model is correct, maar er zit veel mechaniek tussen de Python-aanroep en de loop die de optelling uitvoert, en ik vind het een leuk proces om eens uit elkaar te halen. Vandaag volgen we een enkele aanroep, np.add(a, b) met twee float64-arrays, vanaf het Python-instappunt tot aan de SIMD-kernel, terwijl we de eigenlijke NumPy-broncode lezen.

Alle informatie hieronder is gebaseerd op NumPy 2.5.2. Ik ga ervan uit dat je enigszins comfortabel bent met het lezen van C, maar kennis van de interne werking van NumPy is niet vereist; dat is juist wat we hier gaan ontdekken.

De routekaart

Voordat we beginnen, is hier de "schatkaart", zodat je altijd weet waar we ons bevinden:

np.add(a, b) (Python) $\downarrow$ ufuncgenericfastcall (C: argumenten analyseren) $\downarrow$ arrayufunc override controle (kan afbuigen naar andere bibliotheken) $\downarrow$ Promotie & dispatch (zoek de float64-loop, cache deze) $\downarrow$ Triviale loop óf NpyIter (iteratiestrategie) $\downarrow$ DOUBLEadd (de eigenlijke innerlijke loop, SIMD)

---

np.add is een object

Het eerste dat je moet weten is dat np.add geen normale Python-functie is. Het is een instantie van numpy.ufunc, een in C gedefinieerd type:

>>> 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']

Een ufunc is in essentie een bundel van innerlijke loops. Eén kleine C-functie per ondersteunde type-signatuur, plus metadata over het aantal inputs en outputs. np.add wordt geleverd met 22 daarvan, hoewel de types-lijst alleen de klassieke versies toont; loops die op de moderne manier zijn geregistreerd, leven in een interne mapping die Python nooit ziet.

De loop die we vandaag volgen is dd->d: double, double, naar double. De rest van dit artikel gaat over hoe NumPy van jouw aanroep naar die specifieke entry komt, en wat er gebeurt zodra deze is gevonden.

De duik in C

Wanneer Python np.add(a, b) ziet, wordt het ufunc-object aangeroepen. Het ufunc-type implementeert het vectorcall-protocol, waardoor de aanroep terechtkomt in ufuncgenericvectorcall, die direct doorstuurt naar het echte werkpaard: ufuncgenericfastcall.

Die functie is lang, maar leest als een checklist. Zeer vereenvoudigd ziet het er zo uit:

static PyObject *
ufunc_generic_fastcall(PyUFuncObject *ufunc,
                      PyObject *const *args, Py_ssize_t len_args,
                      PyObject *kwnames, npy_bool outer)
{
    /* ... extraheer inputs, outputs en keyword-argumenten ... */

    /* Controleer op Overrides */
    PyObject *override = NULL;
    errval = PyUFunc_CheckOverride(ufunc, method, full_args.in,
                                   full_args.out, where_obj, args,
                                   len_args, kwnames, &override);

    /* ... als er een override is gevonden, return het resultaat ... */

    /* ... converteer argumenten naar arrays, extraheer hun DTypes ... */
    PyArrayMethodObject *ufuncimpl = promote_and_get_ufuncimpl(ufunc,
                                     operands, signature, operand_DTypes, ...);

    /* Zoek de juiste descriptors voor de operatie */
    if (resolve_descriptors(nop, ufunc, ufuncimpl, ...) < 0) {
        goto fail;
    }

    /* Laatste voorbereidingen en aanroep van de inner-loop */
    errval = PyUFunc_GenericFunctionInternal(ufunc, ufuncimpl,
                                             operation_descrs, operands,
                                             casting, order, wheremask);

    /* ... wrap de outputs en return ze ... */
}

Het proces is: analyseren, controleren op overrides, een loop kiezen, deze uitvoeren en het resultaat wrappen.

Een ontsnappingsluik

Voordat NumPy echt aan het werk gaat, vraagt het de argumenten of zij het liever zelf afhandelen. PyUFuncCheckOverride doorloopt alle inputs en outputs en zoekt naar een niet-standaard arrayufunc methode (het protocol gedefinieerd in NEP 133). Als een argument deze heeft, roept NumPy deze aan en retourneert het resultaat; de rest van de machinery die we hier bespreken wordt dan overgeslagen.

Dit is de hook waardoor np.add(daskarray, cupyarray) werkt met derde partijen. Bibliotheken zoals Dask en CuPy implementeren array_ufunc en nemen de controle over. Dit kunnen we zelf in vier regels nabootsen:

class Diverted:
    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        return f"intercepted {ufunc.__name__}.{method}"

>>> np.add(np.arange(3), Diverted())
'intercepted add.__call__'

Voor onze trace gaan we ervan uit dat beide argumenten gewone ndarrays zijn, dus we gaan verder naar beneden.

Het kiezen van een loop

Nu moet NumPy van "twee float64-arrays" naar "die ene dd->d entry" gaan. Dit is promotion and dispatch, wat zich bevindt in dispatching.cpp. Het proces kan als volgt worden samengevat:

  1. Overwrite eventuele operand_DTypes vanuit de signature.
  2. Controleer of de nieuwe operand_DTypes gecached zijn (zo ja, ga naar stap 4).
  3. Zoek de best passende "loop". Dit gebeurt via multiple dispatching op alle operand_DTypes en loop dtypes. De best passende loop moet beter zijn dan elke andere passende loop. Dit resultaat wordt gecached.
  4. Als de gevonden loop een promoter is: roep de promoter aan. Deze kan de operand_DTypes wijzigen. Ga daarna terug naar stap 2.
  5. De uiteindelijke ArrayMethod is gevonden; de geregistreerde dtypes worden gekopieerd naar de signature.

De "signature" is wat je expliciet vastlegt bij np.add(a, b, dtype=...); in onze aanroep is deze leeg. Een "promoter" is een helper voor gevallen waarin geen loop direct past. Opvallend is dat het alledaagse menggeval, np.add(int32array, float64array), hier geen promoter gebruikt, maar terugvalt op de oude type-resolutie (PyUFunc_AdditionTypeResolver) om de gemeenschappelijke types te kiezen, om vervolgens opnieuw dispatch in te gaan en op dd->d uit te komen.

De cache in stap 2 is cruciaal. De volledige resolutie gebeurt alleen de eerste keer dat een ufunc met een bepaalde combinatie van types wordt aangeroepen. Daarna is het een simpele hash-lookup.

Wat promoteandgetufuncimpl retourneert is een PyArrayMethodObject. Voor float64-optelling is deze ArrayMethod echter een dunne wrapper rond iets veel ouders. Wanneer de eigenlijke loop nodig is, roept getwrappedlegacyufuncloop de PyUFuncDefaultLegacyInnerLoopSelector aan, die simpelweg de types-tabel doorloopt tot hij dd->d vindt en de C-functiepointer ufunc->functions[i] retourneert.

De wrapper die dit aanpast aan de moderne interface is zeer eenvoudig:

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;
}

De aanroep ldata->loop is de "klassieke" ufunc inner-loop interface, PyUFuncGenericFunction, die al decennia ongewijzigd is.

Itereren of niet itereren

We hebben een loop, nu moet hij gevoed worden via PyUFunc_GenericFunctionInternal. Hier wordt een belangrijke beslissing genomen:

/* Controleer of een triviale loop ok is */
int trivial_ok = check_for_trivial_loop(ufuncimpl, op, operation_descrs, casting, buffersize);

if (trivial_ok && context.method->nout == 1) {
    /* Probeer alles af te handelen zonder de (zware) iterator te gebruiken */
    int retval = try_trivial_single_output_loop(&context, op, order, errormask);
    if (retval != -2) {
        return retval;
    }
}
return execute_ufunc_loop(&context, 0, op, order, buffersize, casting, op_flags, errormask);

Het snelle pad komt eerst: als de vormen overeenkomen, er geen broadcasting of casting nodig is, en elke operand 1-D of contigu is, roept trytrivialsingleoutputloop de innerlijke loop één keer aan over de gehele dataset. Er wordt geen iterator aangemaakt.

Alles wat daar niet aan voldoet, gaat via executeufuncloop en NpyIter, de algemene array-iterator van NumPy. Deze regelt broadcasting, buffering van niet-uitgelijnde of cast-operanden, en overlap-detectie tussen inputs en outputs. De uiteindelijke uitvoering is vervolgens simpel:

int res;
do {
    res = strided_loop(context, dataptr, countptr, strides, auxdata);
} while (res == 0 && iternext(iter));

De iterator geeft de loop per keer een contigu-achtig blok en verplaatst de pointers.

Twee belangrijke details:

  1. Tenzij de loop de Python API nodig heeft, laat NumPy de GIL (Global Interpreter Lock) los rondom het geheel.
  2. Floating-point statusvlaggen worden voor de loop gewist en daarna gecontroleerd; dit is waar de bekende RuntimeWarning: overflow encountered in add vandaan komt.

De loop zelf

Laten we kijken naar het prestatieverschil tussen contigue arrays en arrays met strides (sprongen):

n = 10_000_000
a, b, out = (np.random.rand(2 * n) for _ in range(3))
contig_a, contig_b, contig_out = (x[:n].copy() for x in (a, b, out))
strided_a, strided_b, strided_out = (x[::2] for x in (a, b, out))

np.add(contig_a, contig_b, out=contig_out)     # 3.30 ms per aanroep
np.add(strided_a, strided_b, out=strided_out)  # 8.58 ms per aanroep

Op mijn machine is dit een factor 2.6 verschil voor hetzelfde aantal optellingen. Aangezien beide aanroepen via het triviale pad lopen, moet het verschil in DOUBLE_add zelf zitten.

Er is geen bestand in de repository met een functie genaamd DOUBLEadd. De element-wise loops leven in template-bestanden zoals loopsarithmfp.dispatch.c.src. Tijdens het build-proces breidt een script (convtemplate.py) deze templates uit naar functies als FLOATadd, DOUBLEadd, etc.

Binnenin de functie is er een cascade van specialisaties. Er zijn SIMD-versnelde versies voor veelvoorkomende stride-patronen (beide inputs contigu, of één input een scalar), geschreven met "universal intrinsics" (een portable SIMD-abstractie). Wanneer geen specialisatie past, valt het terug op de basis-loop:

loop_scalar:
for (; len > 0; --len, src0 += ssrc0, src1 += ssrc1, dst += sdst) {
    const @type@ a = *((@type@*)src0);
    const @type@ b = *((@type@*)src1);
    *((@type@*)dst) = a @OP@ b;
}

Hier ligt het prestatieverschil: de SIMD-specialisaties vereisen dat operanden contigu of scalar zijn. De strided aanroep faalt bij elke check en belandt in deze scalar-branch, terwijl de contigue aanroep de gevectoriseerde versie krijgt.

Waar de loops vandaan komen

Hoe komen deze gegenereerde functies in de ufunc->functions tabel terecht? Het antwoord is: meer codegeneratie. De volledige inventaris van built-in ufuncs is gedefinieerd in een Python-dictionary in generate_umath.py. De entry voor add bepaalt dat:

  • Booleans logical_or hergebruiken.
  • Float en complex types (fdFD) hun loops krijgen van loopsarithmfp.
  • Integers auto-vectorized loops krijgen.
  • Datetimes handgeschreven special cases hebben.
  • Object-arrays simpelweg Pythons PyNumber_Add aanroepen.

Tijdens het builden wordt deze dictionary omgezet in __umath_generated.c. Een laatste truc: de dispatch= markers betekenen dat de loop meerdere keren wordt gecompileerd, één keer per CPU-feature set (bijv. AVX2, NEON). Bij het importeren van NumPy bepaalt een macro via runtime CPU-checks welke variant in de tabel wordt geplaatst.

Je kunt zelfs aan NumPy vragen welke variant je hebt gekregen:

>>> from numpy.lib import introspect
>>> introspect.opt_func_info(func_name="add", signature="float64")
{'add': {'ddd': {'current': 'baseline(NEON NEON_FP16 NEON_VFPV4 ASIMD)',
                 'available': 'baseline(NEON NEON_FP16 NEON_VFPV4 ASIMD)'}}}

Slotwoord

Laten we de route recapituleren: np.add is een object met een tabel van 22 loops. Een aanroep analyseert argumenten, biedt de kans op overrides, resolvert input-types naar één tabel-entry (gecached), kiest een iteratiestrategie (triviaal of NpyIter), laat de GIL los, en voert een CPU-specifieke, template-gegenereerde C-functie uit die in zijn meest algemene vorm een simpele for-loop is.

Wat me het meeste inspireerde bij dit onderzoek is de zichtbare "geologie" van de code. De innerlijke loop-signatuur en de types-tabel zijn decennia oud en vormen nog steeds de kern. De ArrayMethod-laag, het override-protocol en de SIMD-dispatch zijn daar bovenop gelegd zonder de geschiedenis te wissen. Het bestand dat de oude interface wrapt heet letterlijk legacyarraymethod.c, maar het is geen verouderde machine; het is de machine.

Zelfs deze oude lagen worden nog onderhouden. Ik ontdekte dat de legacy loop-lookup onlangs is gerefactord op de main branch om NumPy voor te bereiden op free-threaded Python, waarbij global state een probleem wordt. Iemand heeft een van de meest uitgevoerde codepaden in numeriek Python, decennia later, herzien omdat de wereld is veranderd. Dat is schoonheid.

---

Voetnoten

  1. De interne werking verandert tussen versies; sommige beschreven mechanismen zijn in de main branch van NumPy alweer aangepast.
  2. De ufunc leeft in de C-extensie module numpy.core.multiarray_umath. De naam is een historisch artefact uit de tijd dat de array-machinery (multiarray) en de ufunc-machinery (umath) aparte modules waren.
  3. Niet te verwarren met arrayfunction (NEP 18), wat het analoge protocol is voor non-ufunc functies zoals np.concatenate. Ufuncs raadplegen alleen arrayufunc.
  4. De ArrayMethod-laag is er zodat nieuwe dtypes (zoals de StringDType in NumPy 2.0) loops kunnen registreren zonder de klassieke tabel te gebruiken.
  5. De overstap naar free-threaded Python vereist dat initiële globale status wordt verplaatst naar initialisatietijd, om concurrency-problemen te voorkomen.