diff --git a/Project.toml b/Project.toml index 8748f219..361d4041 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GPUArrays" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "12.2.0" +version = "12.3.0" [workspace] projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"] @@ -29,7 +29,7 @@ JLD2Ext = "JLD2" [compat] AcceleratedKernels = "0.5.8" -Adapt = "4.7.3" +Adapt = "4.8" GPUArraysCore = "= 0.2.2" JLD2 = "0.4, 0.5, 0.6" KernelAbstractions = "0.10" diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index 3db7a050..952b8f4e 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -31,6 +31,7 @@ include("host/construction.jl") include("host/base.jl") include("host/indexing.jl") include("host/sorting.jl") +include("host/set.jl") include("host/accumulate.jl") include("host/broadcast.jl") include("host/mapreduce.jl") diff --git a/src/host/set.jl b/src/host/set.jl new file mode 100644 index 00000000..df5972f2 --- /dev/null +++ b/src/host/set.jl @@ -0,0 +1,193 @@ +# unique elements and set operations, by sorting +# +# Sorting puts elements that are `isequal` next to each other, as long as the order agrees with +# `isequal`. `isless` does, and so does `isless` on the real and imaginary parts of complex +# numbers. To return elements in the order Base does, the first element of each group is marked +# at its original position, and the marked elements are selected with a mask. + +# Mixed with complex numbers, a real number sorts as a complex number with zero imaginary part. +complex_key(x::Complex) = reim(x) +complex_key(x) = (x, zero(x)) +complex_isless(x, y) = isless(complex_key(x), complex_key(y)) + +has_complex(T) = T isa Union ? any(has_complex, Base.uniontypes(T)) : T <: Complex + +# Plain `isless` where possible, because AcceleratedKernels can only use its faster sorting +# algorithms with the default order. +set_isless(::Type{T}) where {T} = has_complex(T) ? complex_isless : isless + +# Returns the sorting permutation of `keys` and the sorted keys. The sort is stable, so the +# first element of each group of equal keys is the one that occurs first in `keys`. +function sorted_keys(keys::AbstractVector) + lt = set_isless(eltype(keys)) + p = sortperm(keys; lt) + return p, keys[p] +end + +# For each element of the sorted `s`, whether it differs from the one before it +function group_starts(s::AbstractVector) + n = length(s) + head = similar(s, Bool, n) + n == 0 && return head + head[1:1] .= true + head[2:n] .= .!isequal.(view(s, 2:n), view(s, 1:n-1)) + return head +end + +# Moves the flags `heads`, which are in sorted order, back to the original positions +function scatter_mask(p::AbstractVector, heads::AbstractVector{Bool}) + mask = similar(heads) + mask[p] = heads + return mask +end + +function Base.unique(f, A::AnyGPUArray) + v = vec(A) + length(v) <= 1 && return copy(v) + p, s = sorted_keys(f === identity ? v : map(f, v)) + return v[scatter_mask(p, group_starts(s))] +end + +function Base.unique(A::AnyGPUArray; dims::Union{Colon, Integer}=:) + dims isa Colon && return unique(identity, A) + # Base returns a copy when `dims` is out of range + 1 <= dims <= ndims(A) || return copy(A) + return unique_dims(A, Int(dims)) +end + +# The slices along `d` are compared in a 3-d reshape of `A`, whose middle dimension is `d` +function unique_dims(A::AbstractArray, d::Int) + n = size(A, d) + n <= 1 && return copy(A) + B = reshape(A, prod(size(A)[1:d-1]), n, prod(size(A)[d+1:end])) + ltel = set_isless(eltype(A)) + lt = let B = B, ltel = ltel + (i, j) -> begin + @inbounds for k in axes(B, 3), l in axes(B, 1) + a, b = B[l, i, k], B[l, j, k] + ltel(a, b) && return true + ltel(b, a) && return false + end + return false + end + end + eq = let B = B + (i, j) -> begin + @inbounds for k in axes(B, 3), l in axes(B, 1) + isequal(B[l, i, k], B[l, j, k]) || return false + end + return true + end + end + p = similar(A, Int, n) + p .= 1:n + sort!(p; lt) + head = similar(A, Bool, n) + head[1:1] .= true + head[2:n] .= .!eq.(view(p, 1:n-1), view(p, 2:n)) + mask = scatter_mask(p, head) + return A[ntuple(i -> i == d ? mask : Colon(), ndims(A))...] +end + +function Base.unique!(f, v::AnyGPUVector) + u = unique(f, v) + copyto!(v, firstindex(v), u, 1, length(u)) + return resize!(v, length(u)) +end +Base.unique!(v::AnyGPUVector) = unique!(identity, v) + +function Base.allunique(A::AnyGPUArray) + n = length(A) + n <= 1 && return true + lt = set_isless(eltype(A)) + s = sort(vec(A); lt) + return !any(isequal.(view(s, 2:n), view(s, 1:n-1))) +end +# Base has `allunique(f, A)` since Julia 1.11 +if hasmethod(allunique, Tuple{Any, Any}) + Base.allunique(f, A::AnyGPUArray) = allunique(map(f, vec(A))) +end + + +## set operations + +# Converts `x` to a vector with element type `T`. Like Base, which builds a set of type `T`, +# this throws if an element can't be represented exactly. +function as_eltype(::Type{T}, x::AbstractArray) where {T} + v = vec(x) + eltype(v) === T && return v + y = similar(v, T) + y .= v + inexact = .!isequal.(y, v) + if any(inexact) + # throws the same error as Base + push!(Set{T}(), @allowscalar v[findall(inexact)[1]]) + end + return y +end + +# For each element of `x`, whether it is `isequal` to an element of `y`. Base's sets compare +# this way. Base's `issubset` and `isdisjoint` use `==` for arrays of fewer than 70 elements. +function isin(x::AbstractVector, y::AbstractArray) + found = similar(x, Bool) + if isempty(y) + fill!(found, false) + return found + end + # Comparing elements of different types without converting them, as `isequal` does + lt = set_isless(promote_type(eltype(x), eltype(y))) + s = sort(vec(y); lt) + insorted = let s = s, lt = lt + x -> begin + i = searchsortedfirst(s, x; lt) + return i <= length(s) && isequal(@inbounds(s[i]), x) + end + end + found .= insorted.(x) + return found +end + +function Base.union(a::AnyGPUArray, bs::AnyGPUArray...) + T = Base.promote_eltype(a, bs...) + return unique(vcat(as_eltype(T, a), map(b -> as_eltype(T, b), bs)...)) +end + +function Base.intersect(a::AnyGPUArray, bs::AnyGPUArray...) + u = unique(as_eltype(Base.promote_eltype(a, bs...), a)) + keep = fill!(similar(u, Bool), true) + for b in bs + keep .&= isin(u, b) + end + return u[keep] +end + +function Base.setdiff(a::AnyGPUArray, bs::AnyGPUArray...) + # Base promotes the element type before Julia 1.13 (JuliaLang/julia#59439) + T = VERSION >= v"1.13-" ? eltype(a) : Base.promote_eltype(a, bs...) + u = unique(as_eltype(T, a)) + keep = fill!(similar(u, Bool), true) + for b in bs + keep .&= .!isin(u, b) + end + return u[keep] +end + +# Keeps the elements that occur in an odd number of the arguments, counting each argument once +function Base.symdiff(a::AnyGPUArray, bs::AnyGPUArray...) + T = Base.promote_eltype(a, bs...) + c = vcat(unique(as_eltype(T, a)), map(b -> unique(as_eltype(T, b)), bs)...) + length(c) <= 1 && return c + p, s = sorted_keys(c) + head = group_starts(s) + starts = findall(head) + k = length(starts) + lens = similar(starts) + lens[1:k-1] .= view(starts, 2:k) .- view(starts, 1:k-1) + lens[k:k] .= length(s) + 1 .- view(starts, k:k) + head[starts] .= isodd.(lens) + return c[scatter_mask(p, head)] +end + +Base.issubset(a::AnyGPUArray, b::AnyGPUArray) = all(isin(vec(a), b)) +Base.issetequal(a::AnyGPUArray, b::AnyGPUArray) = issubset(a, b) && issubset(b, a) +Base.isdisjoint(a::AnyGPUArray, b::AnyGPUArray) = !any(isin(vec(a), b)) diff --git a/test/testsuite.jl b/test/testsuite.jl index 93beae86..321f3f80 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -129,6 +129,7 @@ include("testsuite/aliasing.jl") include("testsuite/vector.jl") include("testsuite/reductions.jl") include("testsuite/sorting.jl") +include("testsuite/set.jl") include("testsuite/accumulate.jl") include("testsuite/findall.jl") include("testsuite/broadcasting.jl") diff --git a/test/testsuite/set.jl b/test/testsuite/set.jl new file mode 100644 index 00000000..80f8e7e4 --- /dev/null +++ b/test/testsuite/set.jl @@ -0,0 +1,117 @@ +# random elements with many duplicates +dups(ET, dims...) = ET <: Complex ? ET.(rand(1:3, dims...), rand(1:3, dims...)) : + ET.(rand(1:5, dims...)) + +@testsuite "set/unique" (AT, eltypes)->begin + @testset "$ET" for ET in eltypes + # Base returns elements in the order they first occur + for n in (0, 1, 2, 10, 10000) + @test compare_exact(unique, AT, dups(ET, n)) + @test compare_exact(allunique, AT, dups(ET, n)) + end + @test compare_exact(unique, AT, dups(ET, 10, 20)) + @test compare_exact(allunique, AT, ET.(collect(1:100))) + @test compare_exact(A -> unique(abs2, A), AT, dups(ET, 1000)) + @test compare_exact(A -> unique(abs2, A), AT, ET[-2, 2, 1, -1]) + @test compare_exact(A -> unique!(A), AT, dups(ET, 1000)) + @test compare_exact(A -> unique!(abs2, A), AT, dups(ET, 1000)) + if hasmethod(allunique, Tuple{Any, Any}) + @test compare_exact(A -> allunique(abs2, A), AT, dups(ET, 3)) + end + + # unique slices + for dims in (1, 2, 3) + @test compare_exact(A -> unique(A; dims), AT, dups(ET, 100, 2)) + @test compare_exact(A -> unique(A; dims), AT, dups(ET, 2, 3, 100)) + end + @test compare_exact(A -> unique(A; dims=2), AT, dups(ET, 2, 0)) + @test compare_exact(A -> unique(A; dims=1), AT, dups(ET, 1, 3)) + # Base returns a copy when `dims` is out of range + for dims in (-1, 0, 4) + @test compare_exact(A -> unique(A; dims), AT, dups(ET, 2, 3, 4)) + end + + # views and reshaped arrays + @test compare_exact(A -> unique(view(A, 11:90)), AT, dups(ET, 100)) + @test compare_exact(A -> unique(reshape(A, 10, 10); dims=2), AT, dups(ET, 100)) + end + + # with `isequal`, signed zeros differ and NaNs are equal + for ET in filter(isrealfloattype, eltypes) + x = ET[0, -0.0, NaN, 1, -NaN, NaN, -0.0, 1] + @test compare_exact(unique, AT, x) + @test compare_exact(allunique, AT, x) + @test compare_exact(A -> unique(A; dims=2), AT, [x x x]') + end + + @test compare_exact(unique, AT, rand(Bool, 100)) + + # isbits unions, including ones with `missing` + for x in (Union{Int32, Float32}[Int32(1), 1.0f0, Int32(2), 2.0f0, Int32(1)], + Union{Missing, Float32}[1, missing, 2, missing, 1], + Union{Missing, ComplexF32}[1, missing, 1im, missing, 1]) + @test compare_exact(unique, AT, x) + @test compare_exact(allunique, AT, x) + @test compare_exact(A -> unique(A; dims=1), AT, x) + end + + # elements without `isless` are not supported (Base handles them by hashing) + if AT <: AbstractGPUArray + @test_throws Exception unique(AT([Some(Int32(1)), Some(Int32(2))])) + end +end + +@testsuite "set/operations" (AT, eltypes)->begin + @testset "$ET" for ET in eltypes + for (m, n) in ((0, 0), (0, 10), (10, 0), (1, 1), (100, 200)) + a, b, c = dups(ET, m), dups(ET, n), dups(ET, 50) + for f in (union, intersect, setdiff, symdiff, issubset, issetequal, isdisjoint) + @test compare_exact(f, AT, a, b) + end + for f in (union, intersect, setdiff, symdiff) + @test compare_exact(f, AT, a, b, c) + @test compare_exact(f, AT, a) + end + end + @test compare_exact(intersect, AT, dups(ET, 4, 5), dups(ET, 3)) + @test compare_exact(issetequal, AT, ET[1, 2, 2], ET[2, 1]) + end + + # element types are promoted as in Base, but elements are compared without converting them + @test compare_exact(union, AT, Int32[3, 1], Float32[2, 3]) + @test compare_exact(intersect, AT, Int32[3, 1, 2, 1], Float32[2, 1]) + @test compare_exact(setdiff, AT, Int32[3, 1, 2, 1], Float32[2.5, 1]) + @test compare_exact(symdiff, AT, Int16[1, 2], Float32[2, 4]) + big, rounded = Int32[2^24 + 1, 1], Float32[2^24, 1] + for f in (union, intersect, setdiff, symdiff, issubset, isdisjoint) + # Base throws when it would convert an element inexactly, which depends on the order + @test compare_exact(f, AT, big, rounded) + @test compare_exact(f, AT, rounded, big) + end + @test compare_exact(issubset, AT, [-1], UInt[1]) + @test compare_exact(intersect, AT, ComplexF32[1, 2im], Float32[1, 2]) + # isbits unions + m = Union{Missing, Float32}[1, missing, 3] + mc = Union{Missing, ComplexF32}[1, missing, 2im] + for (x, y) in ((Float32[1, 2], m), (m, Float32[1, 2]), (ComplexF32[1, 2im], mc), (mc, Float32[1])) + for f in (union, intersect, setdiff, symdiff) + @test compare_exact(f, AT, x, y) + end + end + + # with `isequal`, signed zeros differ and NaNs are equal + x, y = Float32[0, NaN, 1, -0.0], Float32[-0.0, NaN, 2] + for f in (union, intersect, setdiff, symdiff) + @test compare_exact(f, AT, x, y) + end + # The predicates use `isequal` too. Base's `issubset` and `isdisjoint` use `==` for arrays + # of fewer than 70 elements, so these are longer. + r = Float32.(1:100) + @test compare_exact(issubset, AT, [NaN32; r], [NaN32; r; 200]) + @test compare_exact(issubset, AT, [0.0f0; r], [-0.0f0; r]) + @test compare_exact(isdisjoint, AT, [NaN32; r], [NaN32; r .+ 200]) + @test compare_exact(isdisjoint, AT, [0.0f0; r], [-0.0f0; r .+ 200]) + mr = Union{Missing, Float32}[missing; r] + @test compare_exact(issubset, AT, mr, [mr; 200]) + @test compare_exact(isdisjoint, AT, mr, [missing; r .+ 200]) +end diff --git a/test/testsuite/sorting.jl b/test/testsuite/sorting.jl index a629f552..bac13dc5 100644 --- a/test/testsuite/sorting.jl +++ b/test/testsuite/sorting.jl @@ -57,6 +57,19 @@ end @test compare(A -> sort(A; by=first, alg=MergeSort), AT, tagged) @test compare(A -> sortperm(A; by=first), AT, tagged) + # comparators that capture arrays, e.g. to sort the columns of a matrix + function sortcols(A; kwargs...) + lt = let A = A + (i, j) -> isless((A[1, i], A[2, i]), (A[1, j], A[2, j])) + end + idx = similar(A, Int, size(A, 2)) + idx .= 1:size(A, 2) + return sort(idx; lt, kwargs...), sortperm(idx; lt, kwargs...) + end + @test compare(sortcols, AT, rand(1:3, 2, 1000)) + @test compare(A -> sortcols(A; rev=true), AT, rand(1:3, 2, 1000)) + @test compare(A -> sortcols(A; by=i -> 1001 - i), AT, rand(1:3, 2, 1000)) + # Arrays other than vectors need `dims`, and vectors take none, as in Base @test_throws MethodError sort(AT(x); dims=1) @test_throws MethodError sort!(AT(x); dims=1)