Use GPU blit for large SharedStorage GPU→GPU copies (>32MB) (#716)

Author: KaanKesginLW

src/array.jl

@@ -440,10 +440,19 @@ 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.
+function Base.unsafe_copyto!(dev::MTLDevice, dest::MtlArray{T, <:Any, Metal.SharedStorage}, doffs, src::MtlArray{T, <:Any, Metal.SharedStorage}, soffs, n) where {T}
     synchronize()
-    GC.@preserve src dest unsafe_copyto!(pointer(unsafe_wrap(Array,dest), doffs), pointer(unsafe_wrap(Array,src), soffs), n)
+    bytes = n * sizeof(T)
+    # Use GPU blit for large copies (>32MiB) where it's faster than CPU memcpy.
+    # For small copies, CPU memcpy avoids GPU command buffer overhead.
+    if bytes >= 32 * 2^20 # If changed, also change in tests
+        GC.@preserve src dest unsafe_copyto!(dev, pointer(dest, doffs), pointer(src, soffs), n)
+        if Base.isbitsunion(T)
+            error("Not implemented")
+        end
+    else
+        GC.@preserve src dest unsafe_copyto!(pointer(unsafe_wrap(Array, dest), doffs), pointer(unsafe_wrap(Array, src), soffs), n)
+    end
     return dest
 end
 

test/array.jl

@@ -69,25 +69,33 @@ end
 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)
+    @testset "$S" for S in [Metal.PrivateStorage, Metal.SharedStorage]
+        @testset "$T" for 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
 
-        #gpu -> gpu
-        res = Metal.zeros(T,dim;storage=S)
-        copyto!(res,mtlA)
+        # Large array, only test Float32
+        A = rand(Float32, 32 * 2^20)
+        mtlA = mtl(A; storage = S)
+        res = similar(A)
+        copyto!(res, mtlA)
         @test Array(res) == Array(mtlA)
     end
 end