Fix an ambiguity in copyto!. (#555)

JuliaGPU · Aug 26, 2024 · 13ee44c · 13ee44c
1 parent 60cdac7
commit 13ee44c
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 17 deletions.
diff --git a/src/host/abstractarray.jl b/src/host/abstractarray.jl
@@ -134,17 +134,15 @@ Base.collect(X::AnyGPUArray) = collect_to_cpu(X)
 
 # memory copying
 
+# 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.
+
 function Base.copy!(dst::AbstractGPUVector, src::AbstractGPUVector)
     axes(dst) == axes(src) || throw(ArgumentError(
     "arrays must have the same axes for `copy!`. consider using `copyto!` instead"))
     copyto!(dst, src)
 end
 
-## basic linear copies of identically-typed memory
-
-# 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 ((AnyGPUArray, Array),
                (Array, AnyGPUArray),
                (AnyGPUArray, AnyGPUArray))
@@ -156,18 +154,6 @@ for (D, S) in ((AnyGPUArray, Array),
             copyto!(dest, drange, src, srange)
         end
 
-        function Base.copyto!(dest::$D, d_range::CartesianIndices{1},
-                              src::$S, s_range::CartesianIndices{1})
-            len = length(d_range)
-            if length(s_range) != len
-                throw(ArgumentError("Copy range needs same length. Found: dest: $len, src: $(length(s_range))"))
-            end
-            len == 0 && return dest
-            d_offset = first(d_range)[1]
-            s_offset = first(s_range)[1]
-            copyto!(dest, d_offset, src, s_offset, len)
-        end
-
         Base.copyto!(dest::$D, src::$S) = copyto!(dest, 1, src, 1, length(src))
     end
 end
@@ -260,6 +246,13 @@ function Base.copyto!(dest::AnyGPUArray{<:Any, N}, destcrange::CartesianIndices{
     len = length(destcrange)
     len == 0 && return dest
 
+    # linear copy if we can
+    if N == 1
+        d_offset = first(destcrange)[1]
+        s_offset = first(srccrange)[1]
+        return copyto!(dest, d_offset, src, s_offset, len)
+    end
+
     dest_offsets = first(destcrange) - oneunit(CartesianIndex{N})
     src_offsets = first(srccrange) - oneunit(CartesianIndex{N})
     gpu_call(cartesian_copy_kernel!,
@@ -275,6 +268,15 @@ for (dstTyp, srcTyp) in (AbstractGPUArray=>Array, Array=>AbstractGPUArray)
         if size(dstrange) != size(srcrange)
             throw(ArgumentError("source and destination must have same size (got $(size(srcrange)) and $(size(dstrange)))"))
         end
+        len = length(dstrange)
+        len == 0 && return dest
+
+        # linear copy if we can
+        if N == 1
+            d_offset = first(dstrange)[1]
+            s_offset = first(srcrange)[1]
+            return copyto!(dst, d_offset, src, s_offset, len)
+        end
 
         # figure out how many dimensions of the Cartesian ranges map onto contiguous memory
         # in both source and destination. we will copy these one by one as linear ranges.

diff --git a/test/testsuite/base.jl b/test/testsuite/base.jl
@@ -40,6 +40,23 @@ end
     end
 
     @testset "copyto!" begin
+        x = fill(0f0, 1)
+        y = rand(Float32, 1)
+        a = AT(x)
+        copyto!(x, 1:1, y, 1:1)
+        copyto!(a, 1:1, y, 1:1)
+        @test x == Array(a)
+
+        x = fill(0f0, 10)
+        y = rand(Float32, 20)
+        a = AT(x)
+        b = AT(y)
+        r1 = 1:7
+        r2 = 11:17
+        copyto!(x, r1, y, r2)
+        copyto!(a, r1, b, r2)
+        @test x == Array(a)
+
         x = fill(0f0, (10, 10))
         y = rand(Float32, (20, 10))
         a = AT(x)