GitHub

sort and friends

Sorting algorithms with similar interface and default settings as the Julia Base ones, on GPUs:

  • sort! (in-place), sort (out-of-place)
  • sortperm!, sortperm
  • Other names: sort, sort_team, sort_team_by_key, stable_sort or variations in Kokkos, RAJA, Thrust that I know of.

Function signatures:

AcceleratedKernels.sort! — Function
sort!(
    v::AbstractArray, backend::Backend=get_backend(v);

    lt=isless,
    by=identity,
    rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,

    # CPU settings
    max_tasks=Threads.nthreads(),
    min_elems=1,

    # Algorithm choice
    alg::Union{Nothing, SortAlgorithm}=nothing,

    # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector
    dims::Union{Colon, Integer}=Colon(),

    # GPU settings
    block_size::Union{Nothing, Int}=nothing,

    # Temporary buffer, same size as `v`
    temp::Union{Nothing, AbstractArray}=nothing,
)

Sorts the array v in-place using the specified backend. The lt, by, rev, and order arguments are the same as for Base.sort.

By default (dims=:) the whole array is sorted as one vector, whatever its shape. Pass an integer dims to sort each 1D slice along that dimension independently, like Base.sort!(A; dims). On GPU backends this uses merge sort by default, or BitonicSort when requested; RadixSort does not support dims. On CPU backends each slice is sorted with Base.sort! (also with alg=SampleSort(), and without using temp), slices spread over the tasks.

CPU

CPU settings: use at most max_tasks threads to sort the array such that at least min_elems elements are sorted by each thread. A parallel sample sort is used, processing independent slices of the array and deferring to Base.sort! for the final local sorts.

Note that the Base Julia sort! is mainly memory-bound, so multithreaded sorting only becomes faster if it is a more compute-heavy operation to hide memory latency - that includes:

  • Sorting more complex types, e.g. lexicographic sorting of tuples / structs / strings.
  • More complex comparators, e.g. by=custom_complex_function or lt=custom_lt_function.
  • Less cache-predictable data movement, e.g. sortperm.

GPU

GPU settings: block_size sets the number of threads per block. For RadixSort and BitonicSort, fields on the algorithm take precedence over this keyword, then backend defaults. items_per_thread is set on the algorithm and defaults to 2 for RadixSort, 8 for BitonicSort.

Algorithm choice

By default, sort! uses sample sort on CPU backends and merge sort on GPU backends. Pass alg=SampleSort() for the CPU path, alg=MergeSort() for the GPU merge-sort path, alg=RadixSort() to opt into GPU radix sorting, or alg=BitonicSort() for the GPU sorting network. RadixSort() supports 32-bit and 64-bit integers and floats with default lt/by. BitonicSort() is unstable: fastest for small arrays and short slices, slower than MergeSort/RadixSort for large whole-array sorts.

For both CPU and GPU backends, the temp argument can be used to reuse a temporary buffer of the same size as v to store the sorted output.

Examples

Simple parallel CPU sort using all available threads (as given by julia --threads N):

import AcceleratedKernels as AK
v = rand(1000)
AK.sort!(v)

Parallel GPU sorting, passing a temporary buffer to avoid allocating a new one:

using oneAPI
import AcceleratedKernels as AK
v = oneArray(rand(1000))
temp = similar(v)
AK.sort!(v, temp=temp)
source
AcceleratedKernels.sort — Function
sort(
    v::AbstractArray, backend::Backend=get_backend(v);

    lt=isless,
    by=identity,
    rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,

    # CPU settings
    max_tasks=Threads.nthreads(),
    min_elems=1,

    # Algorithm choice
    alg::Union{Nothing, SortAlgorithm}=nothing,

    # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector
    dims::Union{Colon, Integer}=Colon(),

    # GPU settings
    block_size::Union{Nothing, Int}=nothing,

    # Temporary buffer, same size as `v`
    temp::Union{Nothing, AbstractArray}=nothing,
)

Out-of-place sort, same settings as sort!.

source
AcceleratedKernels.sortperm! — Function
sortperm!(
    ix::AbstractArray,
    v::AbstractArray,
    backend::Backend=get_backend(v);

    lt=isless,
    by=identity,
    rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,

    # CPU settings
    max_tasks=Threads.nthreads(),
    min_elems=1,

    # Algorithm choice
    alg::Union{Nothing, SortAlgorithm}=nothing,

    # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector
    dims::Union{Colon, Integer}=Colon(),

    # GPU settings
    block_size::Union{Nothing, Int}=nothing,

    # Temporary buffer, same size as `v`
    temp::Union{Nothing, AbstractArray}=nothing,
)

