Merge pull request #323 from JuliaGPU/tb/any_alias

Mapreduce on array wrappers

Author: maleadt

src/host/abstractarray.jl

@@ -23,9 +23,7 @@ backend(::Type{<:AbstractGPUDevice}) = error("Not implemented") # COV_EXCL_LINE
 
 const WrappedGPUArray{T,N} = WrappedArray{T,N,AbstractGPUArray,AbstractGPUArray{T,N}}
 
-const AbstractOrWrappedGPUArray{T,N} =
-    Union{AbstractGPUArray{T,N},
-          WrappedArray{T,N,AbstractGPUArray,AbstractGPUArray{T,N}}}
+const AnyGPUArray{T,N} = Union{AbstractGPUArray{T,N}, WrappedGPUArray{T,N}}
 
 
 # input/output
@@ -51,21 +49,21 @@ convert_to_cpu(xs) = adapt(Array, xs)
 ## showing
 
 # display
-Base.print_array(io::IO, X::AbstractOrWrappedGPUArray) =
+Base.print_array(io::IO, X::AnyGPUArray) =
     Base.print_array(io, convert_to_cpu(X))
 
 # show
-Base._show_nonempty(io::IO, X::AbstractOrWrappedGPUArray, prefix::String) =
+Base._show_nonempty(io::IO, X::AnyGPUArray, prefix::String) =
     Base._show_nonempty(io, convert_to_cpu(X), prefix)
-Base._show_empty(io::IO, X::AbstractOrWrappedGPUArray) =
+Base._show_empty(io::IO, X::AnyGPUArray) =
     Base._show_empty(io, convert_to_cpu(X))
-Base.show_vector(io::IO, v::AbstractOrWrappedGPUArray, args...) =
+Base.show_vector(io::IO, v::AnyGPUArray, args...) =
     Base.show_vector(io, convert_to_cpu(v), args...)
 
 ## collect to CPU (discarding wrapper type)
 
 collect_to_cpu(xs::AbstractArray) = collect(convert_to_cpu(xs))
-Base.collect(X::AbstractOrWrappedGPUArray) = collect_to_cpu(X)
+Base.collect(X::AnyGPUArray) = collect_to_cpu(X)
 
 
 # memory copying
@@ -75,9 +73,9 @@ Base.collect(X::AbstractOrWrappedGPUArray) = collect_to_cpu(X)
 # expects the GPU array type to have linear `copyto!` methods (i.e. accepting an integer
 # offset and length) from and to CPU arrays and between GPU arrays.
 
-for (D, S) in ((AbstractOrWrappedGPUArray, Array),
-               (Array, AbstractOrWrappedGPUArray),
-               (AbstractOrWrappedGPUArray, AbstractOrWrappedGPUArray))
+for (D, S) in ((AnyGPUArray, Array),
+               (Array, AnyGPUArray),
+               (AnyGPUArray, AnyGPUArray))
     @eval begin
         function Base.copyto!(dest::$D{<:Any, N}, rdest::UnitRange,
                               src::$S{<:Any, N}, ssrc::UnitRange) where {N}
