并行化 numpy.sort

我需要对长度为1e8-1e9的uint64数组进行排序，这是我当前项目中的性能瓶颈之一。我最近刚刚更新了numpy v2.0版本，其中排序算法得到了显着优化。在我的硬件上测试它，它比 numpy v1.26 版本快大约 5 倍。但目前 numpy 的排序算法即使使用 SIMD，也无法利用多核 CPU。

我尝试将其并行化并同时对多个 np.array 进行排序。一种可能的方法是使用 numba prange，但 numba 对 numpy 排序的支持一直很差。 numba.jit 甚至对 np.sort 有减慢效果，并且 numba v0.60.0 未能跟进 numpy v2.0 对排序的优化（ https://github.com/numba/numba/issues/9611| ||)。替代方案是 cython prange，但 cython 不允许在 nogil 处创建 Python 对象。有没有办法使用 cython 或其他方式并行排序 numpy.array ？如果使用 cpp 的并行排序库，考虑到数据类型转换的开销，它们是否比 numpy 自己的排序更快？ ). The alternative is cython prange, but cython does not allow the creation of Python objects at nogil. Is there a way to sort numpy.array in parallel using cython or otherwise? If using cpp's parallel sorting libraries, are they faster than numpy's own sorting, taking into account the overhead of data type conversions?

arr=np.random.randint(0,2**64,int(3e8),dtype='uint64')  

sorted_arr=np.sort(arr)  # single thread np.sort takes 4 seconds (numpy v2.0.0)

当然，我可以帮助你探索在 Python 中并行化 numpy.sort 的方法。你已经概述了你面临的挑战，即 NumPy 的排序算法本身并不会跨多个核进行并行化。以下是你可以尝试的一些方法：

1. 基于分块的并行化与 multiprocessing

2. 方法： 将你的数组分成若干个块，使用 multiprocessing.Pool 并行排序这些块，然后合并已排序的块。

3. 优点： 相对容易实现；绕过了 NumPy 排序中的 GIL 限制。
4. 缺点： 合并已排序的块（例如，使用 numpy.concatenate 后跟另一个 numpy.sort ）可能会引入开销，特别是在处理大型数组时。

```python import numpy as np from multiprocessing import Pool, cpu_count

def sort_chunk(arr): return np.sort(arr)

def parallel_sort(arr, num_processes=None): if num_processes is None: num_processes = cpu_count() chunk_size = len(arr) // num_processes chunks = [arr[i * chunk_size:(i + 1) * chunk_size] for i in range(num_processes - 1)] chunks.append(arr[(num_processes - 1) * chunk_size:]) # 最后一块 with Pool(num_processes) as pool: sorted_chunks = pool.map(sort_chunk, chunks) sorted_arr = np.concatenate(sorted_chunks) return np.sort(sorted_arr) # 再次排序以合并

arr = np.random.randint(0, 2**64, int(3e8), dtype='uint64') sorted_arr = parallel_sort(arr) ```

1. Cython 与外部并行排序库

2. 方法： 使用 Cython 将你的 NumPy 数组传递给用 C/C++ 编写的并行排序库（例如，来自 STL 的 std::sort 与执行策略或 Intel TBB）。

3. 优点： 可能会带来显著的性能提升，特别是在对排序库进行微调以充分利用你的硬件时。
4. 缺点： 比基于 multiprocessing 的方法更复杂；需要 Cython 知识以及与外部库的链接。

```cython # 需要安装一个外部并行排序库并进行链接 cdef extern from "some_parallel_sort_library.h": void parallel_sort(unsigned long long *arr, size_t n)

def sort_array(np.ndarray[np.uint64_t, ndim=1] arr): cdef size_t n = arr.size with nogil: parallel_sort(&arr[0], n) return arr ```

1. 其他 Python 并行化库

2. 探索诸如 Dask 或 Ray 之类的库，它们提供了对并行计算的高级抽象，并且可以有效地处理大型数组。这些库通常可以很好地与 NumPy 数组配合使用，并且可以提供通过分布式计算进行扩展的途径。

评估和考虑因素

• 数据集大小： 对于非常小的数组，并行化的开销可能超过了性能优势。如果你的数组已经很大（如你的情况），并行化很可能是值得的。
• 硬件： 可用内核的数量和内存带宽将显着影响并行化的有效性。对不同的方法进行基准测试，找到适合你特定硬件的最佳方法。
• 开销： 请注意数据移动、进程间通信和合并已排序结果相关的开销。

请记住，没有一种放之四海而皆准的最佳方法。我建议你试验这些不同的方法，并使用真实数据对它们进行基准测试，以确定哪种方法最适合你的用例。