Sortperm is slow (then also sortslices and sortperm); Julia's oldest performance issue(s)

General Usage Performance
1

Continuing the discussion from 10x faster sortperm():

Sortperm is at least 77.6% slower in Julia than in Python/NumPy.

Note, this is known, the next oldest performance issue (and third oldest issue from 2013):

[It’s similar in speed for Float64 and Int64 as expected, but it doesn’t go much down with smaller types down to Int8 (nor much with NumPy), a bit surprising, because of radix sort?]

We are not (no longer?) “3x slower for sortperm” (as once for MATLAB in 2013, but I didn’t test it, assume it similar speed to NumPy).

I thought AI might be able to help (I could share a link to my answer), even fix if someone has access to more powerful AI than mine…

AI though SortingLab would help but it didn’t maybe in the past:

using BenchmarkTools
using SortingLab

# Generate 10 million random Float64 values
const data = rand(10_000_000)

println("--- Base Julia sortperm ---")
@btime sortperm($data);

println("--- SortingLab fsortperm (Fast Radix) ---")
@btime fsortperm($data);


julia> @btime sortperm($data);
  1.980 s (6 allocations: 152.59 MiB)

--- SortingLab fsortperm (Fast Radix) ---
  1.984 s (6 allocations: 152.59 MiB

vs. Python:

import numpy as np
import timeit

# Generate 10 million random Float64 values
data = np.random.rand(10000000)

# Benchmark NumPy's argsort (defaults to Quicksort, but supports Radix/Mergesort)
t = timeit.timeit(lambda: np.argsort(data), number=3)
print(f"NumPy argsort average time: {t / 3:.4f} seconds")

NumPy argsort average time: 1.1151 seconds
NumPy argsort average time: 1.0962 seconds

Interestingly with as many integers (this time tested with MORE load, i.e. after starting Firefox again):

rng = np.random.default_rng()
array_1d = rng.integers(0, 100, size=10000_000)

NumPy argsort average time: 0.5763 seconds

julia> @btime sortperm($data);
1.903 s (6 allocations: 152.59 MiB)

julia> @btime fsortperm($data);
1.911 s (6 allocations: 152.59 MiB)


I haven't tested:
https://github.com/LSchwerdt/SimultaneousSortperm.jl/pull/3

The oldest (if you want to look at or "fix" the oldest and only other 3-digit performance issue):
https://github.com/JuliaLang/julia/issues/130#issuecomment-1123823801
2

There’s something else going on:

julia> const data = rand(Int,10_000_000);
julia> const shortdata = rand(Int8, 10_000_000);
julia> @btime sort(data);
  237.019 ms (9 allocations: 152.60 MiB)
julia> @btime sort(shortdata);
  4.207 ms (5 allocations: 9.54 MiB)

The sortperm uses sort, like this (though, inside an Ordering inside a Perm) :

ix = Base.copymutable(LinearIndices(data))
sort!(ix, by=i -> data[i])

So, sortperm shouldn’t be much slower, at least not when naively counting operations.

julia> @btime sortperm(data);
  1.290 s (6 allocations: 152.59 MiB)

julia> @btime sortperm(shortdata);
  576.534 ms (6 allocations: 152.59 MiB)

julia> @btime sort!(ix, by=i -> data[i]) setup=(ix = Base.copymutable(LinearIndices(data)));
  1.029 s (3 allocations: 76.30 MiB)

julia> @btime sort!(ix, by=i -> shortdata[i]) setup=(ix = Base.copymutable(LinearIndices(shortdata)));
  456.604 ms (3 allocations: 76.30 MiB)

I guess that sort is specialized on the by or lt argument deep down there for the most common cases, whereas sortperm with its indirect loookup is not:

julia> @btime sort(data, by=fun) setup=(fun = identity);
  234.166 ms (9 allocations: 152.60 MiB)

julia> @btime sort(shortdata, by=fun) setup=(fun = identity);
  4.216 ms (5 allocations: 9.54 MiB)

julia> @btime sort(data, by=fun) setup=(fun = x->x);
  358.734 ms (6 allocations: 152.59 MiB)

julia> @btime sort(shortdata, by=fun) setup=(fun = x->x);
  210.548 ms (6 allocations: 19.07 MiB)

(Btw, python argsort with float64 runs on 0.51s on this machine).

4

*numpy argsort, presumably.

5

sortperm uses sort with by=i->data[i]. This means that comparison costs a spatially non-local memory access :scream:

edit: In-place sort of plain data avoids the spatially non-local memory accesses by rearranging / copying around data during the sorting process. The early comparisons are naturally local; the later ones are local because the early phases of sorting already copied data around.

You can avoid that by

julia> using BenchmarkTools

julia> function sortpermCopy(data)
       zipped = [(data[i], i) for i=1:length(data)]
       sort!(zipped)
       [idx for (_, idx) in zipped]
       end;

julia> r0=rand(Int,10_000_000); s=sortperm(r0); ss = sortpermCopy(r0); isequal(s, ss)
true

julia> @btime sortperm($r0);
  1.007 s (6 allocations: 152.59 MiB)

julia> @btime sortpermCopy($r0);
  414.261 ms (9 allocations: 381.48 MiB)

I thought this was common knowledge? Maybe we should add documentation for that kind of thing.

I understand why the standard library makes the decisions it makes: sortperm needs to work even if the data cannot be copied (array of large keys that are shared read-only memory-mapped into many different processes).

PS.

julia> using StructArrays;

julia> function sortperm2(data)
       s=StructArray((Base.copymutable(data), collect(1:length(data))))
       sort!(s; by=x->x[1])
       s.:2
       end;

julia> @btime sortperm2($r0);
  288.576 ms (9 allocations: 305.18 MiB)
6

And, since allocating is devastating for parallel performance, there should be a sortperm! and a documented way to supply the scratch memory.

7

I think a good API/implementation would be

"""
sort2 sorts `data` in-place, and applies the same permutation to `extra`.
"""
function sort2!(data, extra)
       s=StructArray((data, extra))
       sort!(s;by=x->x[1])
       nothing
       end

An issue is with keyword arguments that need to be forwarded to sort!, and need to be appropriately composed for Ordering. I think that can be done, but slightly painfully.

That is pretty much allocation-free (ok, it allocates O(1), ~ 16 bytes). And this is a very natural operation that I need often: Sort something in-place, and apply the same permutation also to something else.

8

Does this for sure need a new API, or if not, a performance fix can happen in 14.x.

I wouldn’t know. I see sortperm is used in a lot of packages. But your alternative versions reimplement it, using sort!, so I suppose a lot of users don’t know how to do such, and just rely on sortperm as is.

I note while both your versions are faster, they allocate more memory, so is that needed, and might NumPy do such? I don’t know of a similar way to benchmark Python/NumPy as with @btime and get such memory info (easily, or well at all…).

This seems to be a a memory-speed tradeoff, or could it be done in some way to use more memory up to a point then switch to the other method, or a hybrid?

I noticed there is actually sortperm! available, and maybe I should have used it and benchmarked (I was just testing same as in the Julialang issue). Isn’t it exactly what you mean? But while some packages use it, many just use sortperm.

Searching either function shows “500 results found” so can’t be right.

What did you mean by “allocating is devastating for parallel performance”? sortperm! would fix that, but why would it be devastating?

It seems like a quite capable package:

Auto() chooses [CPUThreads.SampleSort](@ref AcceleratedKernels.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.
merge_sortperm_lowmem!

> function _sortperm_impl!(a::MergeSort, ix, v, backend, dims, ord, bufs; lt, by, rev, order)
>     if a.lowmem
>         _merge_sortperm_lowmem!(ix, v, backend, bufs; lt, by, rev, order,
>                                 block_size=a.block_size, dims)
>     else
>         # Copies keys alongside indices, so comparisons never read global memory; the low-memory
>         # path does two global loads per comparison, O(n log²n) global traffic at large n.
>         _merge_sortperm!(ix, v, backend, bufs; lt, by, rev, order, block_size=a.block_size, dims)
>     end
> end

MergeSort(; block_size=nothing, lowmem=false); lowmem seem to be only what you opt into, not what is auto-detected for you.

I’m a bit confused, I see merge_sortperm!(ix, v1) used in the package (and AK.merge_sortperm2!(ix, v1)) but not defined. There is however function _merge_sortperm!

There is benchmark code:

display(@benchmark Base.sortperm!(ix, v) setup=(v = rand(UInt32(0):UInt32(1_000_000), 1_000_000); ix = Vector{Int}(undef, 1_000_000)))
display(@benchmark AK.sortperm!(ix, v) setup=(v = rand(UInt32(0):UInt32(1_000_000), 1_000_000); ix = Vector{Int}(undef, 1_000_000)))
9

Ah, yes, I was a bit quick there. There’s a sortperm!, but it does not support the undocumented keyword parameter scratch to sort! which can be used to supply preallocated scratch memory.
EDIT: Hmm, it seems it does, perhaps I did something wrong when I tried it.

10

More crucially, it does not address the performance issue: Out-of-place sorting is fundamentally slower than in-place sorting for small eltypes, due to spatial memory locality.

If you want the thing to be fast, you need to sort in-place. In many cases this is what the user wants anyway (i.e. “sort data, record the permutation (i.e. apply the permutation to collect(1:length(data)))”). In these cases, we simply lack a good API function.

In many cases, the fastest way is to create a copy, and sort that in-place, while recording the permutation, and then discarding the now-sorted copy. That’s sortperm2.

However, that takes temp-memory length(data)*sizeof(eltype(data)), which might be much much more than the naive algorithm length(data)*8. Therefore sortperm2 cannot be the only supported behavior. (I guess one could do something like: If sizeof(eltype(data)) <= 32, take a temp-copy that is sorted in-place, otherwise sort out-of-place).

11

There is some funny inplace sortperm! that one can use for instance for v::Vector{Int} assuming (roughly) maximum(abs,v) * length(v) < typemax(Int). It works for other Number subtypes, though the criteria to avoid overflows can be a bit more intricate:

function unsafe_sortperm!(v::Vector{Int})
    ceil_log2 = 64 - leading_zeros(length(v)-1)
    v .= (v .<< ceil_log2) .+ eachindex(v)
    sort!(v)
    v .= (v .- 1) .& ((1 << ceil_log2) - 1) .+ 1
end

function test_unsafe_sortperm!(n)
    v = rand(1:1000000, n)
    @time unsafe_sortperm!(v)
    w = rand(1:1000000, n)
    @time sortperm!(v,w)
    return nothing
end

test_unsafe_sortperm!(100000)
# 0.002737 seconds (6 allocations: 789.453 KiB)
# 0.011198 seconds (3 allocations: 781.320 KiB)

As a bonus, it is stable no matter which algorithm one may choose for the sort! call.