Generate methods for array wrappers using Adapt.jl.

Author: maleadt

src/abstractarray.jl

@@ -52,29 +52,19 @@ end
 
 ## showing
 
-for (AT, f) in
-    (GPUArray                                             => Array,
-     SubArray{<:Any,<:Any,<:GPUArray}                     => x->SubArray(Array(parent(x)), parentindices(x)),
-     LinearAlgebra.Adjoint{<:Any,<:GPUArray}              => x->LinearAlgebra.adjoint(Array(parent(x))),
-     LinearAlgebra.Transpose{<:Any,<:GPUArray}            => x->LinearAlgebra.transpose(Array(parent(x))),
-     LinearAlgebra.LowerTriangular{<:Any,<:GPUArray}      => x->LinearAlgebra.LowerTriangular(Array(x.data)),
-     LinearAlgebra.UnitLowerTriangular{<:Any,<:GPUArray}  => x->LinearAlgebra.UnitLowerTriangular(Array(x.data)),
-     LinearAlgebra.UpperTriangular{<:Any,<:GPUArray}      => x->LinearAlgebra.UpperTriangular(Array(x.data)),
-     LinearAlgebra.UnitUpperTriangular{<:Any,<:GPUArray}  => x->LinearAlgebra.UnitUpperTriangular(Array(x.data))
-    )
-  @eval begin
-    # for display
-    Base.print_array(io::IO, X::$AT) =
-        Base.print_array(io,$f(X))
-
-    # for show
-    Base._show_nonempty(io::IO, X::$AT, prefix::String) =
-        Base._show_nonempty(io,$f(X),prefix)
-    Base._show_empty(io::IO, X::$AT) =
-        Base._show_empty(io,$f(X))
-    Base.show_vector(io::IO, v::$AT, args...) =
-        Base.show_vector(io,$f(v),args...)
-  end
+for (W, ctor) in (:AT => (A,mut)->mut(A), Adapt.wrappers...)
+    @eval begin
+        # display
+        Base.print_array(io::IO, X::$W where {AT <: GPUArray}) = Base.print_array(io, $ctor(X, Array))
+
+        # show
+        Base._show_nonempty(io::IO, X::$W where {AT <: GPUArray}, prefix::String) =
+            Base._show_nonempty(io, $ctor(X, Array), prefix)
+        Base._show_empty(io::IO, X::$W where {AT <: GPUArray}) =
+            Base._show_empty(io, $ctor(X, Array))
+        Base.show_vector(io::IO, v::$W where {AT <: GPUArray}, args...) =
+            Base.show_vector(io, $ctor(v, Array), args...)
+    end
 end
 
 # memory operations

src/array.jl

@@ -145,7 +145,7 @@ end
 function LocalMemory(state::JLState, ::Type{T}, ::Val{N}, ::Val{C}) where {T, N, C}
     state.localmem_counter += 1
     lmems = state.localmems[blockidx_x(state)]
-    # first invokation in block
+    # first invocation in block
     if length(lmems) < state.localmem_counter
         lmem = fill(zero(T), N)
         push!(lmems, lmem)

src/broadcast.jl

@@ -12,24 +12,21 @@ import Base.Broadcast: BroadcastStyle, Broadcasted, ArrayStyle
 #       instead of using `ArrayStyle{GPUArray}`, due to the fact how `similar` works.
 BroadcastStyle(::Type{T}) where {T<:GPUArray} = ArrayStyle{T}()
 
-# These wrapper types otherwise forget that they are GPU compatible
+# Wrapper types otherwise forget that they are GPU compatible
 #
 # NOTE: Don't directly use ArrayStyle{GPUArray} here since that would mean that `CuArrays`
 #       customization no longer take effect.
-BroadcastStyle(::Type{<:LinearAlgebra.Transpose{<:Any,T}}) where {T<:GPUArray} = BroadcastStyle(T)
-BroadcastStyle(::Type{<:LinearAlgebra.Adjoint{<:Any,T}}) where {T<:GPUArray} = BroadcastStyle(T)
-BroadcastStyle(::Type{<:SubArray{<:Any,<:Any,T}}) where {T<:GPUArray} = BroadcastStyle(T)
-
-backend(::Type{<:LinearAlgebra.Transpose{<:Any,T}}) where {T<:GPUArray} = backend(T)
-backend(::Type{<:LinearAlgebra.Adjoint{<:Any,T}}) where {T<:GPUArray} = backend(T)
-backend(::Type{<:SubArray{<:Any,<:Any,T}}) where {T<:GPUArray} = backend(T)
+for (W, ctor) in Adapt.wrappers
+  @eval begin
+    BroadcastStyle(::Type{<:$W}) where {AT<:GPUArray} = BroadcastStyle(AT)
+    backend(::Type{<:$W}) where {AT<:GPUArray} = backend(AT)
+  end
+end
 
 # This Union is a hack. Ideally Base would have a Transpose <: WrappedArray <: AbstractArray
 # and we could define our methods in terms of Union{GPUArray, WrappedArray{<:Any, <:GPUArray}}
-const GPUDestArray = Union{GPUArray,
-                           LinearAlgebra.Transpose{<:Any,<:GPUArray},
-                           LinearAlgebra.Adjoint{<:Any,<:GPUArray},
-                           SubArray{<:Any,<:Any,<:GPUArray}}
+@eval const GPUDestArray =
+  Union{GPUArray, $((:($W where {AT <: GPUArray}) for (W, _) in Adapt.wrappers)...)}
 
 # We purposefully only specialize `copyto!`, dependent packages need to make sure that they
 # can handle: