NumPy 2.5 如何让 np.searchsorted 提速最高 25 倍
Making np.searchsorted up to 25× Faster in NumPy 2.5
NumPy 2.5 中的 np.searchsorted 通过向量化二分查找实现最高 25 倍提速。新算法用左边界加全局区间长度替代逐查询维护 lo/hi，将额外内存从 O(K) 降到 O(1)，并已移植回 NumPy 的 C++ 实现。该改动同样惠及 SciPy、scikit-learn 等依赖 searchsorted 的 Python 科学计算库。
np.searchsorted 是 NumPy 对二分查找算法的实现。这是基础搜索算法之一，在 Python 科学计算生态中被用于直方图计算和区间查找等功能。任何优化都会让 SciPy 和 scikit-learn 等库以及更广泛的 Python 科学计算生态受益。
一些案例研究，例如 二分查找变体与批处理的影响 和 Algorithmica 的二分查找案例研究，探索了分支消除、批处理和缓存友好的数据布局等技术来提升二分查找性能。在这篇文章中，我们探讨如何用 NumPy 的向量化原语来表达这些想法。
我们将推导出一种向量化公式，其性能优于 NumPy 2.4 的 searchsorted 实现，然后将所得算法移植回 NumPy。这一改动已包含在 NumPy 2.5 中，在我们的基准测试中实现了最高 25 倍的加速。
searchsorted 也是 Python 数组 API 标准 的一部分。这让我们能够比较不同数组库如何实现同一操作并利用并行性。
问题定义#
给定一个静态有序数组和一组查询键，找出每个键在数组中的插入位置。
一个经典的纯 Python 实现为每个键运行一次二分查找：
def searchsorted_py(a, xs):
res = np.empty(len(xs), dtype=np.int32)
for i, x in enumerate(xs):
lo = 0
hi = len(a)
while lo < hi:
mid = (lo + hi) // 2
if a[mid] < x:
lo = mid + 1
else:
hi = mid
res[i] = lo
return res
随着输入规模增长，每次查询的运行时间呈对数增长（注意 x 轴的对数刻度）。
为了进行基准测试，我们生成了两个由均匀分布的 np.int32 整数组成的随机数组。键（即被搜索的元素）长度固定为 10,000，而有序数组的长度变化范围最高达 $2^{30}$（$4\ GiB$）。键数组和值数组在内存中都是连续的。每个基准测试重复 50 次，我们报告最短执行时间。
基准测试在一台配备 Apple M1 Pro 和 32 GB 内存的 MacBook Pro 上运行。M1 Pro 每个性能核心有 128 KB 的 L1 数据缓存，足以容纳 $2^{15}$ 个 32 位整数，其性能核心共享 12 MB 的 L2 缓存，足以容纳 $1.5 * 2^{21}$ 个 32 位整数。
使用 NumPy 数组进行批处理#
基线实现为每个查询执行一次二分查找。每次搜索都是独立的，但在 Python 中按顺序执行。我们可以调整算法，让多个二分查找以批处理方式共同推进。为此，我们可以将所有搜索的状态表示为数组，并使用向量化操作同时更新它们。
在 NumPy 中，数组操作在编译好的 C++ 循环中执行。这消除了 Python 开销，并让 CPU 能够高效处理大批量的独立工作。
让我们向量化二分查找#
为了向量化该算法，我们将标量变量重新解释为数组状态。我们不再使用单个 lo 和 hi，而是为每个查询维护一个值。在原始实现中，状态由 lo、hi 和 res 组成。由于 lo 最终包含最终结果，我们可以只专注于跟踪 lo 和 hi。
这第一个向量化实现是之前算法的直接翻译。
def searchsorted_py_np(a, xs):
lo = np.zeros(xs.shape, dtype=np.int32)
hi = np.full(xs.shape, len(a), dtype=np.int32)
while True:
# True for each element position where lo_i < hi_i
active = lo < hi
if not np.any(active):
# this is basically `while lo < hi:` in the pure-Python version
break
mid = (lo + hi) // 2
# only index those entries where active is true, so we don't modify already computed positions
mid_a = mid[active]
xs_a = xs[active]
mask = a[mid_a] < xs_a
lo[active] = np.where(mask, mid_a + 1, lo[active])
hi[active] = np.where(mask, hi[active], mid_a)
return lo
在这个实现中，不同查询可以以不同速率缩小其搜索区间，因此它们可能需要不同数量的迭代才能收敛。
例如，考虑在数组 [0, 1] 中搜索键 [-1, 2]。对于查询 2，第一次迭代计算 mid = 1 并设置 lo = mid + 1 = 2，因此区间变为 [2, 2)，搜索一步收敛。对于查询 -1，更新则设置 hi = mid = 1，第一步后留下区间 [0, 1)。这需要额外一次迭代才能将区间收缩到 [0, 0)。
由于各搜索可能在不同时间收敛，我们需要跟踪哪些查询仍然活跃。活跃掩码标识出搜索区间尚未收敛的查询，使我们只更新这些查询。
让所有搜索都执行相同的步数#
重要的观察是，二分搜索实际上不需要为每个键独立终止。我们可以调整每次迭代的更新方式，使得一旦某个搜索收敛，后续迭代可以保持其区间不变。
我们维护插入位置位于 [lo, hi) 中的不变式。在每次迭代中，每个区间都缩小到大约之前大小的一半。经过 np.ceil(np.log2(n)) 次迭代后，每个区间都收缩到单个位置。
为简单起见，此实现假设 a 非空。
def searchsorted_py_np_fixed(a, xs):
n = len(a)
lo = np.zeros(xs.shape, dtype=np.int32)
hi = np.full(xs.shape, n, dtype=np.int32)
for _ in range(int(np.ceil(np.log2(n)))):
mid = (lo + hi) // 2
go_left = xs <= a[mid]
# for each position i:
# if go_left_i is True, we keep the `lo_i` value and `hi_i` is updated to `mid_i`
# if go_left_i is False, we keep the `hi_i` value and `lo_i` is updated to `mid_i`
lo = np.where(go_left, lo, mid)
hi = np.where(go_left, mid, hi)
return hi移除 active 跟踪器可使其速度提升至多 2 倍。
类似的表述在 Python 生态系统中已经存在。例如，JAX 基于 scan 的实现
...
def body_fun(state, _):
low, high = state
mid = low + (high - low) // 2 # use this form to avoid overflow
go_left = op(query, sorted_arr[mid])
return (lax.select(go_left, low, mid), lax.select(go_left, mid, high)), ()
n_levels = int(np.ceil(np.log2(n + 1)))
...NumPy 能击败 NumPy 吗？#
让我们比较这个向量化实现与 NumPy 原生 searchsorted（使用 NumPy 2.4）的性能。
对于包含多个键的输入，我们的向量化 Python 实现可以比原生实现快一个数量级。要理解原因，让我们看看 NumPy 2.4 的实现：
template <class Tag, side_t side>
static void
binsearch(const char *arr, const char *key, char *ret, npy_intp arr_len,
npy_intp key_len, npy_intp arr_str, npy_intp key_str,
npy_intp ret_str, PyArrayObject *)
{
using T = typename Tag::type;
auto cmp = side_to_cmp<Tag, side>::value;
npy_intp min_idx = 0;
npy_intp max_idx = arr_len;
T last_key_val;
if (key_len == 0) {
return;
}
last_key_val = *(const T *)key;
for (; key_len > 0; key_len--, key += key_str, ret += ret_str) {
const T key_val = *(const T *)key;
/*
* Updating only one of the indices based on the previous key
* gives the search a big boost when keys are sorted, but slightly
* slows down things for purely random ones.
*/
if (cmp(last_key_val, key_val)) {
max_idx = arr_len;
}
else {
min_idx = 0;
max_idx = (max_idx < arr_len) ? (max_idx + 1) : arr_len;
}
last_key_val = key_val;
while (min_idx < max_idx) {
const npy_intp mid_idx = min_idx + ((max_idx - min_idx) >> 1);
const T mid_val = *(const T *)(arr + mid_idx * arr_str);
if (cmp(mid_val, key_val)) {
min_idx = mid_idx + 1;
}
else {
max_idx = mid_idx;
}
}
*(npy_intp *)ret = min_idx;
}
}忽略指针运算细节，核心算法是对每个键独立执行的经典二分搜索。还有一个优化，当输入键已排序时复用之前的搜索边界。
此实现为每个键执行一次二分搜索，其中每次搜索都是完全顺序的过程。二分搜索的每次迭代都依赖于上一次的结果（中点决定接下来检查数组的哪一部分）。这在每次搜索中形成依赖链：下一次内存访问依赖于上一次比较的结果。
对于大型数组，二分搜索的读取也往往对缓存不友好，因为每一步可能需要从不同的缓存行读取。缓存未命中对顺序实现的影响更大，因为每一步都可能因等待上一次内存访问完成而停顿。
向量化实现一次性对所有查询执行相同的逻辑步骤（所有查询在每一步一起推进）。在有多个独立搜索的情况下，CPU 可以同时进行多个内存访问。这与数组大小超过 L1 和 L2 缓存大小后观察到的运行时间相符。
我们能优化 NumPy 吗？#
之前的向量化实现维护两个数组 lo 和 hi，来表示每个查询的搜索区间。如果我们将这个确切的实现原生移植到 NumPy 中，它将需要使用 $O(K)$ 的额外内存，其中 K 是查询数量。尽管这种方法可能更快，但对于内存敏感的工作负载来说是不可接受的。
为了减少所需的状态，我们可以用区间边界来重新表述二分查找。不再为每个查询同时跟踪 lo 和 hi，而是用左边界和长度来描述每个区间。
关键观察是，如果我们把算法设计成所有查询在每次迭代中都将各自的区间收缩相同的量，那么每个区间在每次迭代中都具有相同的长度。这意味着我们不需要为每个查询存储单独的 hi：它可以由一个数组 lo 和一个全局长度重建出来。注意这仍然需要 $O(K)$ 的空间来存放输出，但除此之外只需要 $O(1)$ 的内存。
这给出了如下的 Python 实现：
def searchsorted_py_np_fast_where(arr, keys):
K = keys.shape[0]
length = arr.shape[0]
base = np.zeros(K, dtype=np.intp)
# Invariant: the insertion index lies in [base, base + length]
while length > 1:
half = length >> 1
mid = base + half
base = np.where(keys > arr[mid], mid, base)
length -= half
# Final step: result is either base and base + 1
base = np.where(keys > arr[base], base + 1, base)
return base只需要一个数组 base 来存储每个查询的状态，而 length 由所有查询共享。输出直接写入 base，因此不需要额外的状态数组。这个实现仍然需要分配一个临时数组 mid 来保存中点，但在 C++ 移植中可以避免这一点。
这种表述与 Algorithmica 的案例研究 中讨论的无分支二分查找方法密切相关。在我们使用的表述中，不变区间是 [base, base + length]。因此当 length = 1 时需要最后一步来确定插入点落在 base 的左侧还是右侧。
重新表述后的实现明显更快：
将其移植到 C++#
性能结果显示了减少每个查询状态所带来的好处。现在我们可以几乎直接将其翻译成 C++，只需 $O(1)$ 的额外内存。
template <class Tag, side_t side>
static void
binsearch(const char *arr, const char *key, char *ret, npy_intp arr_len,
npy_intp key_len, npy_intp arr_str, npy_intp key_str,
npy_intp ret_str, PyArrayObject *)
{
using T = typename Tag::type;
auto cmp = side_to_cmp<Tag, side>::value;
// If the array length is 0 we return all 0s
if (arr_len <= 0) {
for (npy_intp i = 0; i < key_len; ++i) {
*(npy_intp *)(ret + i * ret_str) = 0;
}
return;
}
/*
base = np.zeros(K, dtype=np.intp)
We unroll the first iteration for the following reasons:
1. ret is not initialized with the bases, so we save |keys| writes
by not having to initialize it with 0s.
2. By assuming the initial base for every key is 0, we also save
|keys| reads.
3. In the first iteration, all elements are compared against the
median. So we can store it in a variable and use it for all keys.
This initial block replaces the initialization loop that is used for the
arr_len==0 case. Note that when arr_len = 1, then half is 0 so the
following block initializes the array as with 0s.
*/
npy_intp interval_length = arr_len;
npy_intp half = interval_length >> 1;
interval_length -= half; // length -> ceil(length / 2)
npy_intp base = 0;
const T mid_val = *(const T *)(arr + (base + half) * arr_str);
for (npy_intp i = 0; i < key_len; ++i) {
const T key_val = *(const T *)(key + i * key_str);
*(npy_intp *)(ret + i * ret_str) = cmp(mid_val, key_val) * half;
}
/*
while length > 1:
half = length >> 1
length -= half
mid = base + half
base = np.where(keys > arr[mid], mid, base)
*/
while (interval_length > 1) {
npy_intp half = interval_length >> 1;
interval_length -= half;
for (npy_intp i = 0; i < key_len; ++i) {
npy_intp &base = *(npy_intp *)(ret + i * ret_str);
const T mid_val = *(const T *)(arr + (base + half) * arr_str);
const T key_val = *(const T *)(key + i * key_str);
base += cmp(mid_val, key_val) * half;
}
}
// base = np.where(keys > arr[base], base + 1, base)
for (npy_intp i = 0; i < key_len; ++i) {
npy_intp &base = *(npy_intp *)(ret + i * ret_str);
const T key_val = *(const T *)(key + i * key_str);
base += cmp(*(const T *)(arr + base * arr_str), key_val);
}
}注意我们利用了二分查找第一次迭代的一个性质。由于每个结果条目的初始值都隐式为零，我们可以在第一次迭代中跳过对这些值的写入和读取。此外，在第一次迭代中所有元素都与同一个中位数比较，因此我们可以只读取一次它的值，而不是 K 次。
这个实现作为 PR #30517 的一部分被直接移植到了 NumPy 中，该 PR 包含在 2.5 版本 中。现在让我们对 NumPy 2.4 和 2.5，以及我们的向量化 Python 实现做最后的比较：
在我们的基准测试中，原生的 2.5 版本比 NumPy 2.4 的实现快最多 25 倍。与向量化的 Python 实现相比，C++ 实现在较小的数组上可以快最多 2 倍。随着数组规模增大，这一差距变得不那么显著。与向量化的 Python 实现相比还有一个内存优势：Python 实现需要额外的数组来存储搜索状态（low 和 mid，或 base 和 base + length），而 C++ 实现将 length 保持为标量。因此，C++ 实现只使用 $O(1)$ 的额外内存，而 NumPy 的表述需要与查询数量成正比的内存。
生态系统比较#
我们可以将优化后的 NumPy 2.5 与生态系统中的其他库进行比较。在这个实验中，我们选择了 Python 库 JAX、TensorFlow 和 PyTorch。
TensorFlow 和 PyTorch 采用与 JAX 和 NumPy 不同的方法。JAX 和 NumPy 利用向量化/批量操作来隐藏内存延迟，而 TensorFlow 和 PyTorch 则跨 CPU 线程并行化独立的搜索。搜索键被划分为批次，由不同的线程处理。更多细节请参见 PyTorch 和 TensorFlow 的实现。
在基准测试中，我们将并行度限制为 8 核，并将查询键的数量从 10,000 增加到 20,000。这为多线程实现提供了足够的独立工作，以摊销线程调度开销。
基准测试表明，在我们的基准测试中，NumPy 与所选库相比具有竞争力。一旦搜索数组超出 CPU 缓存，所有实现都表现出相似的行为。
如果我们禁用多线程，PyTorch 和 TensorFlow 的性能会下降，并且两者都表现出与 NumPy 2.4 实现相似的趋势。一旦搜索数组超出 CPU 缓存，内存访问的成本就占主导地位。
值得对这两种技术是否可以结合进行基准测试：在每个线程内批处理二分搜索。然而，一旦内存子系统饱和，额外的核心可能会竞争相同的内存带宽。此时，改进内存访问模式可能是一个更有前景的方向，例如使用不同的布局，如 Eytzinger 布局（在 Algorithmica 书中有详细讨论）。
结论#
在我们的基准测试中，我们将 np.searchsorted 的速度提升了高达 25 倍。鉴于 NumPy 在 Python 生态系统中的广泛影响，这一优化将使多个依赖它的库受益。Python 生态系统中其他拥有自己二分搜索实现的库也可能从采用类似的批处理技术中受益。
有趣的是，我们使用 NumPy 数组原语推导出了一个初始 Python 实现，其性能超过了 NumPy 2.4 的实现。这表明 NumPy 的数组原语在实现高性能算法方面有多么强大。通过利用独立工作，Python 中的向量化 NumPy 实现可以超越标量原生实现。
诸如 Eytzinger 布局之类的缓存友好布局是使 searchsorted 更快的另一个有趣方向。探索数组 API 是否可以通过像 searchsorted(arr, keys, layout="eytzinger") 这样的接口暴露此类布局将会很有趣，尽管这需要仔细定义 API 语义，因为 Eytzinger 表示不是有序的。
来源：Hacker News · blog.scientific-python.org