Python 中使用自由线程实现NumPy扩展

编者按

本文定位NumPy在free-threaded Python下多线程扩展的五个瓶颈：tracemalloc与ufunc派发缓存的锁竞争、全局PyCapsule引用计数、模块属性查找及内存分配器。修复关键是把派发缓存改为原子加载的无锁读、把全局对象设为immortal、用CPython raw allocator接入mimalloc。32线程复现耗时由44秒降至1.5秒，比多进程快4倍。

来源：labs.quansight.org^[1]

引言

NumPy 是科学计算 Python 生态中的基础数组库。所有数值计算、机器学习和数据分析库要么直接依赖 NumPy，要么与其互操作。随着 CPython 自由线程构建（free-threaded build）日趋成熟，当用户尝试用线程在多个核心上扩展 CPU 密集型数值计算负载时，NumPy 往往是他们最先拿起来的库。

一切的源头

这项调查源于一个 StackOverflow 问题^[2]：有用户报告称，在 CPython 自由线程构建上，使用 ThreadPoolExecutor 的 NumPy 负载明显慢于等价的 ProcessPoolExecutor 版本。该问题记录在 numpy/numpy#30494^[3]。

复现程序非常小巧，而且是 ufunc 密集型负载的典型代表：每个 worker 拿一个数组，在循环里调用若干次 np.sin 和 np.cos，最后对结果做归约。以下是每个 worker 运行的核心代码：

def kernel(x, n_loop):
return sum((np.sin(np.cos(np.sin(np.cos(x + i)))).sum()
for i in range(n_loop)))

每个 worker 使用 ThreadPoolExecutor 或 ProcessPoolExecutor，在自己的数组分片上运行这段核心代码。worker 之间没有任何共享的可变状态，所以原则上应该随核心数线性扩展。但实际上，在同一台机器上，多线程比多进程慢了最多 2 倍。

自由线程构建移除了 GIL，但仅仅移除 GIL 本身还不够。用 samply^[4] 对复现程序进行剖析后，发现了 NumPy 和 CPython 中几处被 GIL 掩盖的隐藏瓶颈——在常规构建上它们无伤大雅，在自由线程构建上却成了严重的扩展性问题。这些扩展性瓶颈主要分为三类：

1. 锁竞争。
2. 共享全局对象上的引用计数竞争。
3. 内存分配器中的竞争。

修复前复现程序的火焰图，标注了主要的扩展性瓶颈。

瓶颈与修复

1. tracemalloc 中的锁竞争

最先浮出水面的瓶颈位于 CPython 的 tracemalloc 模块。tracemalloc 是一个默认关闭的内存跟踪工具。但问题是，即使它没有被启用，每次内存分配和释放时，它仍然会获取一把全局锁，用来检查 tracemalloc 是否在运行时被启用了。

我在 CPython 中修复了这一点：当 tracemalloc 处于禁用状态时不再加锁。现在 CPython 使用原子操作检查 tracemalloc 是否启用，从而提供了一条无锁快速路径。实现见 python/cpython#143065^[5]。

这类瓶颈很容易引入，却很难发现。它也是一个很常见的模式：在执行某个操作前需要先检查一个标志位。加锁固然简单，也能解决线程安全问题，但如果这个标志位在热路径上被频繁检查，就会造成严重的扩展性瓶颈。这种情况下，更好的做法是用原子操作检查标志位，在快速路径上避免加锁，只在需要更新标志位时才去获取锁。

2. ufunc 分发缓存中的锁竞争

NumPy 把 np.sin 这类数学运算实现为 ufunc^[6]（universal function，通用函数）。这里需要理解的关键点是：ufunc 是围绕输入类型来定义的——NumPy 为每一种类型组合都实现了不同的循环。例如，np.add 的实现取决于操作数是整数还是浮点数，因为整数需要基于整数加法的循环，浮点数则需要基于浮点数加法的循环。因此 NumPy 有一套分发系统，根据输入数据类型元组来确定该调用哪个循环。

为了避免每次调用都重复这一解析过程，NumPy 把结果缓存在一个分发缓存中，该缓存将输入类型元组映射到具体的循环实现。这个缓存是所有线程共享的全局状态。它之前由 std::shared_mutex（读写锁）保护，但在大量并发读访问下扩展性依然不佳。

新设计利用了分发缓存的两个特性来解决这个问题：读多写少，且条目不可变——条目一旦插入，就永远不会被修改或删除。这意味着读操作完全不需要防备条目在眼皮底下变化，因此读可以做到完全无锁。现在它实现为一个无锁并发哈希表，读路径完全不需要加锁，只有极少数的缓存未命中需要插入时，写路径才需要一把锁。

无锁分发缓存。读操作顺着原子指针找到当前桶表并查找条目，无需加锁。表扩容后，旧表会一直存活（通过 prev 链连接）直到缓存本身被释放，这样仍在引用旧表的读操作——比如图中的 Reader C——可以安全地完成任务。只有写操作需要获取互斥锁。

该缓存由 PyArrayIdentityHash 实现，它持有一个指向桶的原子指针，真正的条目就存在桶里。查找条目时，线程先原子地加载桶指针，根据输入类型的哈希值索引到相应位置，再原子地加载存储在那里的键。如果键与输入类型匹配，就返回对应的循环。也就是说，查找条目只是一连串原子加载，完全不加锁，即使大量线程同时读缓存也能良好扩展。

