diff --git a/Project.toml b/Project.toml index 77084842..5d5e1f65 100644 --- a/Project.toml +++ b/Project.toml @@ -1,11 +1,12 @@ name = "GPUArrays" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "11.5.15" +version = "12.0.0" [workspace] projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"] [deps] +AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -26,6 +27,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" JLD2Ext = "JLD2" [compat] +AcceleratedKernels = "0.5" Adapt = "4.6.1" GPUArraysCore = "= 0.2.1" JLD2 = "0.4, 0.5, 0.6" diff --git a/lib/JLArrays/Project.toml b/lib/JLArrays/Project.toml index 799f9b5e..10bcd7ac 100644 --- a/lib/JLArrays/Project.toml +++ b/lib/JLArrays/Project.toml @@ -1,6 +1,6 @@ name = "JLArrays" uuid = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" -version = "0.4.0" +version = "0.4.1" authors = ["Tim Besard "] [deps] @@ -13,7 +13,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [compat] Adapt = "2.0, 3.0, 4.0" -GPUArrays = "11.1" +GPUArrays = "12" KernelAbstractions = "0.9" LinearAlgebra = "1" Random = "1" diff --git a/lib/JLArrays/src/JLArrays.jl b/lib/JLArrays/src/JLArrays.jl index 4b034901..c6caf852 100644 --- a/lib/JLArrays/src/JLArrays.jl +++ b/lib/JLArrays/src/JLArrays.jl @@ -36,6 +36,9 @@ include("array.jl") include("sparse.jl") include("broadcast.jl") include("mapreduce.jl") +include("sorting.jl") +include("accumulate.jl") +include("findall.jl") # KernelAbstractions include("JLKernels.jl") diff --git a/lib/JLArrays/src/accumulate.jl b/lib/JLArrays/src/accumulate.jl new file mode 100644 index 00000000..1bf87454 --- /dev/null +++ b/lib/JLArrays/src/accumulate.jl @@ -0,0 +1,20 @@ +# scans + +# As for sorting: Base, applied to the storage, with the method shapes of GPUArrays. + +function Base._accumulate!(op, B::AnyJLArray, A::AnyJLVector, dims::Nothing, init::Nothing) + accumulate!(op, _host(B), _host(A)) + return B +end +function Base._accumulate!(op, B::AnyJLArray, A::AnyJLVector, dims::Nothing, init::Some) + accumulate!(op, _host(B), _host(A); init=something(init)) + return B +end +function Base._accumulate!(op, B::AnyJLArray, A::AnyJLArray, dims::Integer, init::Nothing) + accumulate!(op, _host(B), _host(A); dims) + return B +end +function Base._accumulate!(op, B::AnyJLArray, A::AnyJLArray, dims::Integer, init::Some) + accumulate!(op, _host(B), _host(A); dims, init=something(init)) + return B +end diff --git a/lib/JLArrays/src/findall.jl b/lib/JLArrays/src/findall.jl new file mode 100644 index 00000000..11fdecbf --- /dev/null +++ b/lib/JLArrays/src/findall.jl @@ -0,0 +1,14 @@ +# findall and predicates + +# As for sorting: Base, applied to the storage, with the method shapes of GPUArrays. + +Base.findall(bools::AnyJLArray{Bool}) = JLArray(findall(_host(bools))) +Base.findall(f::Function, A::AnyJLArray) = JLArray(findall(f, _host(A))) +Base.findall(f::Base.Fix2{typeof(in)}, A::AnyJLArray) = JLArray(findall(f, _host(A))) +Base.getindex(A::JLArray, mask::AnyJLArray{Bool}) = JLArray(_host(A)[_host(mask)]) + +for fname in (:any, :all) + @eval Base.$fname(f::Function, A::AnyJLArray; dims=:) = + dims === Colon() ? $fname(f, _host(A)) : + invoke($fname, Tuple{Function, GPUArrays.AnyGPUArray}, f, A; dims) +end diff --git a/lib/JLArrays/src/mapreduce.jl b/lib/JLArrays/src/mapreduce.jl index b2cba874..5d5ccff5 100644 --- a/lib/JLArrays/src/mapreduce.jl +++ b/lib/JLArrays/src/mapreduce.jl @@ -3,11 +3,23 @@ struct ArrayNoCopy end Adapt.adapt_storage(::ArrayNoCopy, x::JLArray) = typed_data(x) -function GPUArrays.mapreducedim!(f, op, R::AnyJLArray, A::Union{AbstractArray,Broadcast.Broadcasted}; - init=nothing) - if init !== nothing - fill!(R, init) - end - @allowscalar Base.reducedim!(op, adapt(ArrayNoCopy(), R), map(f, A)) - R +# GPUArrays reduces with AcceleratedKernels, whose kernels do not run on JLArrays' back-end. The +# reference implementation runs AcceleratedKernels' host algorithms on the arrays' storage +# instead, through GPUArrays' internal indirection for this purpose. (The storage is wrapped +# without a copy, so the arrays must be preserved; a `Broadcasted` source is materialized on the +# host.) +_host_source(A::Broadcast.Broadcasted) = # (`copy` of a 0-dimensional one gives a scalar) + Array(copyto!(similar(A, Broadcast.combine_eltypes(A.f, A.args)), A)) +_host_source(A) = adapt(ArrayNoCopy(), A) + +GPUArrays._ak_mapreduce(f, op, A::Union{AnyJLArray, Broadcast.Broadcasted{<:JLArrayStyle}}; + backend=nothing, kwargs...) = + GC.@preserve A @allowscalar GPUArrays.AK.mapreduce(f, op, _host_source(A); kwargs...) + +function GPUArrays._ak_mapreducedim!(f, op, R::AnyJLArray, + A::Union{AbstractArray, Broadcast.Broadcasted}; + backend=nothing, kwargs...) + GC.@preserve R A @allowscalar GPUArrays.AK.mapreducedim!( + f, op, adapt(ArrayNoCopy(), R), _host_source(A); kwargs...) + return R end diff --git a/lib/JLArrays/src/sorting.jl b/lib/JLArrays/src/sorting.jl new file mode 100644 index 00000000..5aec9ca0 --- /dev/null +++ b/lib/JLArrays/src/sorting.jl @@ -0,0 +1,49 @@ +# sorting and reversing + +# GPUArrays implements these with AcceleratedKernels, whose kernels do not run on JLArrays' +# back-end. The reference implementation applies Base to the arrays' storage instead, with the +# same method shapes, and accepts the same `alg` values as GPU arrays (`GPUArrays._akalg`). + +_host(A) = adapt(ArrayNoCopy(), A) + +function Base.sort!(v::AnyJLVector; alg=nothing, lt=isless, by=identity, rev=nothing, + order::Base.Order.Ordering=Base.Order.Forward, scratch=nothing) + GPUArrays._akalg(alg) + sort!(_host(v); lt, by, rev, order) + return v +end +function Base.sort!(A::AnyJLArray; dims::Integer, alg=nothing, scratch=nothing, kwargs...) + GPUArrays._akalg(alg) + sort!(_host(A); dims, kwargs...) + return A +end + +function Base.sortperm!(ix::AnyJLArray{<:Integer}, v::AnyJLVector; alg=nothing, + scratch=nothing, initialized::Bool=false, dims=nothing, kwargs...) + dims === nothing || throw(ArgumentError("sortperm! of a vector does not accept `dims`")) + axes(ix) == axes(v) || + throw(ArgumentError("index array must have the same axes as the source array")) + GPUArrays._akalg(alg) + sortperm!(_host(ix), _host(v); kwargs...) + return ix +end +function Base.sortperm!(ix::AnyJLArray{<:Integer}, A::AnyJLArray; dims::Integer, alg=nothing, + scratch=nothing, initialized::Bool=false, kwargs...) + axes(ix) == axes(A) || + throw(ArgumentError("index array must have the same axes as the source array")) + GPUArrays._akalg(alg) + sortperm!(_host(ix), _host(A); dims, kwargs...) + return ix +end + +function Base.reverse!(A::AnyJLArray; dims=:) + reverse!(_host(A); dims) + return A +end +Base.reverse(A::AnyJLArray; dims=:) = reverse!(copy(A); dims) +Base.reverse!(v::AnyJLVector; dims=:) = invoke(reverse!, Tuple{AbstractVector}, v; dims) +Base.reverse(v::AnyJLVector; dims=:) = invoke(reverse, Tuple{AbstractVector}, v; dims) +function Base.reverse!(v::AnyJLVector, start::Integer, stop::Integer=lastindex(v)) + reverse!(_host(v), start, stop) + return v +end diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index d634add4..10b5fdb6 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -17,6 +17,8 @@ using Reexport using KernelAbstractions +import AcceleratedKernels as AK + # device functionality include("device/abstractarray.jl") include("device/sparse.jl") @@ -27,6 +29,8 @@ include("host/construction.jl") ## integrations and specialized methods include("host/base.jl") include("host/indexing.jl") +include("host/sorting.jl") +include("host/accumulate.jl") include("host/broadcast.jl") include("host/mapreduce.jl") include("host/gemm.jl") diff --git a/src/host/accumulate.jl b/src/host/accumulate.jl new file mode 100644 index 00000000..9212699f --- /dev/null +++ b/src/host/accumulate.jl @@ -0,0 +1,36 @@ +# scans, implemented by AcceleratedKernels +# +# Base's `accumulate!` and `accumulate` reach these four methods, split as Base's are. `init` is +# passed on only when given (`Some(x)`), so that an explicit `init=nothing` stays an initial +# value. Base's rules are kept here: without `dims`, only vectors are scanned, and a `dims` beyond +# the array's copies it, ignoring `init` (AcceleratedKernels would apply `init` to every element). +# `B`'s element type, which Base's front-end chose, sets the running values' type in +# AcceleratedKernels. + +Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUVector, dims::Nothing, init::Nothing) = + AK.accumulate!(op, B, A) +Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUVector, dims::Nothing, init::Some) = + AK.accumulate!(op, B, A; init=something(init)) +function Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUArray, dims::Integer, init::Nothing) + dims > 0 || throw(ArgumentError("dims must be a positive integer")) + axes(B) == axes(A) || throw(DimensionMismatch("shape of B must match A")) + dims > ndims(A) && return copyto!(B, A) + AK.accumulate!(op, B, A; dims) +end +function Base._accumulate!(op, B::AnyGPUArray, A::AnyGPUArray, dims::Integer, init::Some) + dims > 0 || throw(ArgumentError("dims must be a positive integer")) + axes(B) == axes(A) || throw(DimensionMismatch("shape of B must match A")) + dims > ndims(A) && return copyto!(B, A) + AK.accumulate!(op, B, A; dims, init=something(init)) +end + +# (`cumsum!` of floating-point vectors) +Base.accumulate_pairwise!(op, B::AnyGPUVector, A::AnyGPUVector) = accumulate!(op, B, A) + +# Without `dims`, other arrays are scanned in linear order, keeping their shape +function Base.accumulate(op, A::AnyGPUArray; dims::Union{Nothing,Integer}=nothing, kw...) + if dims === nothing && !(A isa AbstractVector) + return reshape(accumulate(op, vec(A); kw...), size(A)) + end + return invoke(accumulate, Tuple{Any, Any}, op, A; dims, kw...) +end diff --git a/src/host/indexing.jl b/src/host/indexing.jl index 7ab66fe2..2a1b21e6 100644 --- a/src/host/indexing.jl +++ b/src/host/indexing.jl @@ -148,6 +148,11 @@ const IndexGPUArray{T} = Union{AbstractGPUArray{T}, Base.checkindex(Bool, inds, i) end) end +# ... except for a logical mask, which must have the indexed axes (Base's rule) +Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, I::IndexGPUArray{Bool}) = + ndims(I) == 1 && Base.axes1(I) == inds +Base.checkindex(::Type{Bool}, inds::Tuple, I::IndexGPUArray{Bool}) = + length(inds) == ndims(I) && all(map(==, inds, axes(I))) @inline function Base.checkindex(::Type{Bool}, inds::Tuple, I::IndexGPUArray{<:CartesianIndex}) @@ -232,10 +237,48 @@ end Base.findfirst(A::AnyGPUArray{Bool}) = findfirst(identity, A) Base.findlast(A::AnyGPUArray{Bool}) = findlast(identity, A) -function findminmax(binop, f, A::AnyGPUArray; init, dims) - indices = EachIndex(A) - dummy_index = firstindex(A) +# findall, implemented by AcceleratedKernels, which selects `items` of the array's length: Base's +# indices, `keys(A)`, which are linear for vectors and Cartesian otherwise, except that Base's +# predicate form makes those of a 0-dimensional array linear +Base.findall(bools::AnyGPUArray{Bool}) = AK.findall(bools) +Base.findall(f::Function, A::AnyGPUArray) = AK.findall(f, A; items=_findall_items(A)) +Base.findall(f::Base.Fix2{typeof(in)}, A::AnyGPUArray) = # (Base: `keys(A)`) + AK.findall(f, A) +_findall_items(A) = ndims(A) == 0 ? LinearIndices(A) : keys(A) + +# logical indexing: Base's `LogicalIndex` iterates, so the mask becomes the indices it selects. +# Those no longer carry the mask's shape, so a single mask is checked against the array first, as +# Base does; a mask mixed with other indices is not (as before). +Base.to_index(::AnyGPUArray, I::AbstractArray{Bool}) = findall(I) +@static if VERSION >= v"1.11.0-DEV.1157" + Base.to_indices(A::AnyGPUArray, I::Tuple{AbstractArray{Bool}}) = + (checkbounds(A, I[1]); (Base.to_index(A, I[1]),)) +else + # (also reached for the last of several indices, whose `inds` are then not all of `A`'s) + _check_mask(A, inds, mask) = length(inds) == ndims(A) ? checkbounds(A, mask) : nothing + Base.to_indices(A::AnyGPUArray, inds, + I::Tuple{Union{Array{Bool,N}, BitArray{N}}}) where {N} = + (_check_mask(A, inds, I[1]); (Base.to_index(A, I[1]),)) + Base.to_indices(A::AnyGPUArray, inds, I::Tuple{AbstractArray{Bool}}) = + (_check_mask(A, inds, I[1]); (Base.to_index(A, I[1]),)) +end +# ... except that a mask of the array's shape selects the values themselves, in one pass +function Base.getindex(A::AbstractGPUArray, mask::AnyGPUArray{Bool}) + checkbounds(A, mask) + axes(mask) == axes(A) || return invoke(getindex, Tuple{AbstractGPUArray, Vararg{Any}}, A, mask) + return AK.findall(mask; items=A) +end + +# `findmin` and `findmax` reduce `(f(x), i)` pairs, with `i` the position of `x`: without an +# `init`, as partial results start from their first element. Indices are Base's, `keys(A)`; so are +# the errors of empty inputs, from a host stand-in. +struct _FindPair{F} + f::F +end +(p::_FindPair)(x, i) = (p.f(x), i) + +function findminmax(binop, f, A::AnyGPUArray; dims) function reduction(t1, t2) (x, i), (y, j) = t1, t2 @@ -243,25 +286,33 @@ function findminmax(binop, f, A::AnyGPUArray; init, dims) isequal(x, y) && return (x, min(i, j)) return t1 end - - fA = f.(A) - if dims == Colon() - res = mapreduce(tuple, reduction, fA, indices; init = (init, dummy_index)) + if isempty(A) + h = (binop === Base.isless ? findmax : findmin)(f, Array{eltype(A)}(undef, size(A)); dims) + return dims === Colon() ? h : (copyto!(similar(A, eltype(h[1]), size(h[1])), h[1]), + copyto!(similar(A, eltype(h[2]), size(h[2])), h[2])) + end - # out of consistency with Base.findarray, return a CartesianIndex - # when the input is a multidimensional array - return (res[1], ndims(A) == 1 ? res[2] : CartesianIndices(A)[res[2]]) + # (along valid `dims`, a 0-dimensional array reduces nothing; Julia 1.10's `reduced_indices` + # cannot take it) + rdims = ndims(A) == 0 && !(dims isa Colon) && all(d -> d isa Integer && d >= 1, dims) ? () : dims + res = mapreduce(_FindPair(f), reduction, A, LinearIndices(A); dims=rdims) + I = keys(A) + if dims === Colon() + return (res[1], I[res[2]]) else - res = mapreduce(tuple, reduction, fA, indices; - init = (init, dummy_index), dims=dims) - vals = map(x->x[1], res) - inds = map(x->ndims(A) == 1 ? x[2] : CartesianIndices(A)[x[2]], res) + # (`map!`, as `map` of a 0-dimensional array would give a scalar) + vals = map!(first, similar(res, fieldtype(eltype(res), 1)), res) + inds = map!(x -> I[x[2]], similar(res, eltype(I)), res) return (vals, inds) end end -Base.findmax(a::AnyGPUArray; dims=:) = findminmax(Base.isless, identity, a; init=typemin(eltype(a)), dims) -Base.findmin(a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, identity, a; init=typemax(eltype(a)), dims) -Base.findmax(f::Function, a::AnyGPUArray; dims=:) = findminmax(Base.isless, f, a; init=typemin(f(zero(eltype(a)))), dims) -Base.findmin(f::Function, a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, f, a; init=typemax(f(zero(eltype(a)))), dims) +Base.findmax(a::AnyGPUArray; dims=:) = findminmax(Base.isless, identity, a; dims) +Base.findmin(a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, identity, a; dims) +Base.findmax(f::Function, a::AnyGPUArray; dims=:) = findminmax(Base.isless, f, a; dims) +Base.findmin(f::Function, a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, f, a; dims) + +# the element that minimizes or maximizes `f` (Base iterates) +Base.argmax(f::Function, a::AnyGPUArray) = @allowscalar a[findmax(f, a)[2]] +Base.argmin(f::Function, a::AnyGPUArray) = @allowscalar a[findmin(f, a)[2]] diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 2d19fb46..a873f48f 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -1180,11 +1180,8 @@ function Base.isone(x::AbstractGPUMatrix{T}) where {T} bc = Broadcast.broadcasted(x, CartesianIndices(x)) do _x, inds _x - (inds[1] == inds[2] ? one(_x) : zero(_x)) end - # call `GPUArrays.mapreducedim!` directly, which supports Broadcasted inputs - y = similar(x, Bool, 1) - GPUArrays.mapreducedim!(iszero, &, y, Broadcast.instantiate(bc); init=true) - - Array(y)[] + # reduce the Broadcasted object directly + _ak_mapreduce(iszero, &, Broadcast.instantiate(bc); backend=get_backend(x), init=true) end ## Kronecker product diff --git a/src/host/mapreduce.jl b/src/host/mapreduce.jl index b763a9be..6ce7fe48 100644 --- a/src/host/mapreduce.jl +++ b/src/host/mapreduce.jl @@ -1,36 +1,42 @@ # map-reduce +# +# Base's reductions of GPU arrays, implemented with AcceleratedKernels' primitives: `AK.mapreduce` +# for scalar results, `AK.mapreducedim!` into a destination otherwise. AcceleratedKernels applies +# `init` once and needs no neutral element; everything that only exists because Base says so +# (result types, empty results and their errors, the one-element result, `dims` rules) is decided +# here, on the host. const AbstractArrayOrBroadcasted = Union{AbstractArray,Broadcast.Broadcasted} -# GPUArrays' mapreduce methods build on `Base.mapreducedim!`, but with an additional -# argument `init` value to avoid eager initialization of `R` (if set to something). -mapreducedim!(f, op, R::AnyGPUArray, A::AbstractArrayOrBroadcasted; - init=nothing) = error("Not implemented") # COV_EXCL_LINE -# resolve ambiguities -Base.mapreducedim!(f, op, R::AnyGPUArray, A::AbstractArray) = mapreducedim!(f, op, R, A) -Base.mapreducedim!(f, op, R::AnyGPUArray, A::Broadcast.Broadcasted) = mapreducedim!(f, op, R, A) +# AcceleratedKernels' primitives, on the backend `backend` (a `Broadcasted` source may hold arrays +# without one, such as `EachIndex`). The reference back-end (JLArrays) runs them on its storage +# instead; these are internal, not an extension point for back-ends. +_ak_mapreduce(f, op, A; kwargs...) = AK.mapreduce(f, op, A; kwargs...) +_ak_mapreducedim!(f, op, R, A; kwargs...) = AK.mapreducedim!(f, op, R, A; kwargs...) + +# `Base.mapreducedim!` folds into `R`'s values +Base.mapreducedim!(f, op, R::AnyGPUArray, A::AbstractArray) = + _ak_mapreducedim!(f, op, R, A; backend=get_backend(R)) +Base.mapreducedim!(f, op, R::AnyGPUArray, A::Broadcast.Broadcasted) = + _ak_mapreducedim!(f, op, R, A; backend=get_backend(R)) # `neutral_element` lives in GPUArraysCore, so that packages building on GPUArraysCore share it import GPUArraysCore: neutral_element +# `init` when none is given (an explicit `init=nothing` is an initial value, as in Base) +struct _NoInit end + # resolve ambiguities Base.mapreduce(f, op, A::AnyGPUArray, As::AbstractArrayOrBroadcasted...; - dims=:, init=nothing) = _mapreduce(f, op, A, As...; dims=dims, init=init) + dims=:, init=_NoInit()) = _mapreduce(f, op, A, As...; dims=dims, init=init) Base.mapreduce(f, op, A::Broadcast.Broadcasted{<:AbstractGPUArrayStyle}, As::AbstractArrayOrBroadcasted...; - dims=:, init=nothing) = _mapreduce(f, op, A, As...; dims=dims, init=init) + dims=:, init=_NoInit()) = _mapreduce(f, op, A, As...; dims=dims, init=init) function _mapreduce(f::F, op::OP, As::Vararg{Any,N}; dims::D, init) where {F,OP,N,D} - # figure out the destination container type by looking at the initializer element, - # or by relying on inference to reason through the map and reduce functions - if init === nothing - ET = Broadcast.combine_eltypes(f, As) - ET = Base.promote_op(op, ET, ET) - (ET === Union{} || ET === Any) && - error("mapreduce cannot figure the output element type, please pass an explicit init value") - - init = neutral_element(op, ET) - else - ET = typeof(init) + if !(dims isa Colon) + all(d -> d isa Integer, dims) || + throw(ArgumentError("reduced dimension(s) must be integers")) + all(d -> d >= 1, dims) || throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) end # apply the mapping function to the input arrays @@ -52,44 +58,159 @@ function _mapreduce(f::F, op::OP, As::Vararg{Any,N}; dims::D, init) where {F,OP, end f = identity end + S = _source_eltype(A) + M = Base.promote_op(f, S) + backend = _source_backend(As) + + if dims isa Colon + # Base's results for empty and one-element inputs; otherwise AcceleratedKernels' result, + # converted to the type Base's fold settles on + if init isa _NoInit + length(A) == 0 && + return Base.mapreduce_empty_iter(f, op, Array{S}(undef, 0), Base.HasEltype()) + length(A) == 1 && return @allowscalar Base.mapreduce_first(f, op, first(A)) + return convert(_fold_type(op, Union{}, M), _ak_mapreduce(f, op, A; backend)) + else + length(A) == 0 && return init + length(A) == 1 && return @allowscalar op(init, f(first(A))) + return convert(_fold_type(op, typeof(init), M), _ak_mapreduce(f, op, A; backend, init)) + end + end - # allocate an output container - sz = size(A) - red = ntuple(i->(dims==Colon() || i in dims) ? 1 : sz[i], length(sz)) - R = similar(A, ET, red) - - # perform the reduction - if prod(sz) == 0 - fill!(R, init) + # along `dims`: Base's shape, and Base's element type (`typeof(init)`, else the fold type) + rax = Base.reduced_indices(axes(A), dims) + if !(init isa _NoInit) + R = similar(A, typeof(init), length.(rax)) + _ak_mapreducedim!(f, op, R, A; backend, init) + return R + end + if any(d -> d <= ndims(A) && size(A)[d] == 0, dims) + # Base's values for an empty reduction (or its error), from a host stand-in; `f` may be + # called, and may index device arrays + h = @allowscalar Base.reducedim_init(f, op, Array{S}(undef, size(A)), dims) + R = similar(A, eltype(h), size(h)) + isempty(R) || copyto!(R, h) + return R + end + T = _fold_type(op, Union{}, M) + R = similar(A, T, length.(rax)) + z = _reducedim_zero(op, T) + if z === nothing + _ak_mapreducedim!(f, op, R, A; backend, overwrite=true) else - mapreducedim!(f, op, R, A; init) + # Base's sums along `dims` start from zero (`reducedim_init`), so a slice of `-0.0`s sums + # to `0.0`; the other operators' initial values do not change the result + _ak_mapreducedim!(f, op, R, A; backend, init=z) end + return R +end - # return the result - if dims === Colon() - @allowscalar R[] - else - R +_reducedim_zero(op, ::Type) = nothing +_reducedim_zero(::Union{typeof(+), typeof(Base.add_sum)}, ::Type{T}) where {T} = + hasmethod(zero, Tuple{Type{T}}) ? zero(T) : nothing + +# The element type of a reduction source before mapping +_source_eltype(A::AbstractArray) = eltype(A) +_source_eltype(bc::Broadcast.Broadcasted) = Broadcast.combine_eltypes(bc.f, bc.args) + +# The backend of the first GPU array among a reduction's sources +_source_backend(A::AnyGPUArray) = get_backend(A) +_source_backend(bc::Broadcast.Broadcasted) = _source_backend(bc.args) +_source_backend(::Tuple{}) = nothing +function _source_backend(t::Tuple) + b = _source_backend(first(t)) + return b === nothing ? _source_backend(Base.tail(t)) : b +end +_source_backend(_) = nothing + +# The type Base's fold `op(op(init, x₁), x₂)...` settles on, for elements of type `M` (without +# `init`, `I === Union{}`, from `op(x₁, x₂)`): the types `op` returns, not `init`'s +function _fold_type(op, ::Type{I}, ::Type{M}) where {I, M} + T = I === Union{} ? Base.promote_op(op, M, M) : Base.promote_op(op, I, M) + for _ in 1:8 + S = promote_type(T, Base.promote_op(op, T, M)) + S == T && return T + T = S end + return T end -Base.any(A::AnyGPUArray{Bool}) = mapreduce(identity, |, A) -Base.all(A::AnyGPUArray{Bool}) = mapreduce(identity, &, A) +# `any` and `all` follow Base. A scalar call with a predicate that returns a `Bool` uses +# AcceleratedKernels' short-circuiting implementation; one that can return `missing` reduces a code +# of each value (`false` < `missing` < `true`) with `max` or `min`, which gives Base's three-valued +# logic without storing `missing`. Along `dims`, as in Base, the predicate must return a `Bool`. +# Values the predicate cannot return (by inference) are an error before launching, and nothing is +# checked for an empty array, whose elements Base never passes to the predicate. +struct _Bool3{F} <: Function + f::F +end +@inline (c::_Bool3)(x) = _bool3(c.f(x)) +@inline _bool3(x::Bool) = x ? 0x02 : 0x00 +@inline _bool3(::Missing) = 0x01 +@inline _bool3(x) = throw(TypeError(:any, "", Union{Bool, Missing}, x)) +_unbool3(c::UInt8) = c == 0x01 ? missing : c == 0x02 + +struct _BoolOnly{F} <: Function + f::F +end +@inline function (b::_BoolOnly)(x) + y = b.f(x) + y isa Bool || throw(TypeError(:any, "", Bool, y)) + return y +end -Base.any(f::Function, A::AnyGPUArray) = mapreduce(f, |, A) -Base.all(f::Function, A::AnyGPUArray) = mapreduce(f, &, A) +for (fname, op, op3, empty3) in ((:any, :|, :max, 0x00), (:all, :&, :min, 0x02)) + @eval function Base.$fname(f::Function, A::AnyGPUArray; dims=:) + T = Base.promote_op(f, eltype(A)) + if dims === Colon() + isempty(A) && return $(fname === :all) + T === Union{} || Bool <: T || T <: Union{Bool, Missing} || throw(ArgumentError( + "the predicate of `$($fname)` must return a `Bool` or `missing`, not `$T`")) + T <: Bool && return AK.$fname(f, A) + return _unbool3(mapreduce(_Bool3(f), $op3, A; init=$empty3)) + else + # (`reduced_indices` checks `dims` as Base does) + isempty(A) && return fill!(similar(A, Bool, length.(Base.reduced_indices(axes(A), dims))), + $(fname === :all)) + T === Union{} || Bool <: T || throw(ArgumentError( + "the predicate of `$($fname)` along `dims` must return a `Bool`, not `$T`")) + return mapreduce(_BoolOnly(f), $op, A; dims) + end + end +end +Base.any(A::AnyGPUArray{<:Union{Bool, Missing}}; dims=:) = any(identity, A; dims) +Base.all(A::AnyGPUArray{<:Union{Bool, Missing}}; dims=:) = all(identity, A; dims) Base.count(pred::Function, A::AnyGPUArray; dims=:, init=0) = - mapreduce(pred, Base.add_sum, A; init=init, dims=dims) - -# avoid calling into `initarray!` -for (fname, op) in [(:sum, :(Base.add_sum)), (:prod, :(Base.mul_prod)), - (:maximum, :(Base.max)), (:minimum, :(Base.min)), - (:all, :&), (:any, :|)] + mapreduce(Base._bool(pred), Base.add_sum, A; init=init, dims=dims) + +# The in-place reductions, with Base's `init::Bool` but without Base's pass that initializes `r` +# (`initarray!`): `init=true` reduces from the operator's identity, as an `init`, or where it has +# none overwrites `r`; `init=false` folds into `r`'s values. When every slice is empty, Base's +# result (or error) comes from a host stand-in. +for (fname, op, idfun) in ((:sum, :(Base.add_sum), zero), (:prod, :(Base.mul_prod), one), + (:maximum, :max, nothing), (:minimum, :min, nothing), + (:extrema, :(Base._extrema_rf), nothing), + (:any, :|, Returns(false)), (:all, :&, Returns(true)), + (:count, :(Base.add_sum), zero)) fname! = Symbol(fname, '!') - @eval begin - Base.$(fname!)(f::Function, r::AnyGPUArray, A::AnyGPUArray{T}) where T = - GPUArrays.mapreducedim!(f, $(op), r, A; init=neutral_element($(op), T)) + mapf = fname === :extrema ? :(Base.ExtremaMap(f)) : fname === :count ? :(Base._bool(f)) : :f + ftype = fname === :count ? :Any : :Function + @eval function Base.$(fname!)(f::$ftype, r::AnyGPUArray, A::AnyGPUArray; init::Bool=true) + if isempty(A) && !isempty(r) + hr = Array(r) + Base.$(fname!)(f, hr, Array{eltype(A)}(undef, size(A)); init) + return copyto!(r, hr) + end + backend = get_backend(r) + if !init + _ak_mapreducedim!($mapf, $op, r, A; backend) + elseif $(idfun === nothing) + _ak_mapreducedim!($mapf, $op, r, A; backend, overwrite=true) + else + _ak_mapreducedim!($mapf, $op, r, A; backend, init=$idfun(eltype(r))) + end + return r end end diff --git a/src/host/sorting.jl b/src/host/sorting.jl new file mode 100644 index 00000000..b8842459 --- /dev/null +++ b/src/host/sorting.jl @@ -0,0 +1,88 @@ +# sorting and reversing, implemented by AcceleratedKernels +# +# AcceleratedKernels chooses the algorithm and its settings for the array's device. Base's +# algorithm objects become requirements (`_akalg`); an AcceleratedKernels algorithm passed as +# `alg` is used as given, e.g. `sort!(A; alg=AK.RadixSort())`. `scratch` is accepted and ignored: +# the equivalent for GPU arrays is a workspace, `AK.sort!(v; workspace=AK.workspace(AK.sort!, v))`. +# +# Vector and `dims` forms are separate methods, as in Base: vectors take no `dims`, other arrays +# require it. + +# Base's sorting algorithms, as requirements on the algorithm AcceleratedKernels picks +_akalg(::Nothing) = AK.Auto() # Base's default is stable +_akalg(alg::AK.Algorithm) = alg +@static if isdefined(Base.Sort, :DefaultStable) # (Julia 1.11) + _akalg(::Base.Sort.DefaultStable) = AK.Auto(stable=true) + _akalg(::Base.Sort.DefaultUnstable) = AK.Auto(stable=false) +else # (`DEFAULT_UNSTABLE` is the same object) + _akalg(::typeof(Base.Sort.DEFAULT_STABLE)) = AK.Auto(stable=true) +end +_akalg(::Base.Sort.MergeSortAlg) = AK.Auto(stable=true) +_akalg(::Base.Sort.InsertionSortAlg) = AK.Auto(stable=true) +_akalg(::Union{Base.Sort.QuickSortAlg, Base.Sort.PartialQuickSort}) = AK.Auto(stable=false) +_akalg(alg) = throw(ArgumentError( + "sorting algorithm $alg is not supported on GPU arrays; omit `alg`, or pass one of Base's " * + "`MergeSort`, `InsertionSort`, `QuickSort` or `PartialQuickSort`, or an AcceleratedKernels algorithm")) + +# (vectors take Base's keywords only, so that `dims` is an error, as in Base) +function Base.sort!(v::AnyGPUVector; alg=nothing, lt=isless, by=identity, rev=nothing, + order::Base.Order.Ordering=Base.Order.Forward, scratch=nothing) + AK.sort!(v; alg=_akalg(alg), lt, by, rev, order) + return v +end +function Base.sort!(A::AnyGPUArray; dims::Integer, alg=nothing, scratch=nothing, kwargs...) + AK.sort!(A; dims, alg=_akalg(alg), kwargs...) + return A +end + +Base.sort(v::AnyGPUVector; kwargs...) = sort!(copy(v); kwargs...) +Base.sort(A::AnyGPUArray; dims::Integer, kwargs...) = sort!(copy(A); dims, kwargs...) + +# AcceleratedKernels always initialises `ix`; `initialized` is accepted and ignored +function Base.sortperm!(ix::AnyGPUArray{<:Integer}, v::AnyGPUVector; alg=nothing, + scratch=nothing, initialized::Bool=false, dims=nothing, kwargs...) + dims === nothing || throw(ArgumentError("sortperm! of a vector does not accept `dims`")) + axes(ix) == axes(v) || + throw(ArgumentError("index array must have the same axes as the source array")) + AK.sortperm!(ix, v; alg=_akalg(alg), kwargs...) + return ix +end +function Base.sortperm!(ix::AnyGPUArray{<:Integer}, A::AnyGPUArray; dims::Integer, alg=nothing, + scratch=nothing, initialized::Bool=false, kwargs...) + axes(ix) == axes(A) || + throw(ArgumentError("index array must have the same axes as the source array")) + AK.sortperm!(ix, A; dims, alg=_akalg(alg), kwargs...) + return ix +end + +Base.sortperm(v::AnyGPUVector; kwargs...) = sortperm!(similar(v, Int), v; kwargs...) +Base.sortperm(A::AnyGPUArray; dims::Integer, kwargs...) = + sortperm!(similar(A, Int), A; dims, kwargs...) + +# By sorting all of `v`; as Base, an integer `k` gives the element and a range a view +function Base.partialsort!(v::AnyGPUVector, k::Union{Integer, OrdinalRange}; kwargs...) + sort!(v; kwargs...) + return k isa Integer ? @allowscalar(v[k]) : view(v, k) +end +Base.partialsort!(v::AnyGPUVector, k::Union{Integer, OrdinalRange}, o::Base.Order.Ordering) = + partialsort!(v, k; order=o) + +function Base.reverse!(A::AnyGPUArray; dims=:) + AK.reverse!(A; dims) + return A +end +Base.reverse(A::AnyGPUArray; dims=:) = AK.reverse(A; dims) + +# Vectors follow Base's rules for `dims`, and reverse through the ranged form +Base.reverse!(v::AnyGPUVector; dims=:) = invoke(reverse!, Tuple{AbstractVector}, v; dims) +Base.reverse(v::AnyGPUVector; dims=:) = invoke(reverse, Tuple{AbstractVector}, v; dims) +function Base.reverse!(v::AnyGPUVector, start::Integer, stop::Integer=lastindex(v)) + s, n = Int(start), Int(stop) + n > s || return v # as Base, a trivial interval is not checked + checkbounds(v, s) + checkbounds(v, n) + AK.reverse!(view(v, s:n)) + return v +end +Base.reverse(v::AnyGPUVector, start::Integer, stop::Integer=lastindex(v)) = + reverse!(copy(v), start, stop) diff --git a/test/testsuite.jl b/test/testsuite.jl index 0d214cc8..6f8b49eb 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -68,6 +68,25 @@ function compare(@nospecialize(f), AT::Type{<:Array}, @nospecialize(xs...); kwar return true end +# Base's exact result, for rules that `compare`'s approximate comparison does not check: the +# value (with `isequal`, so signed zeros count) and type of a scalar, the element type, size and +# values of an array (or of each array of a tuple), or the type of the error Base throws +function exact_result(@nospecialize(f), @nospecialize(xs...)) + r = try + f(xs...) + catch err + return typeof(err) + end + exact(x::AbstractArray) = (eltype(x), size(x), collect(x)) + exact(x::Tuple) = map(exact, x) + exact(x) = (typeof(x), x) + return exact(r) +end +compare_exact(@nospecialize(f), AT::Type{<:AbstractGPUArray}, @nospecialize(xs...)) = + isequal(exact_result(f, map(deepcopy, xs)...), + exact_result(f, map(x -> x isa AbstractArray ? adapt(AT, x) : x, xs)...)) +compare_exact(@nospecialize(f), AT::Type{<:Array}, @nospecialize(xs...)) = true + # element types that are supported by the array type supported_eltypes(AT, test) = supported_eltypes(AT) supported_eltypes(AT) = supported_eltypes() @@ -108,6 +127,9 @@ include("testsuite/indexing.jl") include("testsuite/base.jl") include("testsuite/vector.jl") include("testsuite/reductions.jl") +include("testsuite/sorting.jl") +include("testsuite/accumulate.jl") +include("testsuite/findall.jl") include("testsuite/broadcasting.jl") include("testsuite/linalg.jl") include("testsuite/math.jl") diff --git a/test/testsuite/accumulate.jl b/test/testsuite/accumulate.jl new file mode 100644 index 00000000..d35a3352 --- /dev/null +++ b/test/testsuite/accumulate.jl @@ -0,0 +1,64 @@ +@testsuite "accumulate" (AT, eltypes)->begin + @testset "$ET" for ET in eltypes + range = ET <: Real ? (ET(1):ET(10)) : ET + sizes = ET in (Float16, ComplexF16) ? (0, 1, 10, 100) : (0, 1, 10, 1000, 100_000) + for n in sizes + @test compare(A -> cumsum(A), AT, rand(range, n)) + @test compare(A -> accumulate(+, A), AT, rand(range, n)) + @test compare(A -> accumulate(+, A; init=one(ET)), AT, rand(range, n)) + @test compare((B, A) -> accumulate!(+, B, A), AT, zeros(ET, n), rand(range, n)) + # (into the element type, which small integers overflow) + n <= 1000 && @test compare((B, A) -> cumsum!(B, A), AT, zeros(ET, n), rand(range, n)) + end + @test compare(A -> cumprod(A), AT, rand(range, 10)) + # (into the element type, whose small integers overflow a product of larger values) + @test compare((B, A) -> cumprod!(B, A), AT, zeros(ET, 10), rand(ET <: Real ? (ET(1):ET(2)) : ET, 10)) + + # along dimensions, including one beyond the array's + for dims in (1, 2, 3) + @test compare(A -> cumsum(A; dims), AT, rand(range, 10, 20)) + @test compare(A -> accumulate(+, A; dims, init=one(ET)), AT, rand(range, 10, 20)) + end + + # views and reshaped arrays + @test compare(A -> cumsum(view(A, 2:9)), AT, rand(range, 10)) + @test compare(A -> cumsum(reshape(A, 4, 5); dims=2), AT, rand(range, 20)) + end + + # Base's shape rules: an allocating scan without `dims` runs in linear order and keeps the + # shape; an in-place one needs `dims` for arrays other than vectors + M = rand(Float32, 4, 5) + @test compare(A -> accumulate(+, A), AT, M) + @test_throws ArgumentError accumulate!(+, AT(similar(M)), AT(M)) + @test_throws ArgumentError accumulate(+, AT(M); dims=0) + @test_throws TypeError accumulate(+, AT(M); dims=:) + @test_throws DimensionMismatch accumulate!(+, AT(zeros(Float32, 5, 4)), AT(M); dims=1) + @test_throws DimensionMismatch accumulate!(+, AT(zeros(Float32, 5, 4)), AT(M); dims=3) + @test_throws DimensionMismatch accumulate!(+, AT(zeros(Float32, 5, 4)), AT(M); dims=3, init=1f0) + + # An explicit `init=nothing` is an initial value + something_add(a, b) = something(a, 0) + something(b, 0) + @test compare(A -> accumulate(something_add, A; init=nothing), AT, rand(1:10, 100)) + + # Base's rules, which AcceleratedKernels leaves to GPUArrays: element types (which differ + # between Julia versions), a `dims` beyond the array's (which copies, ignoring `init`), and + # the running type Base's front-end chooses through the destination + for (f, x) in ((A -> accumulate(+, A), Int8[1, 2, 100]), (A -> accumulate(+, A; init=0), Int8[1, 2, 100]), + (cumsum, Int8[1, 2, 100]), (cumsum, Bool[1, 1, 0]), (cumprod, Int8[2, 3]), + (A -> accumulate(*, A), UInt8[16, 16, 16]), + (A -> accumulate(+, A; dims=3, init=10), rand(1:10, 3, 4)), + (A -> accumulate(+, A; dims=3), rand(1:10, 3, 4)), + (A -> cumsum(A; dims=2), rand(Int8(1):Int8(50), 3, 40))) + @test compare_exact(f, AT, x) + end + @test compare_exact((B, A) -> accumulate!(+, B, A; init=0.5f0), AT, zeros(Int32, 3), + Float32[0.5, 1.0, 2.0]) + @test compare_exact((B, A) -> accumulate!(+, B, A; init=0.5f0, dims=2), AT, + zeros(Int32, 1, 3), Float32[0.5 1.0 2.0]) + @test compare_exact((B, A) -> accumulate!(+, B, A; dims=3, init=10), AT, zeros(Int, 3, 4), + rand(1:10, 3, 4)) + + # Associative operators need not be commutative + @test compare(A -> accumulate((a, b) -> a, A), AT, rand(1:10, 1000)) + @test compare(A -> accumulate((a, b) -> a, A; dims=2), AT, rand(1:10, 10, 100)) +end diff --git a/test/testsuite/findall.jl b/test/testsuite/findall.jl new file mode 100644 index 00000000..e6d35143 --- /dev/null +++ b/test/testsuite/findall.jl @@ -0,0 +1,116 @@ +@testsuite "findall" (AT, eltypes)->begin + # indices as in Base: `Int` for vectors, `CartesianIndex` otherwise + # (`Array`, since `CartesianIndex` results are compared elementwise) + for sz in ((0,), (1,), (1000,), (10, 20), (4, 5, 6)) + @test compare(A -> Array(findall(A)), AT, rand(Bool, sz)) + @test compare(A -> Array(findall(x -> x > 0.5f0, A)), AT, rand(Float32, sz)) + end + # ... also for 0-d arrays, where only the predicate form gives linear indices + @test compare(A -> Array(findall(A)), AT, fill(true)) + @test compare(A -> Array(findall(identity, A)), AT, fill(true)) + @test compare(A -> Array(findall(A)), AT, fill(false)) + + # views and reshaped arrays + @test compare(A -> findall(view(A, 3:90)), AT, rand(Bool, 100)) + @test compare(A -> Array(findall(isodd, reshape(A, 10, 10))), AT, rand(1:10, 100)) + + @test compare(A -> findall(in((2, 3)), A), AT, rand(1:5, 100)) + + # the predicate must return a Bool, as in Base + @test_throws Union{TypeError, ArgumentError} findall(x -> 1, AT([1, 2])) + + # Base's index types exactly (AcceleratedKernels selects the items GPUArrays passes) + for x in (rand(Bool, 100), rand(Bool, 10, 20), fill(true), fill(false), Bool[]) + @test compare_exact(findall, AT, x) + @test compare_exact(A -> findall(!, A), AT, x) + end + @test compare_exact(A -> findall(isodd, view(A, 2:2:10, :)), AT, rand(1:9, 10, 3)) + for x in (fill(2), fill(5), rand(1:5, 100), rand(1:5, 4, 5)) + @test compare_exact(A -> findall(in((2, 3)), A), AT, x) + end +end + +@testsuite "indexing logical" (AT, eltypes)->begin + # a mask on the device or on the host + @test compare((A, m) -> A[m], AT, rand(Float32, 100), rand(Bool, 100)) + @test compare(A -> A[Array(A) .> 0.5f0], AT, rand(Float32, 100)) + @test compare((A, m) -> A[m], AT, rand(Float32, 10, 10), rand(Bool, 10, 10)) + # mixed with other indices + @test compare((A, m) -> A[m, :], AT, rand(Float32, 10, 5), rand(Bool, 10)) + @test compare((A, m) -> A[:, m], AT, rand(Float32, 5, 10), rand(Bool, 10)) + # assignment + @test compare((A, m) -> (A[m] .= 0; A), AT, rand(Float32, 100), rand(Bool, 100)) + # views and reshaped arrays + @test compare((A, m) -> view(A, 1:50)[m], AT, rand(Float32, 100), rand(Bool, 50)) + @test compare((A, m) -> reshape(A, 10, 10)[m], AT, rand(Float32, 100), rand(Bool, 10, 10)) + # the selected values and their type, from masks of the array's shape or of another one + for (a, m) in ((rand(Float32, 100), rand(Bool, 100)), (rand(Int8, 10, 10), rand(Bool, 10, 10)), + (rand(Float32, 10, 10), rand(Bool, 100)), (rand(Float32, 0), Bool[]), + (rand(Float32, 4, 5), falses(4, 5))) + @test compare_exact((A, m) -> A[m], AT, a, m) + end + @test compare_exact((A, m) -> view(A, 1:5, :)[m], AT, rand(Float32, 10, 4), rand(Bool, 5, 4)) + # ... and a mask that does not fit the array + for (a, m) in ((rand(Float32, 3), rand(Bool, 2)), (rand(Float32, 2, 3), rand(Bool, 3, 2))) + @test compare_exact((A, m) -> A[m], AT, a, m) + @test compare_exact(A -> A[m], AT, a) # (a host mask) + @test compare_exact(A -> A[view(m, :)], AT, a) # (a host view as the mask) + end + @test compare_exact((A, m) -> view(A, 1:2:5)[m], AT, rand(Float32, 5), rand(Bool, 2)) +end + +@testsuite "reductions/any all predicates" (AT, eltypes)->begin + for sz in ((0,), (1,), (1000,), (10, 20)) + @test compare(A -> any(A), AT, rand(Bool, sz)) + @test compare(A -> all(A), AT, rand(Bool, sz)) + @test compare(A -> any(x -> x > 0.9f0, A), AT, rand(Float32, sz)) + @test compare(A -> all(x -> x > 0.1f0, A), AT, rand(Float32, sz)) + end + @test compare(A -> all(A), AT, trues(1000)) + @test compare(A -> any(A), AT, falses(1000)) + + # along dimensions + for dims in (1, 2, (1, 2)) + @test compare(A -> any(A; dims), AT, rand(Bool, 10, 20)) + @test compare(A -> all(x -> x > 0.1f0, A; dims), AT, rand(Float32, 10, 20)) + end + + # views and reshaped arrays + @test compare(A -> any(x -> x > 0.9f0, view(A, 2:90)), AT, rand(Float32, 100)) + @test compare(A -> all(x -> x > 0.1f0, reshape(A, 10, 10); dims=2), AT, rand(Float32, 100)) + + # the predicate must return a Bool (or missing, without `dims`), as in Base, but is never + # called on an empty array + @test_throws Union{TypeError, ArgumentError} any(x -> 1, AT([1, 2])) + @test_throws Union{TypeError, ArgumentError} all(x -> 1, AT([1, 2])) + if AT <: AbstractGPUArray # (Julia 1.10's Base accepts this) + @test_throws Union{TypeError, ArgumentError} any(x -> 1, AT([1 2; 3 4]); dims=1) + end + @test any(x -> 1, AT(Int[])) === false + @test all(x -> 1, AT(Int[])) === true + @test compare(A -> any(x -> 1, A; dims=1), AT, zeros(Int, 0, 3)) + for dims in (0, -1, 1.5, (1, 0)), f in (any, all) # Base's checks of `dims`, also when empty + @test compare_exact(A -> f(A; dims), AT, zeros(Bool, 0, 3)) + end + @test compare(A -> all(x -> missing, A; dims=1), AT, zeros(Int, 0, 3)) + + # three-valued logic with missing values + for (f, xs) in ((x -> x > 1 ? missing : false, [1, 2]), (x -> x > 1 ? missing : true, [1, 2]), + (x -> x > 1 ? missing : x == 1, [1, 2, 3]), (x -> missing, [1, 2])) + @test isequal(any(f, AT(xs)), any(f, xs)) + @test isequal(all(f, AT(xs)), all(f, xs)) + end + # ... also stored in the array, where the back-end supports isbits unions + supports_unions = try + Array(AT(Union{Missing,Bool}[missing, true]) .| false) + true + catch + false + end + if supports_unions + for (a, b) in ((missing, true), (missing, false), (true, true)) + @test isequal(any(AT(Union{Missing,Bool}[a, b])), any([a, b])) + @test isequal(all(AT(Union{Missing,Bool}[a, b])), all([a, b])) + end + end +end diff --git a/test/testsuite/reductions.jl b/test/testsuite/reductions.jl index 83f39c39..c69013e1 100644 --- a/test/testsuite/reductions.jl +++ b/test/testsuite/reductions.jl @@ -233,6 +233,193 @@ end @test compare((A, B) -> isequal(A, B), AT, [missing], [missing]) end +@testsuite "reductions/contract" (AT, eltypes)->begin + M = rand(1:10, 10, 20) + + # Without `init`, `mapreducedim!` folds into the destination's values + @test compare((R, A) -> Base.mapreducedim!(identity, +, R, A), AT, rand(1:3, 1, 20), M) + @test compare((R, A) -> Base.mapreducedim!(abs2, +, R, A), AT, rand(1:3, 10), M) + + # A user `init` that is not neutral is applied once + @test compare(A -> sum(A; init=10), AT, rand(1:10, 100_000)) + @test compare(A -> sum(A; dims=1, init=10), AT, rand(1:10, 1000, 3)) + + # Operators without a known neutral element + @test compare(A -> reduce((a, b) -> a + b, A), AT, rand(1:10, 100_000)) + @test compare(A -> mapreduce(abs2, (a, b) -> max(a, b), A), AT, rand(-10:10, 1000)) + if AT <: AbstractGPUArray # (Base cannot, along `dims`) + @test Array(mapreduce(abs2, (a, b) -> a + b, AT(M); dims=2)) == sum(abs2, M; dims=2) + end + + # Tuple and named-tuple accumulators + @test compare(A -> findmin(A), AT, rand(Float32, 1000)) + @test compare(A -> findmax(A), AT, rand(Float32, 100, 10)) + @test compare((A, B) -> A == B, AT, [1, 2, 3], [1, 2, 3]) + @test compare((A, B) -> A == B, AT, [1, 2, 3], [1, 5, 3]) + + # Result types follow Base (compared with Base itself, whose rules differ between Julia + # versions): small integers widen, a scalar result has the type the fold settles on, and a + # reduction along `dims` with `init` has `init`'s type + I = Int32[1 2; 3 5] + for (red, args) in ((A -> sum(A), (Int8[100, 100],)), (A -> sum(A; dims=2), (Int8[100 100],)), + (A -> sum(A; init=Int8(0)), (Int32[1, 2],)), + (A -> sum(A; init=1.5f0), (I,)), + (A -> sum(A; dims=1, init=Int8(0)), (I,)), + (A -> count(isodd, A; init=Int8(0)), (I,)), + (A -> reduce((a, b) -> floor(Int32, a) + floor(Int32, b), A; init=0.5f0), + (Int16[1, 2],))) + cpu, gpu = red(args...), red(AT(args...)) + @test gpu isa AbstractArray ? eltype(gpu) === eltype(cpu) && Array(gpu) == cpu : + gpu === cpu + end + # An explicit `init=nothing` is an initial value + @test_throws Exception sum(AT(Int32[1, 2]); init=nothing) + something_add(a, b) = something(a, Int32(0)) + something(b, Int32(0)) + @test reduce(something_add, AT(Int32[1, 2]); init=nothing) === Int32(3) + # Base's checks of `dims` + @test_throws ArgumentError sum(AT(Int32[1, 2]); dims=0) + @test_throws ArgumentError sum(AT(Int32[1 2; 3 4]); dims=1.5) + + # Empty reductions follow Base + @test sum(AT(Int[])) === 0 + @test prod(AT(Float32[])) === 1f0 + @test sum(AT(Int[]); init=Int8(0)) === Int8(0) + @test reduce((a, b) -> a + b, AT(Int[]); init=0) === 0 + # (the same error as Base's, whose type differs between Julia versions) + errtype(f) = try f(); nothing catch err; typeof(err) end + @test errtype(() -> maximum(AT(Int[]))) === errtype(() -> maximum(Int[])) !== nothing + @test errtype(() -> reduce((a, b) -> a + b, AT(Int[]))) === + errtype(() -> reduce((a, b) -> a + b, Int[])) !== nothing + @test compare(A -> sum(A; dims=1), AT, zeros(Int, 0, 3)) + @test compare(A -> sum(A; dims=2), AT, zeros(Int, 0, 3)) + @test_throws ArgumentError maximum(AT(zeros(Int, 0, 3)); dims=1) + + # `Broadcasted` sources, and several arrays + @test compare((A, B) -> sum(Broadcast.instantiate(Broadcast.broadcasted(*, A, B))), AT, + rand(Float32, 100), rand(Float32, 100)) + @test compare((A, B) -> mapreduce(*, +, A, B), AT, rand(Float32, 100), rand(Float32, 100)) + @test compare((A, B) -> mapreduce(*, +, A, B), AT, rand(Float32, 100), rand(Float32, 50)) + @test compare((A, B) -> mapreduce(*, +, A, B; dims=1), AT, rand(Float32, 10, 10), rand(Float32, 10, 10)) + # ... including arrays without a backend, such as the indices `findfirst` reduces with + @test compare(A -> mapreduce(+, (a, b) -> a + b, A, GPUArrays.EachIndex(A)), AT, Int32[1, 2]) + @test compare(A -> mapreduce(+, +, A, GPUArrays.EachIndex(A); init=10), AT, Int32[1, 2]) +end + +@testsuite "reductions/base" (AT, eltypes)->begin + # Base's exact results (see `compare_exact`), which AcceleratedKernels leaves to GPUArrays + same(f, xs...) = compare_exact(f, AT, xs...) + + # empty inputs: Base's value, else its error + for (f, x) in ((sum, Int32[]), (prod, Int32[]), (count, Bool[]), (minimum, Int32[]), + (A -> reduce(max, A), Int32[]), (A -> reduce((a, b) -> a + b, A), Int32[]), + (A -> sum(A; init=1), Float32[]), (A -> maximum(A; init=Int32(-1)), Int32[]), + (A -> reduce(+, A; init=nothing), Int32[]), + (A -> mapreduce(x -> error("never called"), +, A; init=7), Int32[])) + @test same(f, x) + end + # one element: Base's `mapreduce_first`, or `op(init, x)` + for (f, x) in ((A -> reduce((a, b) -> a + b, A), [true]), (sum, [true]), (sum, Int8[3]), + (A -> reduce(+, A; init=Int8(0)), [true]), (A -> mapreduce(x -> x + 1, +, A), Int8[1]), + (maximum, Int8[3]), (A -> reduce(*, A; init=0.5f0), Int32[3])) + @test same(f, x) + end + # ... whose map may index device arrays + w = AT(Int32[11, 22]) + @test mapreduce(x -> w[x], +, AT(Int32[2])) === Int32(22) + # scalar result types + h8 = Int8[100, 100, 27] + for f in (sum, A -> reduce(+, A), A -> reduce(+, A; init=0), A -> sum(A; init=Int16(0)), + A -> reduce((a, b) -> Base.add_sum(a, b), A), maximum, A -> count(>(50), A)) + @test same(f, h8) + end + @test same(A -> sum(A; init=Int8(0)), [1, 2]) + if VERSION >= v"1.13-" # (Julia 1.10's Base wraps these in the small type) + @test same(A -> sum(A; init=UInt8(0)), Int8[-1, -2]) + @test same(A -> count(A; init=UInt8(0)), [true, true]) + end + + # along `dims`: `typeof(init)`, else the fold type + m8 = rand(Int8(-9):Int8(9), 40, 30) + for dims in (1, 2, (1, 2), 3) + for f in (A -> sum(A; dims), A -> sum(A; dims, init=Int16(1)), A -> maximum(A; dims), + A -> count(x -> x > 0, A; dims)) + @test same(f, m8) + end + # (Base's pairwise path reduces `Int8`s in `Int8`, which the values here do not overflow) + @test same(A -> reduce(+, A; dims, init=0.5), Int16.(m8)) + end + # ... and empty reduced dimensions: `init`, else Base's initial value or error + e8 = zeros(Int8, 0, 3) + for (f, x) in ((A -> sum(A; dims=1), e8), (A -> prod(A; dims=1), e8), + (A -> minimum(A; dims=1, init=Int8(7)), e8), (A -> minimum(A; dims=1), e8), + (A -> maximum(A; dims=1), zeros(Int8, 3, 0)), + (A -> minimum(A; dims=1), zeros(Int32, 0, 0)), + (A -> mapreduce(x -> x + 1, +, A; dims=1), zeros(Int32, 0, 2)), + (A -> mapreduce(x -> x + 1, *, A; dims=1), zeros(Int32, 0, 2)), + (A -> count(A; dims=2), zeros(Bool, 3, 0))) + @test same(f, x) + end + # ... whose map may index device arrays + wh = Int32[11, 22] + @test Array(mapreduce(x -> w[x + 1], +, AT(zeros(Int32, 0, 2)); dims=1)) == + mapreduce(x -> wh[x + 1], +, zeros(Int32, 0, 2); dims=1) + + # signed zeros: Base's sums along `dims` start from zero, whole-array sums do not + z = fill(-0.0f0, 40, 3) + for f in (sum, A -> sum(A; dims=1), A -> sum(A; dims=2), A -> sum(A; init=0.0f0), + A -> prod(A; dims=1), A -> maximum(A; dims=1), A -> reduce(+, A; dims=(1, 2))) + @test same(f, z) + end + + # the in-place reductions, with `init=true` and `false`, into destinations of another type + A = rand(1:9, 20, 30) + B = rand(Bool, 20, 30) + for init in (true, false), dims in (1, 2) + sz = dims == 1 ? (1, 30) : (20, 1) + for (f!, r, x) in ((sum!, rand(1:3, sz), A), (sum!, Float32.(rand(1:3, sz)), A), + (prod!, rand(1.0f0:2.0f0, sz), A .% 2 .+ 1), + (maximum!, rand(1:3, sz), A), (minimum!, rand(Int16(1):Int16(3), sz), A), + (any!, rand(Bool, sz), B), (all!, rand(Bool, sz), B), + (count!, rand(1:3, sz), B), + (extrema!, fill((5, 5), sz), A)) + @test same((r, x) -> f!(r, x; init), r, x) + end + @test same((r, x) -> sum!(abs2, r, x; init), rand(1:3, sz), A) + end + # ... of empty inputs + for (f!, r) in ((sum!, ones(Float32, 1, 3)), (prod!, zeros(Float32, 1, 3)), + (maximum!, zeros(Float32, 1, 3)), (count!, ones(Int, 1, 3))), init in (true, false) + @test same((r, x) -> f!(r, x; init), r, zeros(f! === count! ? Bool : Float32, 0, 3)) + end + # ... and `Base.mapreducedim!`, which folds into the destination + @test same((r, x) -> Base.mapreducedim!(abs2, +, r, x), rand(1:3, 1, 30), A) + + # findmin and findmax: Base's indices, without an `init` + for (f, x) in ((findmin, Float32[3, 1, 2]), (findmax, Float32[3, 1, 2]), + (A -> findmax(A; dims=1), Float32[3, 1, 2]), + (A -> findmin(A; dims=2), rand(Float32, 20, 30)), + (A -> findmax(A; dims=(1, 2)), rand(Float32, 20, 30)), + (A -> findmax(x -> (x > 0.5f0, -x), A), rand(Float32, 100)), + (argmin, rand(Float32, 100)), (A -> argmax(A; dims=1), rand(Float32, 20, 30)), + (findmin, Int[]), (A -> findmax(A; dims=1), zeros(Float32, 0, 3)), + (A -> findmax(A; dims=1), zeros(Float32, 3, 0)), + (A -> findmin(A; dims=1), fill(3.0f0)), (A -> findmax(A; dims=2), fill(3.0f0)), + (A -> argmin(abs, A), Int32[-4, 2, 1, -1]), (A -> argmax(abs, A), Int32[-4, 2, 1, -1]), + (A -> argmax(abs, A), Int32[])) + @test same(f, x) + end + + # 0-dimensional arrays, views and reshapes + for (f, x) in ((sum, fill(3)), (A -> sum(A; dims=1), fill(3)), (A -> mapreduce(abs2, +, A; init=1), fill(3)), + (findmax, fill(3.0f0)), + (A -> sum(view(A, 2:9, :); dims=1), rand(1:9, 10, 4)), + (A -> maximum(view(A, 1:2:9)), rand(1:9, 10)), + (A -> sum(reshape(A, 4, 5); dims=2), rand(1:9, 20))) + @test same(f, x) + end + @test same((r, x) -> sum!(view(r, 1:1, :), x), zeros(Int, 2, 4), rand(1:9, 3, 4)) +end + @testsuite "reductions/neutral_element" (AT, eltypes)->begin # GPUArrays extends GPUArraysCore's function, so every package shares one set of methods @test GPUArrays.neutral_element === GPUArrays.GPUArraysCore.neutral_element diff --git a/test/testsuite/sorting.jl b/test/testsuite/sorting.jl new file mode 100644 index 00000000..48680cf7 --- /dev/null +++ b/test/testsuite/sorting.jl @@ -0,0 +1,153 @@ +@testsuite "sorting/sort" (AT, eltypes)->begin + @testset "$ET" for ET in eltypes + ET <: Real || continue # only orderable element types + + range = ET <: AbstractFloat ? ET : (ET(1):ET(100)) + + # flat 1-D sort, in- and out-of-place, forward and reverse + for n in (0, 1, 2, 10, 1000, 100000) + @test compare(A -> sort(A), AT, rand(range, n)) + @test compare(A -> sort(A; rev=true), AT, rand(range, n)) + @test compare(A -> sort!(copy(A)), AT, rand(range, n)) + end + + # by / order + @test compare(A -> sort(A; by=abs), AT, rand(range, 1000)) + @test compare(A -> sort(A; order=Base.Order.Reverse), AT, rand(range, 1000)) + + # heavy ties + @test compare(A -> sort(A), AT, rand(ET(1):ET(3), 1000)) + + # per-slice sort along a dimension + for dims in (1, 2) + @test compare(A -> sort(A; dims), AT, rand(range, 100, 50)) + @test compare(A -> sort(A; dims, rev=true), AT, rand(range, 100, 50)) + end + @test compare(A -> sort(A; dims=2), AT, rand(range, 8, 16, 4)) + + # views and reshaped arrays + @test compare(A -> (sort!(view(A, 11:90)); A), AT, rand(range, 100)) + @test compare(A -> (sort!(view(A, 2:9, :); dims=1); A), AT, rand(range, 10, 10)) + @test compare(A -> sort(reshape(A, 10, 10); dims=2), AT, rand(range, 100)) + end +end + +@testsuite "sorting/algorithms" (AT, eltypes)->begin + x = rand(Float32, 1000) + + # Base's algorithms are requirements on the algorithm the implementation picks + for alg in (QuickSort, MergeSort, InsertionSort, Base.Sort.DEFAULT_STABLE, + Base.Sort.DEFAULT_UNSTABLE) + @test compare(A -> sort(A; alg), AT, x) + end + @test compare(A -> Array(sort(A; alg=PartialQuickSort(1:10)))[1:10], AT, x) + if AT <: AbstractGPUArray # (Base supports every algorithm of its own, and no others) + @test_throws ArgumentError sort(AT(x); alg=Base.Sort.ScratchQuickSort()) + @test_throws ArgumentError sortperm(AT(x); alg=Base.Sort.ScratchQuickSort()) + # ... or AcceleratedKernels' own + @test Array(sort(AT(x); alg=GPUArrays.AK.MergeSort())) == sort(x) + @test Array(sortperm(AT(x); alg=GPUArrays.AK.MergeSort())) == sortperm(x) + end + # `scratch` is accepted + @test compare(A -> sort(A; scratch=nothing), AT, x) + + # Stable by default, as Base: tagged ties keep their order + tagged = [(rand(1:3), i) for i in 1:1000] + @test compare(A -> sort(A; by=first), AT, tagged) + @test compare(A -> sort(A; by=first, alg=MergeSort), AT, tagged) + @test compare(A -> sortperm(A; by=first), AT, tagged) + + # 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) + @test_throws UndefKeywordError sort(AT(rand(Float32, 4, 4))) + @test_throws UndefKeywordError sort!(AT(rand(Float32, 4, 4))) + @test_throws UndefKeywordError sortperm(AT(rand(Float32, 4, 4))) +end + +@testsuite "sorting/sortperm" (AT, eltypes)->begin + @testset "$ET" for ET in eltypes + ET <: Real || continue + + range = ET <: AbstractFloat ? ET : (ET(1):ET(100)) + + for n in (1, 2, 10, 1000) + @test compare(A -> sortperm(A), AT, rand(range, n)) + @test compare(A -> sortperm(A; rev=true), AT, rand(range, n)) + end + + # along a dimension, with linear indices as in Base + for dims in (1, 2) + @test compare(A -> sortperm(A; dims), AT, rand(range, 20, 30)) + @test compare((ix, A) -> sortperm!(ix, A; dims), AT, zeros(Int, 20, 30), rand(range, 20, 30)) + end + # into an existing index array, and through views + @test compare((ix, A) -> sortperm!(ix, A), AT, zeros(Int, 100), rand(range, 100)) + @test compare((ix, A) -> sortperm!(ix, view(A, 1:50)), AT, zeros(Int, 50), rand(range, 100)) + end + + # Base's rules: vectors take no `dims`, and the index array must match + x = AT(rand(Float32, 10)) + @test_throws ArgumentError sortperm!(similar(x, Int), x; dims=1) + @test_throws ArgumentError sortperm!(similar(x, Int, 11), x) + @test_throws ArgumentError sortperm!(similar(x, Int, 4, 5), AT(rand(Float32, 5, 4)); dims=1) +end + +@testsuite "sorting/partialsort" (AT, eltypes)->begin + N = 10000 + @testset "$ET" for ET in eltypes + ET <: Real || continue + range = ET <: AbstractFloat ? ET : (ET(1):ET(100)) + + @test compare(A -> partialsort!(A, 1), AT, rand(range, N)) + @test compare(A -> partialsort!(A, N), AT, rand(range, N)) + @test compare(A -> partialsort!(A, N ÷ 2; rev=true), AT, rand(range, N)) + @test compare(A -> partialsort!(A, (N ÷ 10):(2N ÷ 10)), AT, rand(range, N)) + @test compare(A -> partialsort(A, N ÷ 2), AT, rand(range, N)) + end + + # As Base: an element for an integer, a view for a range + x = AT(rand(Float32, 100)) + @test partialsort!(copy(x), 3) isa Float32 + y = copy(x) + r = partialsort!(y, 3:5) + fill!(r, 0) + @test all(iszero, Array(y)[3:5]) + @test Array(partialsort(x, 3:5)) == partialsort(Array(x), 3:5) + # ... also with a positional ordering + @test compare(A -> partialsort!(A, 3, Base.Order.Reverse), AT, rand(Float32, 100)) +end + +@testsuite "reverse" (AT, eltypes)->begin + @testset "$ET" for ET in eltypes + # vectors: whole, and ranged + for n in (0, 1, 2, 7, 1000) + @test compare(A -> reverse(A), AT, rand(ET, n)) + @test compare(A -> reverse!(A), AT, rand(ET, n)) + end + @test compare(A -> reverse(A; dims=1), AT, rand(ET, 10)) + @test compare(A -> reverse!(A, 3), AT, rand(ET, 10)) + @test compare(A -> reverse!(A, 3, 8), AT, rand(ET, 10)) + @test compare(A -> reverse(A, 3, 8), AT, rand(ET, 10)) + + # along dimensions + for dims in (1, 2, (1, 2), :) + @test compare(A -> reverse(A; dims), AT, rand(ET, 5, 6)) + @test compare(A -> reverse!(A; dims), AT, rand(ET, 5, 6)) + end + @test compare(A -> reverse!(A; dims=(1, 3)), AT, rand(ET, 3, 4, 5)) + + # views and reshaped arrays + @test compare(A -> (reverse!(view(A, 2:9)); A), AT, rand(ET, 10)) + @test compare(A -> reverse(reshape(A, 4, 5); dims=2), AT, rand(ET, 20)) + end + + # As Base: a trivial interval is a no-op, even out of bounds; others are checked + x = AT(collect(1:10)) + @test Array(reverse!(copy(x), 7, 6)) == 1:10 + @test Array(reverse!(copy(x), 12, 11)) == 1:10 + @test_throws BoundsError reverse!(copy(x), 0, 3) + @test_throws BoundsError reverse!(copy(x), 5, 11) + @test_throws ArgumentError reverse(x; dims=2) + @test_throws ArgumentError reverse(AT(rand(Float32, 3, 3)); dims=3) +end