numpy 数组的并行就地排序

2024-01-07

我经常需要对大型 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]

编辑:修复了下面评论中注意到的错误

本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

numpy 数组的并行就地排序 的相关文章

随机推荐