sort!(
    v::AbstractArray;
    backend=nothing,
    alg::Algorithm=Auto(),
    dims::Union{Colon, Integer}=:,
    lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,
    workspace=nothing,
) -> v

Sort v in place. lt, by, rev and order are those of Base.sort!, and so is the result: by default the sort is stable. dims=: sorts the whole array as one vector (Base.sort! requires dims for arrays of more than one dimension); an integer dims sorts each slice along that dimension independently.

alg is Auto() by default, which chooses an algorithm for the backend and device, the element type, the ordering, and the length of the array or slices: on the host, CPUThreads.SampleSort; on GPUs, MergeSort, or BitonicSort for short inputs or RadixSort for long ones where the device's tuning enables them and they give the same result. Auto(stable=false) also allows the unstable BitonicSort for elements that compare equal without being identical (e.g. floats). An explicit algorithm, with any of its settings, is used as given or rejected with an ArgumentError.

backend is derived from v; pass it only for arrays that do not determine their backend.

workspace takes the scratch memory of a workspace made for the same call, so that the sort allocates none of its own.

Examples

import AcceleratedKernels as AK
using CUDA

v = CuArray(rand(Float32, 100_000))
AK.sort!(v)                                         # Auto: stable, chosen for the device
AK.sort!(v; rev=true, alg=AK.Auto(stable=false))    # may use an unstable algorithm
AK.sort!(v; alg=AK.RadixSort(block_size=512))       # this algorithm, with this setting

A = CuArray(rand(Int32, 64, 10_000))
AK.sort!(A; dims=1)                                 # each column

AK.sort!(rand(1000))                                # host array: Julia threads

sort(v::AbstractArray; kwargs...)

Out-of-place sort!: sort a copy of v, with the same keywords.

sortperm!(
    ix::AbstractArray,
    v::AbstractArray;
    backend=nothing,
    alg::Algorithm=Auto(),
    dims::Union{Colon, Integer}=:,
    lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,
    workspace=nothing,
) -> ix

Write into ix the stable permutation that sorts v, so that v[ix] is sorted; ix is always overwritten. The keywords are those of sort!.

With dims=:, ix needs as many elements as v and receives indices 1:length(v). With an integer dims, ix must have the same axes as v and receives linear indices into v, so that v[ix] is sorted along dims, like Base.sortperm!(ix, A; dims).

Auto() chooses CPUThreads.SampleSort on the host and MergeSort on GPUs; MergeSort(lowmem=true) avoids copying the keys. RadixSort and BitonicSort have no permutation path. backend is derived from ix and v.

sortperm(v::AbstractArray; kwargs...)

Out-of-place sortperm!: return an Int array shaped like v holding the permutation, with the same keywords.

sort_by_key!(
    keys::AbstractArray,
    values::AbstractArray;
    backend=nothing,
    alg::Algorithm=Auto(),
    dims::Union{Colon, Integer}=:,
    lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing,
    order::Base.Order.Ordering=Base.Order.Forward,
    workspace=nothing,
) -> (keys, values)

Sort keys in place and apply the same permutation to values, stably: values with equal keys keep their relative order. The ordering keywords and dims are those of sort!; values must have as many elements as keys (the same axes with an integer dims).

Auto() chooses CPUThreads.SampleSort on the host and MergeSort on GPUs, the algorithms that support key/value sorting. backend is derived from keys and values.

Thrust, oneDPL and Kokkos call this operation sort_by_key, CUB SortPairs.

Examples

import AcceleratedKernels as AK
using Metal

keys = MtlArray(rand(Int32(1):Int32(10), 1000))
values = MtlArray(Int32.(1:1000))
AK.sort_by_key!(keys, values)      # values of equal keys stay in ascending order

MergeSort(; block_size=nothing, lowmem=false)

GPU merge sort: stable, for every element type and ordering, whole arrays and dims, and the only kernel algorithm for sortperm! and sort_by_key!. Each block of block_size threads (any positive number) sorts a tile of 2 * block_size elements in local memory, then global passes merge the tiles. lowmem=true selects a sortperm! path that does not copy the keys, at the cost of reading them from global memory in every comparison; other operations reject it.

RadixSort(; block_size=nothing, items_per_thread=nothing, chunk_size=nothing)

GPU LSD radix sort for 32- and 64-bit integers and floats (UInt32, Int32, Float32, UInt64, Int64, Float64) under the default ordering or its reverse, for whole arrays and dims; it is stable and orders floats like isless. It does not support custom lt/by, sortperm! or sort_by_key!, nor arrays or slices of more than 2^32 elements. block_size must be a power of two up to 1024, items_per_thread between 1 and 64. items_per_thread and chunk_size apply where radix sort's chunked kernels run, which needs atomics, a block_size that is a multiple of chunk_size, and tiles that fit local memory; elsewhere its portable kernels sort one item per thread. The chunked kernels rank each element among the earlier elements of its chunk: chunk_size, a power of two up to 1024, trades those comparisons against the chunk histograms in local memory.

Along dims, each pass sorts all slices at once. It pays off for long slices: every slice takes at least one block of the passes, so short ones leave most threads idle, and slices along other dimensions than the first are read with a stride. Where the chunked kernels run and a slice of up to 2 * block_size elements fits in local memory, one block sorts each slice there instead.

BitonicSort(; block_size=nothing, items_per_thread=nothing)

GPU bitonic sorting network, for whole arrays and dims, every element type and ordering. It is unstable and does not support sortperm! or sort_by_key!. Fastest for small arrays and short slices: tiles of block_size * items_per_thread elements sort in local memory, larger inputs need global passes. Both settings must be positive powers of two.

CPUThreads.SampleSort(; max_tasks=nothing, min_elems=nothing)

Parallel sample sort on Julia threads, deferring to Base.sort! for the local sorts; it is stable, and also provides sortperm! and sort_by_key!. Uses at most max_tasks tasks (default Threads.nthreads()), each with at least min_elems elements (default 1). Only runs on the host backend.