插入操作很少发生，只在缓存未命中时出现。写操作会获取缓存的互斥锁，保证同一时刻只有一个线程能修改缓存，然后原子地把新条目发布到桶中，这样并发的读操作也能无数据竞争地读到更新后的条目。如果缓存需要扩容，写操作会创建一张新的桶表，把所有条目重新哈希进去，再原子地更新指针指向新表。旧表会一直存活到缓存本身被释放，因此仍在引用旧表的读操作可以安全地完成，不用担心读取过程中内存被释放。这种设计让分发缓存在多线程环境下既线程安全，又能良好扩展。实现见 numpy/numpy#30593^[7]。

3. 全局 PyCapsule 引用计数竞争

NumPy 允许用户通过可配置的内存处理器^[8]自定义数组的内存分配方式。为了让这个处理器能从 Python 侧访问并且在运行时可以替换，NumPy 把它包装在一个 PyCapsule^[9] 对象中——这种对象携带一个不透明的 C 指针。

NumPy 把默认内存处理器存在一个全局 PyCapsule 对象里。在 GIL 构建下，对这些 capsule 对象做额外的 Py_INCREF/Py_DECREF 几乎零成本：引用计数只是一个普通整数，同一时刻只有一个线程能碰它。而在自由线程构建下，共享对象的每次引用计数增减都必须是对引用计数的原子操作。当大量线程对同一个对象做引用计数时，持有该计数的缓存行会在核心之间来回弹跳，而不是留在某个核心本地。这导致对共享对象做引用计数的开销远高于线程本地对象，当大量线程同时引用同一个作为默认内存处理器的全局 PyCapsule 对象时，性能会严重退化。

我的修复方案是让全局默认内存处理器变成 immortal。immortal 对象永远不会被释放，也完全不需要做引用计数，竞争问题就此彻底绕开。CPython 早已在 None、True、小整数等内部单例上广泛使用 immortal 对象，原因恰恰是它能让这些对象跨线程共享而不产生引用计数竞争。

问题在于，Python 3.15 之前，像 NumPy 这样的库没有公开的 C API 能把自己的对象变成 immortal。这个使用场景正好为添加这样一个 API 提供了充分理由，于是我在 CPython 中加入了 PyUnstable_SetImmortal。实现见 python/cpython#144543^[10] 和 numpy/numpy#30826^[11]。该 API 从 Python 3.15 起公开可用，不过 NumPy 已经可以通过 pythoncapi-compat^[12] 头文件在 Python 3.14 上使用它——这些头文件借助 3.14 中已有的私有 C API 实现了该函数。

4. 模块属性查找竞争

当你写 np.sin 时，那是对 numpy 模块对象的一次属性访问——Python 需要在模块上查找 sin 属性。NumPy 通过模块级的 __getattr__ 把 np.sin、np.cos 这类 ufunc 解析到它们的实际实现。在 CPython 中，对于定义了 __getattr__ 的模块，模块属性查找的字节码特化不会被启用，导致属性查找走上慢路径：获取 import 锁，然后执行查找。

我在 CPython 上游修复了这个问题（python/cpython#143470^[13]）：即使定义了 __getattr__，也照常为模块属性查找启用字节码特化。该修复从 Python 3.15 起可用。

5. 内存分配器竞争

NumPy 直接通过 malloc 和 free 从系统分配器为数组分配内存。当大量线程并发分配和释放内存时，扩展性很差，因为某些系统分配器会在内部锁后面把分配操作串行化。这在 macOS 上尤其糟糕——多线程同时分配和释放内存时，系统分配器的扩展性非常差。

CPython 暴露了自己的底层 “raw”内存分配例程^[14]，例如 PyMem_RawMalloc 和 PyMem_RawFree。然而这些例程同样在调用系统分配器。CPython 自由线程构建已经用 mimalloc^[15] 来分配 Python 对象，mimalloc 针对多线程负载做了高度优化。但 NumPy 并没有使用它，因为 NumPy 是直接调用系统分配器的。

修复分两部分：

• 在 CPython 中，我把原始分配器 API 改为在自由线程构建上使用 mimalloc 作为底层内存分配器。实现见 python/cpython#144916^[16]。
• 在 NumPy 中，数组分配 API 则改为使用 CPython 的原始分配器，不再直接调用系统分配器。让数组分配走这些例程，NumPy 就能自动受益于 mimalloc——完全不需要知道 mimalloc 的任何细节。实现见 numpy/numpy#30846^[17]。

基准测试

基准测试用的正是引发本次调查的 StackOverflow 问题中的复现程序：每个 worker 在循环中对各自的数组调用若干次 np.sin 和 np.cos，最后归约结果，worker 之间没有共享可变状态。下面是在一台 32 核 Linux 机器上，应用上述全部修复前后，自由线程构建上多线程复现程序的性能对比：

复现程序耗时随 worker 数量增长的变化，对比修复前多线程、修复后多线程和多进程。

修复前，多线程版本在 18 个线程以内扩展良好，但之后由于上述瓶颈，性能急剧恶化：32 个 worker 时耗时约 44 秒，比多进程版本慢了约 7 倍。修复后，多线程版本在全部 32 个核心上都能良好扩展。同样的负载现在只需约 1.5 秒——比修复前快了约 30 倍，比耗时约 6 秒的多进程版本也快了约 4 倍。

总结

修复了 NumPy 和 CPython 中的多个瓶颈之后，NumPy ufunc 现在可以在 CPython 自由线程构建上良好扩展。我在 CPython 中针对 tracemalloc、内存分配器和模块属性查找所做的修复，也会让自由线程构建上的其他库和负载受益，而不仅仅是 NumPy。此外，这个项目还完成了一些基础性工作，比如新增了让对象 immortal 的 C API、让原始分配器改用 mimalloc——这将使更多库能够轻松修复自己代码中类似的瓶颈，并在自由线程构建上获得良好扩展性。

#NumPy #自由线程 #多线程 #性能优化 #CPython