追踪 np.add：一路到底

# 追踪 np.add，一路到底 来源：https://blog.veitheller.de/numpy.html 这篇博客的笔记在我的草稿箱里已经躺了半年了。过去一年里我对NumPy本身做了一点工作。虽然不算什么成就，但足以让我去熟悉它的源代码。这让我萌生了写这篇文章的想法，但后来，随着其他事务压过了我对NumPy的贡献，它就悄悄地开始腐烂了。又过了一两个NumPy版本后，我终于重新拾起它，重溯我的脚步，于是就有了现在这篇文章。 前提是：`np.add(a, b)`很可能是世界上执行次数最多的数值Python代码行之一，我们大多数人对它都有一个等效于“它在C语言里快速地将数组相加”的心智模型。这个模型是正确的，但在Python调用和实际执行加法的循环之间，存在大量的机制，我认为拆解这个机器会很有趣。所以今天我们将追踪单次调用`np.add(a, b)`——使用两个`float64`数组——从Python入口点一直到底层的SIMD内核，同时阅读实际的NumPy源码。希望我们能学到很多！ 以下内容均基于NumPy 2.5.2版本（https://github.com/numpy/numpy/tree/v2.5.2），这是我撰写本文时的当前版本，所有链接都指向该标签。内部实现在不同版本间会有所变动1 (https://blog.veitheller.de/numpy.html#1)，因此如果你在家跟着一起探索，请检出对应的标签。我假设你至少对阅读C代码比较熟悉，但不需要任何NumPy内部知识，这正是我们要来学习的地方。 ## 地图 在深入之前，先看一下这张寻宝地图，这样你总能知道我们在哪里： ``` np.add(a, b) (Python) │ ▼ ufunc_generic_fastcall (C：解析参数) │ ▼ __array_ufunc__ 覆盖检查 (可能会转交给其他库) │ ▼ 类型提升与分发 (找到float64循环，缓存它) │ ▼ 简单循环或 NpyIter (迭代策略) │ ▼ DOUBLE_add (实际的内部循环，SIMD) ``` 以下每一节对应其中一步。让我们从顶部开始。 ## np.add 是一个对象 首先要了解的是，`np.add`并不是一个普通的Python函数。它是`numpy.ufunc`类型的一个实例，这是一种C语言定义的类型2 (https://blog.veitheller.de/numpy.html#2)： ```python >>> type(np.add) >>> np.add.nin, np.add.nout (2, 1) >>> len(np.add.types) 22 >>> np.add.types[11:14] ['ee->e', 'ff->f', 'dd->d'] ``` ufunc的核心是一系列内部循环的捆绑包。每个支持的类型签名对应一个C小函数，外加关于输入和输出数量的元数据。`np.add`提供了22个这样的循环，不过我应该说明，`types`只列出了传统的循环：以现代方式注册的循环（稍后会详细讨论这种区别）存在于ufunc的一个内部映射中，Python永远看不到。 我们今天要追踪的是`dd->d`：double, double, to double。本文其余的所有内容都是关于NumPy如何从你的调用到达这个条目，以及找到它之后发生了什么。其他路径可能有所不同，这是一个庞大的工具包！ 让我们潜入洞穴。 ## 进入C语言当Python看到`np.add(a, b)`时，它调用ufunc对象。`ufunc`类型实现了向量调用协议，因此调用会落到`ufunc_generic_vectorcall` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/ufunc_object.c#L4905)，它立即转发给真正的工作马——`ufunc_generic_fastcall` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/ufunc_object.c#L4561)。这个函数很长，但它读起来像一个清单，并且它*就是*整个操作的骨架。高度简化后： ```c static PyObject * ufunc_generic_fastcall(PyUFuncObject *ufunc, PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames, npy_bool outer) { /* ... 提取输入、输出和关键字参数 ... */ /* 我们现在拥有了检查覆盖所需的所有信息 */ PyObject *override = NULL; errval = PyUFunc_CheckOverride(ufunc, method, full_args.in, full_args.out, where_obj, args, len_args, kwnames, &override); /* ... 如果找到覆盖，则返回其结果 ... */ /* ... 将参数转换为数组，提取它们的DType ... */ PyArrayMethodObject *ufuncimpl = promote_and_get_ufuncimpl(ufunc, operands, signature, operand_DTypes, ...); /* 为操作找到正确的描述符 */ if (resolve_descriptors(nop, ufunc, ufuncimpl, ...) < 0) { goto fail; } /* * 做最后的准备工作并调用内部循环。 */ errval = PyUFunc_GenericFunctionInternal(ufunc, ufuncimpl, operation_descrs, operands, casting, order, wheremask); /* ... 包装输出并返回 ... */ } ``` 解析、检查覆盖、选择循环、运行它、包装结果。看来我们有了计划！现在看看有趣的部分。我们洞穴里的光线越来越暗了。 ## 一个逃生通道 在NumPy承诺做任何工作之前，它会询问参数是否希望自己来做（这总是一个好的操作模式）。`PyUFunc_CheckOverride` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/override.c#L206)遍历所有输入和输出，寻找非默认的`__array_ufunc__`方法，这是在NEP 13 (https://numpy.org/neps/nep-0013-ufunc-overrides.html)3 (https://blog.veitheller.de/numpy.html#3)中定义的协议。如果任何参数有这个方法，NumPy会调用它并返回它产生的任何内容，而下面我们要讨论的机制则一个都不会运行。这就是让`np.add(dask_array, cupy_array)`能够与第三方库配合工作的钩子！像Dask和CuPy这样的库实现了`__array_ufunc__`并接管了控制权。我们自己也可以用四行代码演奏同样的曲调： ```python 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__' ``` 我们的对象什么也不做，只是赢得了关于谁来计算的争论。对于我们的追踪，我们假设两个参数都是普通的`ndarray`，所以没有覆盖，我们继续向下垂降。 ## 选择循环 接下来，NumPy必须从“两个`float64`数组”到达“那个`dd->d`条目”。这就是提升和分发，它位于`dispatching.cpp` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/dispatching.cpp)，其头部注释是我在任何地方找到的关于该过程的最佳文档，所以让我直接引用它，连同拼写错误一起： ``` 分发和提升的过程可以总结为以下步骤： 1. 用`signature`覆盖任何`operand_DTypes`。 2. 检查新的`operand_DTypes`是否已缓存（如果已缓存，转到第4步）。 3. 查找最佳匹配的"循环"。这是通过对所有`operand_DTypes`和循环`dtypes`进行多重分发来完成的。匹配的循环必须是其DType是`operand_DTypes`（已定义）的超类的循环。最佳匹配的循环必须优于任何其他匹配的循环。此结果被缓存。 4. 如果找到的循环是promoter：我们调用promoter。它可以修改当前的`operand_DTypes`。然后返回第2步。 5. 找到最终的`ArrayMethod`，其注册的`dtypes`被复制到`signature`中，以便ufunc循环可以使用。 ``` 需要做几个翻译。`signature`是你在调用`np.add(a, b, dtype=...)`时显式固定的；在我们的调用中它是空的。“promoter”是一个注册的帮助器，用于处理没有直接循环匹配的情况，通过重写请求的类型并让分发重新运行来处理。令人困惑的是，日常的混合情况，如`np.add(int32_array, float64_array)`，甚至不使用promoter：当分发在那里找不到结果时，它会回退到ufunc的*旧*类型解析机制（在我们的情况下是`PyUFunc_AdditionTypeResolver`）来选择公共类型，然后带着这些类型重新进入分发，最终落在`dd->d`上。第2步中的缓存非常重要！完整的解析只在你第一次用给定的类型组合调用ufunc时发生。对于我们这样普通的可缓存情况，后续每次相同类型的调用都只是对DType类进行一次哈希查找。这纯粹是机制，但很有用。 为了理解我们现在拥有什么，我们需要进行一些考古和文字解读。`promote_and_get_ufuncimpl` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/dispatching.cpp#L1047)返回的是一个`PyArrayMethodObject`，这是“针对具体DType的一个操作的具体实现”的现代（NEP 43之后(https://numpy.org/neps/nep-0043-extensible-ufuncs.html)）表示。然而，对于`float64`加法，ArrayMethod只是一个更古老事物的薄包装器。当实际需要循环时，`get_wrapped_legacy_ufunc_loop` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/legacy_array_method.c#L203)调用`PyUFunc_DefaultLegacyInnerLoopSelector` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/ufunc_type_resolution.c#L1531)，它做的正是人们一开始会写的东西！它逐条遍历ufunc的`types`表，直到找到`dd->d`，然后返回`ufunc->functions[i]`，一个普通的C函数指针。将其适配到现代接口的包装器很可爱： ```c 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; } ``` 那个`ldata->loop`调用就是“经典”的ufunc内部循环接口`PyUFuncGenericFunction`，几十年来都没有改变。它是一个数据指针数组、一个元素计数、每个操作数的步长和一个不透明的有效载荷。NumPy中的每个内置数值循环仍然具有这种形状。它之上的所有存在都是为了排列内存，使得调用它是正确的4 (https://blog.veitheller.de/numpy.html#4)。我们现在深入洞穴。几乎完全黑暗了。 ## 是否迭代我们有了一个循环，现在需要喂它。这就是`PyUFunc_GenericFunctionInternal` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/ufunc_object.c#L2192)，它做出了一个有趣的决策： ```c /* * 这检查简单循环是否可行，如果该操作有助于提高性能， * 则对标量和一维操作数进行复制。 */ int trivial_ok = check_for_trivial_loop(ufuncimpl, op, operation_descrs, casting, buffersize); /* ... */ if (trivial_ok && context.method->nout == 1) { /* 尝试不使用（沉重的）迭代器来处理所有情况 */ 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); ``` 快速路径在前面：如果形状匹配，不需要广播或类型转换，并且每个操作数都是1维或连续的，`try_trivial_single_output_loop` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/ufunc_object.c#L869)会对整个数据调用一次内部循环。根本不构造迭代器。对于日常的`np.add(a, b)`——两个行为良好、形状相同的数组——这就是你走的路径。其他所有情况都通过`execute_ufunc_loop` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/ufunc_object.c#L1061)和`NpyIter`（NumPy的通用数组迭代器）处理。构造标志是它为内部循环处理的所有事情的一个紧凑总结： ```c npy_uint32 iter_flags = ufunc->iter_flags | NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK | NPY_ITER_ZEROSIZE_OK | NPY_ITER_BUFFERED | NPY_ITER_GROWINNER | NPY_ITER_DELAY_BUFALLOC | NPY_ITER_COPY_IF_OVERLAP; ``` 广播、缓冲未对齐或类型转换的操作数、输入输出之间的重叠检测，所有这些都是迭代器的问题。之后的实际执行几乎是反高潮的： ```c int res; do { res = strided_loop(context, dataptr, countptr, strides, auxdata); } while (res == 0 && iternext(iter)); ``` 迭代器一次交给循环一个大致连续的块，并在之间推进指针。拍拍手，你完成了。 但是：这两条路径也共享两个值得注意的簿记细节。第一，除非循环需要Python API（我们的`float64`循环不需要，`object` dtype循环需要），NumPy会围绕整个操作释放GIL，使用一个很小的阈值，这样小数组就不会付出额外开销。第二，浮点数状态标志在循环之前清除并在循环之后检查，这就是我们熟悉的`RuntimeWarning: overflow encountered in add`的来源。内部循环本身从不检查任何东西，它只是计算并设置CPU标志。 在我们进一步下降之前，有一个警告（它在这篇博文的基准测试中坑了我）。步长本身并不会让你退出简单路径。简单循环接受任何1维数组，并将实际步长直接传递给内部循环，因此即使是`a[::2]`也会在一次调用中处理，看不到迭代器。我重新构建了NumPy并启用了ufunc跟踪来检查，它需要一些简单路径实际拒绝的东西，比如非连续的N维视图（想想`base[:, ::2]`）或一个与输入不安全重叠的`out=`，才会打印“Making iterator”。（重叠检查比你想象的更精细：`out=a[1:]`从`a[:-1]`读取会启动迭代器，而相反方向——写入安全地跟随读取——则不会。）记住这一点，这在一分钟后会很重要。我们现在位于洞穴的中央房间。只低声耳语，天花板非常脆弱。 ## 循环本身在我们看内部循环之前，先给它一些东西来咀嚼。相同数量的元素，相同的操作，但一个版本是每隔一个元素进行步进： ```python 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 call np.add(strided_a, strided_b, out=strided_out) # 8.58 ms per call ``` 在我的机器上，相同的加法次数，性能差了2.6倍。从上一节我们知道两次调用都是以相同的方式到达这里的。简单路径，一次调用进入内部循环，真实的步长等（这就是你之前应该记住的，现在可以放松了）。所以解释这种差距的原因必须存在于`DOUBLE_add`本身内部。部分原因是内存带宽——步进版本触及了两倍的内存！——但让我们看看其余部分。 那么`DOUBLE_add`到底是什么样子呢？在这里我得先让你失望一下，因为仓库中没有任何文件包含名为`DOUBLE_add`的函数。逐元素的循环位于像`loops_arithm_fp.dispatch.c.src` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_core/src/umath/loops_arithm_fp.dispatch.c.src)这样的模板文件中，它们使用了NumPy的内部模板语言，因为每个大型项目都会发展出自己的配置格式和模板语言： ```c /**begin repeat * 浮点类型 * #type = npy_float, npy_double# * #TYPE = FLOAT, DOUBLE# * #sfx = f32, f64# */ /**begin repeat1 * 算术运算 * # kind = add, subtract, multiply, divide# * # intrin = add, sub, mul, div# * # OP = +, -, *, /# */ NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@) (char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func)) { npy_intp len = dimensions[0]; char *src0 = args[0], *src1 = args[1], *dst = args[2]; npy_intp ssrc0 = steps[0], ssrc1 = steps[1], sdst = steps[2]; ... ``` 在构建时，一个小脚本（`conv_template.py` (https://github.com/numpy/numpy/blob/v2.5.2/numpy/_build_utils/conv_template.py)）会为每个值组合扩展每个`repeat`块，因此这一个函数体变成了`FLOAT_add`、`DOUBLE_add`、`FLOAT_subtract`等等，从一个模板生成了八个函数。注意签名：这就是我们两节前遇到的`PyUFuncGenericFunction`。我们到了。