@@ -112,8 +110,8 @@ function linear_copy_kernel!(ctx::AbstractKernelContext, dest, dstart, src, ssta
     return
 end
 
-function Base.copyto!(dest::AbstractOrWrappedGPUArray, dstart::Integer,
-                      src::AbstractOrWrappedGPUArray, sstart::Integer, n::Integer)
+function Base.copyto!(dest::AnyGPUArray, dstart::Integer,
+                      src::AnyGPUArray, sstart::Integer, n::Integer)
     n == 0 && return dest
     n < 0 && throw(ArgumentError(string("tried to copy n=", n, " elements, but n should be nonnegative")))
     destinds, srcinds = LinearIndices(dest), LinearIndices(src)
@@ -152,7 +150,7 @@ end
 # to quickly perform these very lightweight conversions
 
 function Base.copyto!(dest::Array{T}, dstart::Integer,
-                      src::AbstractOrWrappedGPUArray{U}, sstart::Integer,
+                      src::AnyGPUArray{U}, sstart::Integer,
                       n::Integer) where {T,U}
     n == 0 && return dest
     temp = Vector{U}(undef, n)
@@ -161,7 +159,7 @@ function Base.copyto!(dest::Array{T}, dstart::Integer,
     return dest
 end
 
-function Base.copyto!(dest::AbstractOrWrappedGPUArray{T}, dstart::Integer,
+function Base.copyto!(dest::AnyGPUArray{T}, dstart::Integer,
                       src::Array{U}, sstart::Integer, n::Integer) where {T,U}
     n == 0 && return dest
     temp = Vector{T}(undef, n)
@@ -181,8 +179,8 @@ function cartesian_copy_kernel!(ctx::AbstractKernelContext, dest, dest_offsets,
     return
 end
 
-function Base.copyto!(dest::AbstractOrWrappedGPUArray{<:Any, N}, destcrange::CartesianIndices{N},
-                      src::AbstractOrWrappedGPUArray{<:Any, N}, srccrange::CartesianIndices{N}) where {N}
+function Base.copyto!(dest::AnyGPUArray{<:Any, N}, destcrange::CartesianIndices{N},
+                      src::AnyGPUArray{<:Any, N}, srccrange::CartesianIndices{N}) where {N}
     shape = size(destcrange)
     if shape != size(srccrange)
         throw(ArgumentError("Ranges don't match their size. Found: $shape, $(size(srccrange))"))

src/host/broadcast.jl

@@ -6,7 +6,7 @@ using Base.Broadcast
 
 import Base.Broadcast: BroadcastStyle, Broadcasted, AbstractArrayStyle
 
-const BroadcastGPUArray{T} = Union{AbstractOrWrappedGPUArray{T},
+const BroadcastGPUArray{T} = Union{AnyGPUArray{T},
                                    Base.RefValue{<:AbstractGPUArray{T}}}
 
 """

src/host/construction.jl

@@ -1,6 +1,6 @@
 # constructors and conversions
 
-function Base.fill!(A::AbstractOrWrappedGPUArray{T}, x) where T
+function Base.fill!(A::AnyGPUArray{T}, x) where T
     length(A) == 0 && return A
     gpu_call(A, convert(T, x)) do ctx, a, val
         idx = @linearidx(a)
@@ -18,16 +18,16 @@ function uniformscaling_kernel(ctx::AbstractKernelContext, res::AbstractArray{T}
     return
 end
 
-function (T::Type{<: AbstractOrWrappedGPUArray{U}})(s::UniformScaling, dims::Dims{2}) where {U}
+function (T::Type{<: AnyGPUArray{U}})(s::UniformScaling, dims::Dims{2}) where {U}
     res = similar(T, dims)
     fill!(res, zero(U))
     gpu_call(uniformscaling_kernel, res, size(res, 1), s; total_threads=minimum(dims))
     res
 end
 
-(T::Type{<: AbstractOrWrappedGPUArray})(s::UniformScaling{U}, dims::Dims{2}) where U = T{U}(s, dims)
+(T::Type{<: AnyGPUArray})(s::UniformScaling{U}, dims::Dims{2}) where U = T{U}(s, dims)
 
-(T::Type{<: AbstractOrWrappedGPUArray})(s::UniformScaling, m::Integer, n::Integer) = T(s, Dims((m, n)))
+(T::Type{<: AnyGPUArray})(s::UniformScaling, m::Integer, n::Integer) = T(s, Dims((m, n)))
 
 function Base.copyto!(A::AbstractGPUMatrix{T}, s::UniformScaling) where T
     fill!(A, zero(T))

src/host/linalg.jl

@@ -99,7 +99,7 @@ end
 
 ## matrix multiplication
 
-function generic_matmatmul!(C::AbstractOrWrappedGPUArray{R}, A::AbstractOrWrappedGPUArray{T}, B::AbstractOrWrappedGPUArray{S}, a::Number, b::Number) where {T,S,R}
+function generic_matmatmul!(C::AnyGPUArray{R}, A::AnyGPUArray{T}, B::AnyGPUArray{S}, a::Number, b::Number) where {T,S,R}
     if size(A,2) != size(B,1)
         throw(DimensionMismatch("matrix A has dimensions $(size(A)), matrix B has dimensions $(size(B))"))
     end

src/host/mapreduce.jl

@@ -4,11 +4,11 @@ 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::AbstractGPUArray, A::AbstractArrayOrBroadcasted;
+mapreducedim!(f, op, R::AnyGPUArray, A::AbstractArrayOrBroadcasted;
               init=nothing) = error("Not implemented") # COV_EXCL_LINE
 # resolve ambiguities
-Base.mapreducedim!(f, op, R::AbstractGPUArray, A::AbstractArray) = mapreducedim!(f, op, R, A)
-Base.mapreducedim!(f, op, R::AbstractGPUArray, A::Broadcast.Broadcasted) = mapreducedim!(f, op, R, A)
+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)
 
 neutral_element(op, T) =
     error("""GPUArrays.jl needs to know the neutral element for your operator `$op`.
@@ -24,7 +24,7 @@ neutral_element(::typeof(Base.min), T) = typemax(T)
 neutral_element(::typeof(Base.max), T) = typemin(T)
 
 # resolve ambiguities
-Base.mapreduce(f, op, A::AbstractGPUArray, As::AbstractArrayOrBroadcasted...;
+Base.mapreduce(f, op, A::AnyGPUArray, As::AbstractArrayOrBroadcasted...;
                dims=:, init=nothing) = _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)
@@ -68,24 +68,24 @@ function _mapreduce(f::F, op::OP, As::Vararg{Any,N}; dims::D, init) where {F,OP,
     end
 end
 
-Base.any(A::AbstractGPUArray{Bool}) = mapreduce(identity, |, A)
-Base.all(A::AbstractGPUArray{Bool}) = mapreduce(identity, &, A)
+Base.any(A::AnyGPUArray{Bool}) = mapreduce(identity, |, A)
+Base.all(A::AnyGPUArray{Bool}) = mapreduce(identity, &, A)
 
-Base.any(f::Function, A::AbstractGPUArray) = mapreduce(f, |, A)
-Base.all(f::Function, A::AbstractGPUArray) = mapreduce(f, &, A)
+Base.any(f::Function, A::AnyGPUArray) = mapreduce(f, |, A)
+Base.all(f::Function, A::AnyGPUArray) = mapreduce(f, &, A)
 
-Base.count(pred::Function, A::AbstractGPUArray; dims=:) =
+Base.count(pred::Function, A::AnyGPUArray; dims=:) =
     mapreduce(pred, Base.add_sum, A; init=0, dims=dims)
 
-Base.:(==)(A::AbstractGPUArray, B::AbstractGPUArray) = Bool(mapreduce(==, &, A, B))
+Base.:(==)(A::AnyGPUArray, B::AnyGPUArray) = Bool(mapreduce(==, &, A, B))
 
 # avoid calling into `initarray!`
 for (fname, op) in [(:sum, :(Base.add_sum)), (:prod, :(Base.mul_prod)),
                     (:maximum, :(Base.max)), (:minimum, :(Base.min)),
                     (:all, :&),              (:any, :|)]
     fname! = Symbol(fname, '!')
     @eval begin
-        Base.$(fname!)(f::Function, r::AbstractGPUArray, A::AbstractGPUArray{T}) where T =
+        Base.$(fname!)(f::Function, r::AnyGPUArray, A::AnyGPUArray{T}) where T =
             GPUArrays.mapreducedim!(f, $(op), r, A; init=neutral_element($(op), T))
     end
 end

src/host/math.jl

@@ -1,6 +1,6 @@
 # Base mathematical operations
 
-function Base.clamp!(A::AbstractOrWrappedGPUArray, low, high)
+function Base.clamp!(A::AnyGPUArray, low, high)
     gpu_call(A, low, high) do ctx, A, low, high
         I = @cartesianidx A
         A[I] = clamp(A[I], low, high)

src/host/random.jl

@@ -67,7 +67,7 @@ struct RNG <: AbstractRNG
 end
 
 # return an instance of GPUArrays.RNG suitable for the requested array type
-default_rng(::Type{<:AbstractGPUArray}) = error("Not implemented") # COV_EXCL_LINE
+default_rng(::Type{<:AnyGPUArray}) = error("Not implemented") # COV_EXCL_LINE
 
 make_seed(rng::RNG) = make_seed(rng, rand(UInt))
 function make_seed(rng::RNG, n::Integer)
@@ -81,7 +81,7 @@ function Random.seed!(rng::RNG, seed::Vector{UInt32})
     return
 end
 
-function Random.rand!(rng::RNG, A::AbstractGPUArray{T}) where T <: Number
+function Random.rand!(rng::RNG, A::AnyGPUArray{T}) where T <: Number
     gpu_call(A, rng.state) do ctx, a, randstates
         idx = linear_index(ctx)
         idx > length(a) && return
@@ -91,7 +91,7 @@ function Random.rand!(rng::RNG, A::AbstractGPUArray{T}) where T <: Number
     A
 end
 
-function Random.randn!(rng::RNG, A::AbstractGPUArray{T}) where T <: Number
+function Random.randn!(rng::RNG, A::AnyGPUArray{T}) where T <: Number
     threads = (length(A) - 1) ÷ 2 + 1
     length(A) == 0 && return
     gpu_call(A, rng.state; total_threads = threads) do ctx, a, randstates