Introduce AsyncNumber to lazily copy numeric mapreduce results to the host

Project.toml

@@ -3,6 +3,7 @@ uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 version = "10.2.3"
 
 [deps]
+AbstractNumbers = "85c772de-338a-5e7f-b815-41e76c26ac1f"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
@@ -14,6 +15,7 @@ Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+AbstractNumbers = "0.2"
 Adapt = "4.0"
 GPUArraysCore = "= 0.1.6"
 LLVM = "3.9, 4, 5, 6, 7, 8"

src/GPUArrays.jl

@@ -1,5 +1,7 @@
 module GPUArrays
 
+import AbstractNumbers as AN
+
 using Serialization
 using Random
 using LinearAlgebra
@@ -25,6 +27,7 @@ include("device/synchronization.jl")
 # host abstractions
 include("host/abstractarray.jl")
 include("host/construction.jl")
+include("host/gpunumber.jl")
 ## integrations and specialized methods
 include("host/base.jl")
 include("host/indexing.jl")

src/host/gpunumber.jl

@@ -0,0 +1,42 @@
+# Custom GPU-compatible `Number` interface.
+struct GPUNumber{T <: AbstractGPUArray} <: AN.AbstractNumber{T}
+    val::T
+
+    function GPUNumber(val::T) where T <: AbstractGPUArray
+        length(val) != 1 && error(
+            "`GPUNumber` accepts only 1-element GPU arrays, " *
+            "instead `$(length(val))`-element array was given.")
+        new{T}(val)
+    end
+end
+
+AN.number(g::GPUNumber) = @allowscalar g.val[]
+maybe_number(g::GPUNumber) = AN.number(g)
+maybe_number(g) = g
+
+number_type(::GPUNumber{T}) where T = eltype(T)
+
+# When operations involve other `::Number` types,
+# do not convert back to `GPUNumber`.
+AN.like(::Type{<: GPUNumber}, x) = x
+
+# When broadcasting, just pass the array itself.
+Base.broadcastable(g::GPUNumber) = g.val
+
+# Overload to avoid copies.
+Base.one(g::GPUNumber) = one(number_type(g))
+Base.one(::Type{GPUNumber{T}}) where T = one(eltype(T))
+Base.zero(g::GPUNumber) = zero(number_type(g))
+Base.zero(::Type{GPUNumber{T}}) where T = zero(eltype(T))
+Base.identity(g::GPUNumber) = g
+
+Base.getindex(g::GPUNumber) = AN.number(g)
+
+Base.isequal(g::GPUNumber, v::Number) = isequal(g[], v)
+Base.isequal(v::Number, g::GPUNumber) = isequal(v, g[])
+
+Base.nextpow(a, x::GPUNumber) = nextpow(a, x[])
+Base.nextpow(a::GPUNumber, x) = nextpow(a[], x)
+Base.nextpow(a::GPUNumber, x::GPUNumber) = nextpow(a[], x[])
+
+Base.convert(::Type{Number}, g::GPUNumber) = g[]

src/host/indexing.jl

@@ -204,7 +204,7 @@ function Base.findfirst(f::Function, A::AnyGPUArray)
     end
 
     res = mapreduce((x, y)->(f(x), y), reduction, A, indices;
-                    init = (false, dummy_index))
+                    init = (false, dummy_index)) |> AN.number
     if res[1]
         # out of consistency with Base.findarray, return a CartesianIndex
         # when the input is a multidimensional array
@@ -230,14 +230,15 @@ function findminmax(binop, A::AnyGPUArray; init, dims)
     end
 
     if dims == Colon()
-        res = mapreduce(tuple, reduction, A, indices; init = (init, dummy_index))
+        res = mapreduce(tuple, reduction, A, indices;
+            init = (init, dummy_index)) |> AN.number
 
         # 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]])
     else
         res = mapreduce(tuple, reduction, A, indices;
-                        init = (init, dummy_index), dims=dims)
+                        init = (init, dummy_index), dims=dims) |> maybe_number
         vals = map(x->x[1], res)
         inds = map(x->ndims(A) == 1 ? x[2] : CartesianIndices(A)[x[2]], res)
         return (vals, inds)

src/host/mapreduce.jl

@@ -68,20 +68,20 @@ function _mapreduce(f::F, op::OP, As::Vararg{Any,N}; dims::D, init) where {F,OP,
     end
 
     if dims === Colon()
-        @allowscalar R[]
+        GPUNumber(R)
     else
         R
     end
 end
 
-Base.any(A::AnyGPUArray{Bool}) = mapreduce(identity, |, A)
-Base.all(A::AnyGPUArray{Bool}) = mapreduce(identity, &, A)
+Base.any(A::AnyGPUArray{Bool}) = mapreduce(identity, |, A) |> AN.number
+Base.all(A::AnyGPUArray{Bool}) = mapreduce(identity, &, A) |> AN.number
 
-Base.any(f::Function, A::AnyGPUArray) = mapreduce(f, |, A)
-Base.all(f::Function, A::AnyGPUArray) = mapreduce(f, &, A)
+Base.any(f::Function, A::AnyGPUArray) = mapreduce(f, |, A) |> AN.number
+Base.all(f::Function, A::AnyGPUArray) = mapreduce(f, &, A) |> AN.number
 
 Base.count(pred::Function, A::AnyGPUArray; dims=:, init=0) =
-    mapreduce(pred, Base.add_sum, A; init=init, dims=dims)
+    mapreduce(pred, Base.add_sum, A; init=init, dims=dims) |> maybe_number
 
 # avoid calling into `initarray!`
 for (fname, op) in [(:sum, :(Base.add_sum)), (:prod, :(Base.mul_prod)),
@@ -94,7 +94,8 @@ for (fname, op) in [(:sum, :(Base.add_sum)), (:prod, :(Base.mul_prod)),
     end
 end
 
-LinearAlgebra.ishermitian(A::AbstractGPUMatrix) = mapreduce(==, &, A, adjoint(A))
+LinearAlgebra.ishermitian(A::AbstractGPUMatrix) =
+    mapreduce(==, &, A, adjoint(A)) |> AN.number
 
 
 # comparisons
@@ -105,7 +106,7 @@ function Base.isequal(A::AnyGPUArray, B::AnyGPUArray)
     if axes(A) != axes(B)
         return false
     end
-    mapreduce(isequal, &, A, B; init=true)
+    mapreduce(isequal, &, A, B; init=true) |> AN.number
 end
 
 # returns `missing` when missing values are involved
@@ -129,6 +130,7 @@ function Base.:(==)(A::AnyGPUArray, B::AnyGPUArray)
             (; is_missing=false, is_equal=a.is_equal & b.is_equal)
         end
     end
-    res = mapreduce(mapper, reducer, A, B; init=(; is_missing=false, is_equal=true))
+    res = mapreduce(mapper, reducer, A, B;
+        init=(; is_missing=false, is_equal=true)) |> AN.number
     res.is_missing ? missing : res.is_equal
 end

test/testsuite.jl

@@ -34,6 +34,9 @@ function test_result(as::NTuple{N,Any}, bs::NTuple{N,Any}; kwargs...) where {N}
         test_result(a, b; kwargs...)
     end
 end
+# Special case for `extrema` accross all dims.
+test_result(as::NTuple{N,Any}, bs::GPUArrays.GPUNumber; kwargs...) where {N} =
+    test_result(as, GPUArrays.maybe_number(bs))
 
 function compare(f, AT::Type{<:AbstractGPUArray}, xs...; kwargs...)
     # copy on the CPU, adapt on the GPU, but keep Ref's