GPU Implementation of partialsort!

[btmit]

It would be great if there was an implementation of sort! and partialsort! for CuArrays.

A simple (1D array, radix-2) native Julia implementation would be really helpful as a good starting point.

Other issues have considered wrapping Thrust or CUB JuliaGPU/CuArrays.jl#525.

This should enable a number of higher level functions, such as the ability to generate statistics about data (median, quartiles, etc) on the GPU.

A few (possibly) helpful links:
https://www.cse.chalmers.se/~tsigas/papers/GPU-Quicksort-jea.pdf
https://onlinelibrary.wiley.com/doi/full/10.1002/cpe.3611
https://sourceforge.net/projects/cuda-quicksort/

[wongalvis14]

I looked into a couple parallel sorting algorithms each with their own setbacks:

Bitonic sort - for arrays with lengths not a power of 2, we need to fill in dummy values to make it so, the asymptotic complexity is the same, but the overhead and extra memory will be large;
or we can devise a non-recursive version that works based on https://www.inf.hs-flensburg.de/lang/algorithmen/sortieren/bitonic/oddn.htm, if an efficient solution exists

Merge sort - extra memory, more work for each thread (binary search) on the merge step, special handling for arrays with lengths not a power of 2

Radix sort - handling for negative numbers

Quick sort - bad worst case complexity

[maleadt]

I had tried a quicksort implementation to test dynamic parallelism, but it's slow: https://gist.github.com/maleadt/08bd84906382ed8b9c57e154356f6a0b

[wongalvis14]

CUDA's quicksort examples use dynamic parallelism, my experience with dynamic parallelism was very... not ideal... but I did come across a paper recently that showed a speedup of dynamic parallelism over device-host sync+dispatch

[wongalvis14]

I benchmarked my implementation of bitonic sort and @maleadt 's quicksort

It seems that bitonic sort is a better baseline to build upon?

julia> @benchmark test_bitonic_sort(1000000)
BenchmarkTools.Trial: 
  memory estimate:  253.53 KiB
  allocs estimate:  8858
  --------------
  minimum time:     2.811 ms (0.00% GC)
  median time:      29.530 ms (0.00% GC)
  mean time:        29.380 ms (0.00% GC)
  maximum time:     30.135 ms (0.00% GC)
  --------------
  samples:          171
  evals/sample:     1

julia> @benchmark test_quicksort(1000000)
BenchmarkTools.Trial: 
  memory estimate:  1.53 KiB
  allocs estimate:  57
  --------------
  minimum time:     995.534 ms (0.00% GC)
  median time:      1.536 s (0.00% GC)
  mean time:        1.578 s (0.00% GC)
  maximum time:     2.243 s (0.00% GC)
  --------------
  samples:          4
  evals/sample:     1

julia> @benchmark test_quicksort(100000)
BenchmarkTools.Trial: 
  memory estimate:  1.53 KiB
  allocs estimate:  57
  --------------
  minimum time:     202.019 μs (0.00% GC)
  median time:      43.582 ms (0.00% GC)
  mean time:        52.244 ms (0.00% GC)
  maximum time:     148.446 ms (0.00% GC)
  --------------
  samples:          96
  evals/sample:     1

[maleadt]

It seems that bitonic sort is a better baseline to build upon?

Definitely, but the pow2 requirement is problematic. Padding the array is a no-go, and using sentinel values to denote out-of-bounds elements is generally not possible either.

[wongalvis14]

For comparison, this is the standard Julia CPU sort

julia> @benchmark test_sort_cpu(1000000)
BenchmarkTools.Trial: 
  memory estimate:  15.26 MiB
  allocs estimate:  4
  --------------
  minimum time:     131.498 ms (0.00% GC)
  median time:      134.018 ms (0.00% GC)
  mean time:        136.183 ms (1.00% GC)
  maximum time:     150.196 ms (1.20% GC)
  --------------
  samples:          37
  evals/sample:     1

the pow2 requirement is problematic

Indeed, in the beginning I was in search of a generic solution for all array lengths. Until I saw NVIDIA's implementation, in which they added a memory buffer (I learnt it from them). I realized this is probably the "best solution" so far. On the other hand, I think a lot of times we sort arrays that are either pow2 or pow10 which would be relatively close to a pow2 anyway.

[xaellison]

I have an implementation of quicksort which may be of interest. It uses dynamic parallelism and gets pretty good thread usage when partitioning. On my machine, for random input, it slightly outperforms sorting a CuArray compared to sorting the equivalent Array on the host (so it would really beat transferring to host, sorting, and transferring back).

Using @maleadt 's implementation above for comparison

function main(T, N)
    b = @benchmarkable (CUDA.@sync(sort!(c))) setup=(c = CuArray(rand($T, $N))) samples=7 seconds=30
    run(b)
end

main(Float64, 1000000)

gives

BenchmarkTools.Trial: 
  memory estimate:  544 bytes
  allocs estimate:  24
  --------------
  minimum time:     1.740 s (0.00% GC)
  median time:      3.407 s (0.00% GC)
  mean time:        4.920 s (0.00% GC)
  maximum time:     12.968 s (0.00% GC)
  --------------
  samples:          7
  evals/sample:     1

The same test on my quicksort is about roughly 50x faster

julia> speed_test(Float64, Int(1e6))
BenchmarkTools.Trial: 
  memory estimate:  720 bytes
  allocs estimate:  32
  --------------
  minimum time:     62.126 ms (0.00% GC)
  median time:      63.724 ms (0.00% GC)
  mean time:        67.051 ms (0.00% GC)
  maximum time:     82.177 ms (0.00% GC)
  --------------
  samples:          10
  evals/sample:     1

and it scales well

julia> speed_test(Float64, Int(1e8))
BenchmarkTools.Trial: 
  memory estimate:  720 bytes
  allocs estimate:  32
  --------------
  minimum time:     7.658 s (0.00% GC)
  median time:      7.663 s (0.00% GC)
  mean time:        7.663 s (0.00% GC)
  maximum time:     7.669 s (0.00% GC)
  --------------
  samples:          4
  evals/sample:     1

One of the last things I want to improve is make it handle the worst case of all elements are equal. This should be feasible without killing performance.

It has plenty of room to improve but I think it could be a solid starting point. Comparing to bullet points above:

• It can do negative numbers
• No padding the length
• No extra memory
• Worst case can be well controlled by the strategy linked above, plus median is drawn from a pool of 1024.

Happy to open a branch if you want to consider it!

[maleadt]

That's great! Please do open a branch, even if it's an initial version. We can still do development there.

Going back to the bitonic timings from @wongalvis14, which are another order of magnitude better, maybe we should dispatch at run-time based on the input sizes.

[wongalvis14]

@xaellison I'm curious what's your approach to efficiently use dynamic parallelism. In my experience with CUDA it introduces way too much overhead during run-time that turns any efficient algorithm inefficient.

[xaellison]

@maleadt was about to commit when Github told me I need to be authorized if you don't mind setting me up.

@wongalvis14 Each quicksort kernel only launches one kernel (partition) that it needs to wait on. Partition happens in default stream (no stream arg is provided). Recursive quicksort kernels are launched on their own streams but aren't explicitly waited on. A lot of work happens in the quicksort kernel: selecting a pivot, half the partitioning logic, and sorting small sublists. That allows me to launch fewer kernels per recursive step. I hope that the performance I'm getting won't be limited to my hardware.

[vchuravy]

was about to commit when Github told me I need to be authorized if you don't mind setting me up.

You will need to fork the project, into your own account and push to there. Then you can open a PR.