Use CPU memcpy for SharedStorage copyto! (#445)

Author: christiangnrd

src/array.jl

@@ -401,6 +401,12 @@ function Base.unsafe_copyto!(dev::MTLDevice, dest::MtlArray{T}, doffs, src::Arra
     end
     return dest
 end
+function Base.unsafe_copyto!(::MTLDevice, dest::MtlArray{T,<:Any,Metal.SharedStorage}, doffs, src::Array{T}, soffs, n) where T
+    # these copies are implemented using pure memcpy's, not API calls, so aren't ordered.
+    synchronize()
+    GC.@preserve src dest unsafe_copyto!(pointer(unsafe_wrap(Array,dest), doffs), pointer(src, soffs), n)
+    return dest
+end
 
 # GPU -> CPU
 function Base.unsafe_copyto!(dev::MTLDevice, dest::Array{T}, doffs, src::MtlArray{T}, soffs, n) where T
@@ -414,6 +420,12 @@ function Base.unsafe_copyto!(dev::MTLDevice, dest::Array{T}, doffs, src::MtlArra
     end
     return dest
 end
+function Base.unsafe_copyto!(::MTLDevice, dest::Array{T}, doffs, src::MtlArray{T,<:Any,Metal.SharedStorage}, soffs, n) where T
+    # these copies are implemented using pure memcpy's, not API calls, so aren't ordered.
+    synchronize()
+    GC.@preserve src dest unsafe_copyto!(pointer(dest, doffs), pointer(unsafe_wrap(Array,src), soffs), n)
+    return dest
+end
 
 # GPU -> GPU
 function Base.unsafe_copyto!(dev::MTLDevice, dest::MtlArray{T}, doffs, src::MtlArray{T}, soffs, n) where T
@@ -427,6 +439,12 @@ function Base.unsafe_copyto!(dev::MTLDevice, dest::MtlArray{T}, doffs, src::MtlA
     end
     return dest
 end
+function Base.unsafe_copyto!(::MTLDevice, dest::MtlArray{T,<:Any,Metal.SharedStorage}, doffs, src::MtlArray{T,<:Any,Metal.SharedStorage}, soffs, n) where T
+    # these copies are implemented using pure memcpy's, not API calls, so aren't ordered.
+    synchronize()
+    GC.@preserve src dest unsafe_copyto!(pointer(unsafe_wrap(Array,dest), doffs), pointer(unsafe_wrap(Array,src), soffs), n)
+    return dest
+end
 
 
 ## regular gpu array adaptor

test/array.jl

@@ -69,6 +69,30 @@ end
     @test collect(Metal.fill(1, 2, 2)) == ones(Float32, 2, 2)
 end
 
+@testset "copyto!" begin
+    @testset "$T, $S" for S in [Metal.PrivateStorage, Metal.SharedStorage],
+                          T in [Float16, Float32, Bool, Int16, Int32, Int64, Int8, UInt16, UInt32, UInt64, UInt8]
+        dim = (1000,17,10)
+        A = rand(T,dim)
+        mtlA = mtl(A;storage=S)
+
+        #cpu -> gpu
+        res = Metal.zeros(T,dim;storage=S)
+        copyto!(res,A)
+        @test Array(res) == Array(A)
+
+        #gpu -> cpu
+        res = zeros(T,dim)
+        copyto!(res,mtlA)
+        @test Array(res) == Array(mtlA)
+
+        #gpu -> gpu
+        res = Metal.zeros(T,dim;storage=S)
+        copyto!(res,mtlA)
+        @test Array(res) == Array(mtlA)
+    end
+end
+
 check_storagemode(arr, smode) = Metal.storagemode(arr) == smode
 
 # There is some repetition to the GPUArrays tests to test for different storagemodes