Skip to content

Replace dynamic dispatch with runtime branch on rev keyword in sortperm #47966

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KristofferC merged 9 commits into from
Dec 28, 2022
Merged

Replace dynamic dispatch with runtime branch on rev keyword in sortperm #47966

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
70d9aef
add strange hacks to avoid compiler heuristics?
Dec 21, 2022
62a5729
add allocation tests
Dec 21, 2022
3ae8d9c
elimiate allocations in the rev::Bool case as well
Dec 22, 2022
f06d197
compile before allocation test
Dec 22, 2022
4191e27
fixup: fix typo in added tests
Dec 22, 2022
16136d0
Merge branch 'master' into sortperm-allocs
LilithHafner Dec 27, 2022
9616d05
Remove the hacks
LilithHafner Dec 27, 2022
da6d39d
revert more superfluous changes
Dec 27, 2022
2a98fbb
fixup & put back `@inline` because it is sometimes necessary for perf…
Dec 27, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 26 additions & 8 deletions base/sort.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1560,8 +1560,14 @@ function sortperm(A::AbstractArray;
order::Ordering=Forward,
scratch::Union{Vector{<:Integer}, Nothing}=nothing,
dims...) #to optionally specify dims argument
ordr = ord(lt,by,rev,order)
if ordr === Forward && isa(A,Vector) && eltype(A)<:Integer
if rev === true
_sortperm(A; alg, order=ord(lt, by, true, order), scratch, dims...)
else
_sortperm(A; alg, order=ord(lt, by, nothing, order), scratch, dims...)
end
end
function _sortperm(A::AbstractArray; alg, order, scratch, dims...)
if order === Forward && isa(A,Vector) && eltype(A)<:Integer
n = length(A)
if n > 1
min, max = extrema(A)
Expand All @@ -1572,8 +1578,8 @@ function sortperm(A::AbstractArray;
end
end
end
ix = copymutable(LinearIndices(A))
sort!(ix; alg, order = Perm(ordr, vec(A)), scratch, dims...)
ix = unsafe_copymutable_LinearIndices(A)
sort!(ix; alg, order = Perm(order, vec(A)), scratch, dims...)
end


Expand Down Expand Up @@ -1615,7 +1621,7 @@ julia> sortperm!(p, A; dims=2); p
2 4
```
"""
function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
@inline function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
alg::Algorithm=DEFAULT_UNSTABLE,
lt=isless,
by=identity,
Expand All @@ -1625,14 +1631,26 @@ function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
scratch::Union{Vector{T}, Nothing}=nothing,
dims...) where T <: Integer #to optionally specify dims argument
(typeof(A) <: AbstractVector) == (:dims in keys(dims)) && throw(ArgumentError("Dims argument incorrect for type $(typeof(A))"))
axes(ix) == axes(A) || throw(ArgumentError("index array must have the same size/axes as the source array, $(axes(ix)) != $(axes(A))"))
axes(ix) == axes(A) || @noinline throw_sortperm_axis_mismatch_error(ix, A)

if !initialized
ix .= LinearIndices(A)
very_unsafe_copyto!(ix, LinearIndices(A))
end

if rev === true
sort!(ix; alg, order=Perm(ord(lt, by, true, order), vec(A)), scratch, dims...)
else
sort!(ix; alg, order=Perm(ord(lt, by, nothing, order), vec(A)), scratch, dims...)
end
sort!(ix; alg, order = Perm(ord(lt, by, rev, order), vec(A)), scratch, dims...)
end

# TODO stop using these three hacks
# but check performance, especially unexpected allocations, when removing
Base.@assume_effects :nothrow very_unsafe_copyto!(a, b) = a .= b
Copy link
Member

@oscardssmith oscardssmith Dec 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is scary

Copy link
Member Author

@LilithHafner LilithHafner Dec 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah... @aviatesk, is this sort of thing acceptable or is there another way?

Copy link
Member

@N5N3 N5N3 Dec 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My local test shows that perhaps @noinline is enough?

Copy link
Member Author

@LilithHafner LilithHafner Dec 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is why we do code review. Thanks!

I think I introduced these when they were necessary and then added branching on rev which is more impactful and fixes the root of the issue so now none of these hacks are necessary at all, and CI checks for unexpected allocations automatically :)

Base.@assume_effects :nothrow unsafe_copymutable_LinearIndices(A) = copymutable(LinearIndices(A))
throw_sortperm_axis_mismatch_error(ix, A) =
throw(ArgumentError("index array must have the same size/axes as the source array, $(axes(ix)) != $(axes(A))"))

# sortperm for vectors of few unique integers
function sortperm_int_range(x::Vector{<:Integer}, rangelen, minval)
offs = 1 - minval
Expand Down
22 changes: 22 additions & 0 deletions test/sorting.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -918,6 +918,28 @@ end
@test bsqs() === bsqs(missing, missing, InsertionSort)
end

function test_allocs()
v = rand(10)
i = randperm(length(v))
@test 1 == @allocations sort(v)
@test 0 == @allocations sortperm!(i, v)
@test 0 == @allocations sort!(i)
@test 0 == @allocations sortperm!(i, v, rev=true)
@test 1 == @allocations sortperm(v, rev=true)
@test 1 == @allocations sortperm(v, rev=false)
@test 0 == @allocations sortperm!(i, v, order=Base.Reverse)
@test 1 == @allocations sortperm(v)
@test 1 == @allocations sortperm(i, by=sqrt)
@test 0 == @allocations sort!(v, lt=(a, b) -> hash(a) < hash(b))
sort!(Int[], rev=false) # compile
@test 0 == @allocations sort!(i, rev=false)
rand!(i)
@test 0 == @allocations sort!(i, order=Base.Reverse)
end
@testset "Small calls do not unnecessarily allocate" begin
test_allocs()
end

# This testset is at the end of the file because it is slow.
@testset "searchsorted" begin
numTypes = [ Int8, Int16, Int32, Int64, Int128,
Expand Down