我经常需要对大型 numpy 数组(几十亿个元素)进行排序,这成为了我的代码的瓶颈。我正在寻找一种并行化它的方法。
该功能是否有任何并行实现ndarray.sort()
?Numexpr 模块为 numpy 数组上的大多数数学运算提供并行实现,但缺乏排序功能。
也许,可以围绕并行排序的 C++ 实现制作一个简单的包装器,并通过 Cython 使用它?
我最终包装了 GCC 并行排序。这是代码:
并行排序.pyx
# cython: wraparound = False
# cython: boundscheck = False
import numpy as np
cimport numpy as np
import cython
cimport cython
ctypedef fused real:
cython.char
cython.uchar
cython.short
cython.ushort
cython.int
cython.uint
cython.long
cython.ulong
cython.longlong
cython.ulonglong
cython.float
cython.double
cdef extern from "<parallel/algorithm>" namespace "__gnu_parallel":
cdef void sort[T](T first, T last) nogil
def numpyParallelSort(real[:] a):
"In-place parallel sort for numpy types"
sort(&a[0], &a[a.shape[0]])
额外的编译器参数:-fopenmp(编译)和-lgomp(链接)
这个makefile会做到这一点:
all:
cython --cplus parallelSort.pyx
g++ -g -march=native -Ofast -fpic -c parallelSort.cpp -o parallelSort.o -fopenmp `python-config --includes`
g++ -g -march=native -Ofast -shared -o parallelSort.so parallelSort.o `python-config --libs` -lgomp
clean:
rm -f parallelSort.cpp *.o *.so
这表明它有效:
from parallelSort import numpyParallelSort
import numpy as np
a = np.random.random(100000000)
numpyParallelSort(a)
print a[:10]
编辑:修复了下面评论中发现的错误
合并排序很自然地并行化。只需让每个工作人员预先对任意块进行排序,然后在其上运行一个合并通道。最终的合并应该只需要 O(N) 操作,并且在 numba 或类似的东西中编写一个这样做的函数是微不足道的。