Skip to content

Instantly share code, notes, and snippets.

@nico-zck
Last active February 1, 2023 05:07
Show Gist options
  • Save nico-zck/d166cd362c397fbc2245cf47cfbdcb44 to your computer and use it in GitHub Desktop.
Save nico-zck/d166cd362c397fbc2245cf47cfbdcb44 to your computer and use it in GitHub Desktop.
Parallel argsort via c++ built-in parallel mode.
// #include <algorithm> // std::sort, std::stable_sort
#include <parallel/algorithm>
void argsort(unsigned long *pidx, const double *parr, unsigned long length)
{
// Sort indexes based on comparing values in v
// Because Cython pass data blob as pointer, so access the array as pointer style.
unsigned long *pidx2 = pidx + length;
// double *parr2 = parr + length;
__gnu_parallel::sort(pidx, pidx2,
[parr](size_t i1, size_t i2)
{ return *(parr + i1) < *(parr + i2); });
}

compile:

python3 setup.py build_ext -if

using:

from pargsort import pargsort
import numpy as np

arr = np.random.uniform(0,100,100000000)

%timeit pargsort(arr)
%timeit np.argsort(arr)

Unkonw problem: When processing a very large number of numbers, such as arr = np.random.uniform(0,100,100000000) (an extra zero), the function will degrade to single core.

import numpy as np
cimport numpy as np
import cython
cimport cython
cdef extern from "argsort.h":
cdef void argsort(unsigned long *pidx, double *parr, unsigned long length) nogil
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False) # turn off negative index wrapping for entire function
def pargsort(np.ndarray arr):
"""
In-place parallel sort for numpy types
"""
cdef unsigned long l = arr.shape[0]
cdef double[:] arr2 = arr.astype(np.double)
cdef unsigned long[:] idx2 = np.arange(l, dtype=np.uint)
argsort(&idx2[0], &arr2[0], l) # c++ only consider memview as a pointer, so pass the whole is unnecessary.
return np.asarray(idx2)
from setuptools import Extension, setup
from Cython.Build import cythonize
ext_modules = [
Extension(
"pargsort",
["pargsort.pyx"],
language="c++",
extra_compile_args=["-fopenmp"],
extra_link_args=["-fopenmp"],
)
]
setup(
name="pargsort",
ext_modules=cythonize(ext_modules, compiler_directives={"language_level": "3"}),
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment