Merge pull request #172 from JuliaGPU/tb/wrappers

Use Adapt.jl to generate array wrapper methods

Author: maleadt

Project.toml

@@ -2,6 +2,7 @@ name = "GPUArrays"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

REQUIRE

@@ -2,3 +2,4 @@ julia 1.0
 StaticArrays
 FFTW
 FillArrays 0.3
+Adapt 0.4.1

src/GPUArrays.jl

@@ -14,6 +14,8 @@ using Base.Cartesian
 
 using FFTW
 
+using Adapt
+
 include("abstractarray.jl")
 include("abstract_gpu_interface.jl")
 include("ondevice.jl")

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:

test/runtests.jl

@@ -1,36 +1,29 @@
-using GPUArrays, Test
-using GPUArrays.TestSuite
-
+# GPUArrays development often happens in lockstep with other packages, so try to match branches
 using Pkg
-
-@testset "JLArray" begin
-    GPUArrays.test(JLArray)
+function match_package(package, var)
+    try
+        branch = ENV[var]
+        Pkg.add(PackageSpec(name=package, rev=branch))
+        @info "Installed $package from branch $branch"
+    catch ex
+        @warn "Could not install $package from $branch branch, trying master" exception=ex
+        Pkg.add(PackageSpec(name=package, rev="master"))
+        @info "Installed $package from master"
+    end
 end
+haskey(ENV, "TRAVIS")    && match_package("Adapt", "TRAVIS_PULL_REQUEST_BRANCH")
+haskey(ENV, "APPVEYOR")  && match_package("Adapt", "APPVEYOR_PULL_REQUEST_HEAD_REPO_BRANCH")
+haskey(ENV, "GITLAB_CI") && match_package("Adapt", "CI_COMMIT_REF_NAME")
 
-function test_package(package, branch=nothing)
-    mktempdir() do devdir
-        withenv("JULIA_PKG_DEVDIR" => devdir) do
-            # try to install from the same branch of GPUArrays
-            try
-                if branch === nothing
-                    branch = chomp(read(`git -C $(@__DIR__) rev-parse --abbrev-ref HEAD`, String))
-                    branch == "HEAD" && error("in detached HEAD state")
-                end
-                Pkg.add(PackageSpec(name=package, rev=String(branch)))
-                @info "Installed $package from $branch branch"
-            catch ex
-                @warn "Could not install $package from same branch as GPUArrays, trying master branch" exception=ex
-                Pkg.add(PackageSpec(name=package, rev="master"))
-            end
+using GPUArrays, Test
 
-            Pkg.test(package)
-        end
-    end
+@testset "JLArray" begin
+    GPUArrays.test(JLArray)
 end
 
 if haskey(ENV, "GITLAB_CI")
-    branch = ENV["CI_COMMIT_REF_NAME"]
+    match_package("CuArrays", "CI_COMMIT_REF_NAME")
     @testset "CuArray" begin
-        test_package("CuArrays", branch)
+        Pkg.test("CuArrays")
     end
 end