Save into ix the index permutation of v such that v[ix] is sorted. The lt, by, rev, and order arguments are the same as for Base.sortperm. The same algorithms are used as for sort! with custom by-index comparators.

By default (dims=:) the whole array is permuted as one vector. Pass an integer dims to permute each 1D slice along that dimension independently, like Base.sortperm!(ix, A; dims): ix must then have the same axes as v and receives linear indices into v, so that v[ix] is sorted along dims. The permutation is stable in both cases.

Algorithm choice

By default, sortperm! uses sample sort on CPU backends and merge sort on GPU backends. Pass alg=MergeSort(lowmem=true) to use the lower-memory GPU permutation path. RadixSort() and BitonicSort() do not provide a permutation path.

source
AcceleratedKernels.sortperm — Function
sortperm(
    v::AbstractArray,
    backend::Backend=get_backend(v);

    lt=isless,
    by=identity,
    rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,

    # CPU settings
    max_tasks=Threads.nthreads(),
    min_elems=1,

    # Algorithm choice
    alg::Union{Nothing, SortAlgorithm}=nothing,

    # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector
    dims::Union{Colon, Integer}=Colon(),

    # GPU settings
    block_size::Union{Nothing, Int}=nothing,

    # Temporary buffer, same size as `v`
    temp::Union{Nothing, AbstractArray}=nothing,
)

Out-of-place sortperm, same settings as sortperm!.

source

Algorithm choice is available on sort! / sort / sortperm! / sortperm with alg=AK.MergeSort(), alg=AK.MergeSort(lowmem=true), alg=AK.RadixSort(), alg=AK.BitonicSort(), or alg=AK.SampleSort(), depending on the backend and operation.

Function signatures:

AcceleratedKernels.MergeSort — Type
MergeSort(; lowmem=false)

Use GPU merge sort for sort! and sort. For sortperm!, lowmem=true selects the lower-memory permutation path.

source
AcceleratedKernels.RadixSort — Type
RadixSort(; block_size=nothing, items_per_thread=nothing)

Use GPU radix sort for sort! and sort. Supports UInt32, Int32, Float32, UInt64, Int64, and Float64 with forward or reverse ordering. This algorithm does not support sortperm!.

source
AcceleratedKernels.BitonicSort — Type
BitonicSort(; block_size=nothing, items_per_thread=nothing)

Use GPU bitonic sort for sort! and sort, whole-array or along dims. Supports GPU-compatible element types and lt/by/rev/order; by is evaluated at each comparison. The sort is unstable and does not support sortperm! or sortperm.

Fastest for small arrays and short slices, slower than MergeSort/RadixSort for large whole-array sorts. Tiles of up to block_size * items_per_thread elements sort in local memory; larger inputs need global passes.

Both settings must be positive powers of two and default to bitonic_defaults(backend). For block_size, the sort! keyword takes precedence over the backend default.

source

Example:

import AcceleratedKernels as AK
using AMDGPU

v = ROCArray(rand(Int32, 100_000))
AK.sort!(v)

Multidimensional arrays are sorted as one flat vector by default; pass dims to sort each 1D slice along that dimension independently, like Base.sort!(A; dims). sortperm along dims returns linear indices into the array, so A[ix] is sorted along dims:

A = ROCArray(rand(Float32, 1000, 1000))
AK.sort!(A; dims=1)             # each column sorted
ix = AK.sortperm(A; dims=2)     # A[ix] has each row sorted

On GPU backends dims uses merge sort by default (RadixSort() does not support it); on CPU backends each slice is sorted with Base.sort!. BitonicSort() is unstable and supports sort! and sort, including dims, but not sortperm! or sortperm. It is fastest for small arrays and short slices, slower than MergeSort/RadixSort for large whole-array sorts:

A = ROCArray(rand(Float32, 64, 100_000))
AK.sort!(A; dims=1, alg=AK.BitonicSort())

As GPU memory is more expensive, all functions in AcceleratedKernels.jl expose any temporary arrays they will use (the temp argument); you can supply your own buffers to make the algorithms not allocate additional GPU storage, e.g.:

v = ROCArray(rand(Float32, 100_000))
temp = similar(v)
AK.sort!(v, temp=temp)