在 NumPy 2.5 中让 np.searchsorted 速度提升最高 25 倍

# 让 np.searchsorted 在 NumPy 2.5 中提速最高达 25 倍 来源：https://blog.scientific-python.org/numpy/searchsorted/ `np.searchsorted` 是 NumPy 对二分查找算法的实现。这是最基础的搜索算法之一，被广泛用于 Python 科学生态系统中的直方图计算、区间查找等功能。对它的任何优化都将惠及 SciPy、scikit-learn 等库以及更广泛的 Python 科学生态系统。 已有多个案例研究探索了提升二分查找性能的技术，包括[二分查找变体与批处理的效果](https://curiouscoding.nl/posts/binsearch/)以及 [Algorithmica 的二分查找案例研究](https://en.algorithmica.org/hpc/data-structures/binary-search/)，涉及分支消除、批处理和缓存友好数据布局等技术。本文将探讨如何利用 NumPy 的向量化原语来表达这些思想。我们将推导出一种向量化方案，其性能超过 NumPy 2.4 的 `searchsorted` 实现，然后将最终算法移植回 NumPy。该变更已包含在 NumPy 2.5 中，在我们的基准测试中最高实现了 25 倍的加速。 `searchsorted` 也是 [Python Array API 标准](https://data-apis.org/array-api/latest/API_specification/generated/array_api.searchsorted.html#array_api.searchsorted)的一部分。这使我们能够比较不同数组库如何实现同一操作并利用并行性。 ## 问题定义 给定一个静态有序数组和一系列查询键，找出每个键在数组中的插入位置。 经典的纯 Python 实现对每个键执行一次二分查找： ```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`。 第一个向量化实现是之前算法的直接翻译： ```python 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: # 对于每个 lo_i < hi_i 的元素位置，为 True active = lo < hi if not np.any(active): # 这基本上等同于纯 Python 版本中的 `while lo < hi:` break mid = (lo + hi) // 2 # 只对 active 为 true 的条目进行索引，这样就不会修改已经计算完成的位置 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` 非空。 ```python 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] # 对于每个位置 i： # 如果 go_left_i 为 True，保留 `lo_i` 的值，将 `hi_i` 更新为 `mid_i` # 如果 go_left_i 为 False，保留 `hi_i` 的值，将 `lo_i` 更新为 `mid_i` lo = np.where(go_left, lo, mid) hi = np.where(go_left, mid, hi) return hi ``` 移除 `active` 追踪器使速度提升了最高 2 倍。 类似的公式化方法已经在 Python 生态系统中存在。例如，[JAX 基于 scan 的实现](https://github.com/jax-ml/jax/blob/a6e4a8b95a731269bdf23e5b3e30da2f8494bb28/jax/_src/numpy/hijax.py#L330)： ```python ... def body_fun(state, _): low, high = state mid = low + (high - low) // 2 # 使用这种形式以避免溢出 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` 的实现： ```cpp template <...> 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 *NPY_UNUSED(context)) { using T = typename Tag::type; auto cmp = side_to_cmp<T, 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; /* * 根据前一个键只更新一个索引，当键已排序时 * 这给搜索带来了很大提升，但对完全随机的键 * 会略微降低速度。 */ 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 实现： ```python def searchsorted_py_np_fast_where(arr, keys): K = keys.shape[0] length = arr.shape[0] base = np.zeros(K, dtype=np.intp) # 不变式：插入索引位于 [base, base + length] 中 while length > 1: half = length >> 1 mid = base + half base = np.where(keys > arr[mid], mid, base) length -= half # 最后一步：结果是 base 或 base + 1 base = np.where(keys > arr[base], base + 1, base) return base ``` 只需要单个数组 `base` 来存储每个查询的状态，而 `length` 由所有查询共享。输出直接写入 `base`，因此不需要额外的状态数组。此实现仍然需要分配一个临时的 `mid` 来存储中点值，但在 C++ 移植中可以避免这一点。 这种公式化方法与 [Algorithmica 案例研究](https://en.algorithmica.org/hpc/data-structures/binary-search/#removing-branches)中讨论的无分支二分查找方法密切相关。在我们使用的公式中，不变式范围是 `[base, base + length]`。因此当 `length = 1` 时需要一个最终步骤来确定插入点是落在 `base` 的左边还是右边。 重新表述后的实现显著更快： ### 移植到 C++ 性能结果展示了减少每个查询状态的好处。现在我们可以以 $O(1)$ 的额外内存几乎直接将其翻译为 C++。 ```cpp template <...> 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 *NPY_UNUSED(context)) { using T = typename Tag::type; auto cmp = side_to_cmp<T, Side>::value; // 如果数组长度为 0，我们返回全 0 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) * 我们展开第一次迭代，原因如下： * 1. ret 未用 base 初始化，因此通过不需要将其用 0 初始化，节省了 |keys| 次写入。 * 2. 通过假设每个键的初始 base 为 0，我们还节省了 |keys| 次读取。 * 3. 在第一次迭代中，所有元素都与中位数进行比较。 * 因此我们可以将其存储在一个变量中并用于所有键。 * 这个初始块替代了 arr_len==0 情况下使用的初始化循环。 * 注意当 arr_len = 1 时，half 为 0，因此后续块会像用 0 初始化那样初始化数组。 */ 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` 次。 此实现已直接移植到 NumPy，作为 [PR #30517](https://github.com/numpy/numpy/pull/30517) 的一部分，包含在 [2.5 版本](https://numpy.org/devdocs/release/2.5.0-notes.html#improved-performance-of-numpy-searchsorted)中。 现在让我们在 NumPy 2.4 和 2.5 以及我们的向量化 Python 实现之间做最终比较： 在我们的基准测试中，原生 2.5 版本比 NumPy 2.4 的实现最高快了 25 倍。与向量化 Python 实现相比，C++ 实现在较小数组上可以快至 2 倍。随着数组大小的增长，这种差异变得不那么显著。 此外，在内存使用方面，C++ 实现也优于向量化 Python 实现：Python 实现需要额外的数组来存储搜索状态（`low` 和 `mid`，或 `base` 和 `base + length`），而 C++ 实现将 `length` 保持为标量。因此，C++ 实现仅使用 $O(1)$ 的额外内存，而 NumPy 公式化方法需要与查询数量成比例的内存。