Implement sorting, reductions, scans and findall with AcceleratedKernels

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"

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 <tim.besard@gmail.com>"]
 
 [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"

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")

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

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

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

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

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")

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

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,36 +237,82 @@ 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
 
         binop(x, y) && return t2
         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]]

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