Implement cupy.partition

[takagi]

cupy.partition, the counterpart of numpy.partition, creates a copy of an array with its elements rearranged in such a way that the value of the element in k-th position is in the position it would be in a sorted array. This function is equivalent to std::nth_element of C++ STL.

numpy.partition
https://docs.scipy.org/doc/numpy/reference/generated/numpy.partition.html

std::nth_element
http://en.cppreference.com/w/cpp/algorithm/nth_element

Unfortunately, neither of the following parallel algorithm libraries implement this kind of parallel selection algorithm:

• Thrust (https://github.com/thrust/thrust)
• Modern GPU (https://github.com/moderngpu/moderngpu)
• CUB (https://nvlabs.github.io/cub/)

For now, the most likely is to implement radix select based on CUB's radix sort implementation.

[takagi]

cupy.partiotion as well as numpy.partition belong to a problem domain called unordered partial sort, which is a very similar problem with selection of kth element.

https://en.wikipedia.org/wiki/Selection_algorithm#Unordered_partial_sorting

There are the following algorithms to solve such kind of problems:

• Quickselect
• Median of median
• Heap select
• Introselect
• Radixselect

numpy.partition uses introselect, which is a hybrid algorithm of quickselect and median of medians and does not paralellize well. For GPU environment with CuPy, radixselect would be the most likely option.

• https://docs.scipy.org/doc/numpy/reference/generated/numpy.partition.html
• Functionality similar to std::nth_element? moderngpu/moderngpu#10 (comment)

The following papers mention on radixselect:

• Fast K-selection Algorithms for Graphics Processing Units (http://www.math.grin.edu/~blanchaj/Research/ABGS_KSelection.pdf)

• Radix Selection Algorithm for the kth Order Statistics (http://csis.pace.edu/~ctappert/srd2012/a4.pdf)

The first paper has its public source code and Tensorflow guys have some discussion on it recently.

• https://code.google.com/archive/p/ggks/downloads
• [FEATURE REQUEST]Calculate topK on GPU with k-selection algos tensorflow/tensorflow#5719

[takagi]

I compared the performance of:

• radixselect on CPU
• std::sort & choose on CPU
• thrust::sort & choose on GPU

(Celeron G3900 @ 2.80GHz and GeForce GTX 750 Ti)

Just thrust::sort & choose may be enough at first because its performance fit well to O(n) in time complexity.

https://gist.github.com/sonots/7e11231af2f2cc6b1f519ad832632566

$ nvcc radixselect.cu -O2 -std=c++11 -Wno-deprecated-gpu-targets && ./a.out

----- n = 1000, k = 100
radixselect : 59233 Time elapsed 0.000006 sec
std::sort   : 59233 Time elapsed 0.000043 sec
thrust::sort: 59233 Time elapsed 0.000583 sec

----- n = 10000, k = 1000
radixselect : 59017 Time elapsed 0.000032 sec
std::sort   : 59017 Time elapsed 0.000597 sec
thrust::sort: 59017 Time elapsed 0.000586 sec

----- n = 100000, k = 10000
radixselect : 58992 Time elapsed 0.000347 sec
std::sort   : 58992 Time elapsed 0.006836 sec
thrust::sort: 58992 Time elapsed 0.000436 sec

----- n = 1000000, k = 100000
radixselect : 58953 Time elapsed 0.004255 sec
std::sort   : 58953 Time elapsed 0.074907 sec
thrust::sort: 58953 Time elapsed 0.002840 sec

----- n = 10000000, k = 1000000
radixselect : 58986 Time elapsed 0.039464 sec
std::sort   : 58986 Time elapsed 0.745077 sec
thrust::sort: 58986 Time elapsed 0.023981 sec

----- n = 100000000, k = 10000000
radixselect : 58983 Time elapsed 0.489075 sec
std::sort   : 58983 Time elapsed 7.569387 sec
thrust::sort: 58983 Time elapsed 0.235514 sec

[takagi]

According to Fig. 7 in the paper,

• Fast K-selection Algorithms for Graphics Processing Units (http://www.math.grin.edu/~blanchaj/Research/ABGS_KSelection.pdf)

sort & choose has O(N) complexity in time for vector length N as theoretically shown, and is at least just four times slower than radixselect in case vector length is 2^27.

It seems reasonable to first implement cupy.partition with naive sort & choose.