Split out copyto for texture arrays and add more tests (#2719)

Author: kshyatt

src/texture.jl

@@ -104,29 +104,65 @@ CuTextureArray(A::AbstractArray{T,N}) where {T,N} = CuTextureArray{T,N}(A)
 
 ## memory operations
 
-function Base.copyto!(dst::CuTextureArray{T,1}, src::Union{Array{T,1}, CuArray{T,1}}) where {T}
+function Base.copyto!(dst::CuTextureArray{T,1}, src::Union{Array{T,1}, CuTextureArray{T,1}}) where {T}
     size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
     Base.unsafe_copyto!(pointer(dst), pointer(src), length(dst))
     return dst
 end
 
-function Base.copyto!(dst::CuTextureArray{T,2}, src::Union{Array{T,2}, CuArray{T,2}}) where {T}
+function Base.copyto!(dst::CuTextureArray{T,1}, src::CuArray{T,1,M}) where {T, M}
     size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
-    unsafe_copy2d!(pointer(dst), ArrayMemory,
-                   pointer(src), isa(src, Array) ? HostMemory : DeviceMemory,
-                   size(dst)...)
+    if M <: Union{HostMemory, DeviceMemory}
+        Base.unsafe_copyto!(pointer(dst), pointer(src; type=M), length(dst))
+    else
+        Base.unsafe_copyto!(pointer(dst), pointer(src), length(dst))
+    end
     return dst
 end
 
-function Base.copyto!(dst::CuTextureArray{T,3}, src::Union{Array{T,3}, CuArray{T,3}}) where {T}
+function Base.copyto!(dst::CuTextureArray{T,2}, src::Array{T,2}) where {T}
     size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
-    unsafe_copy3d!(pointer(dst), ArrayMemory,
-                   pointer(src), isa(src, Array) ? HostMemory : DeviceMemory,
-                   size(dst)...)
+    unsafe_copy2d!(pointer(dst), ArrayMemory, pointer(src), HostMemory, size(dst)...)
     return dst
 end
 
+function Base.copyto!(dst::CuTextureArray{T,2}, src::CuArray{T,2,M}) where {T, M}
+    size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
+    if M <: Union{HostMemory, DeviceMemory}
+        unsafe_copy2d!(pointer(dst), ArrayMemory, pointer(src; type=M), M, size(dst)...)
+    else
+        unsafe_copy2d!(pointer(dst), ArrayMemory, pointer(src), M, size(dst)...)
+    end
+    return dst
+end
 
+function Base.copyto!(dst::CuTextureArray{T,2}, src::CuTextureArray{T,2}) where {T}
+    size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
+    unsafe_copy2d!(pointer(dst), ArrayMemory, pointer(src), ArrayMemory, size(dst)...)
+    return dst
+end
+
+function Base.copyto!(dst::CuTextureArray{T,3}, src::Array{T,3}) where {T}
+    size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
+    unsafe_copy3d!(pointer(dst), ArrayMemory, pointer(src), HostMemory, size(dst)...)
+    return dst
+end
+
+function Base.copyto!(dst::CuTextureArray{T,3}, src::CuArray{T,3,M}) where {T, M}
+    size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
+    if M <: Union{HostMemory, DeviceMemory}
+        unsafe_copy3d!(pointer(dst), ArrayMemory, pointer(src; type=M), M, size(dst)...)
+    else
+        unsafe_copy3d!(pointer(dst), ArrayMemory, pointer(src), M, size(dst)...)
+    end
+    return dst
+end
+
+function Base.copyto!(dst::CuTextureArray{T,3}, src::CuTextureArray{T,3}) where {T}
+    size(dst) == size(src) || throw(DimensionMismatch("source and destination sizes must match"))
+    unsafe_copy3d!(pointer(dst), ArrayMemory, pointer(src), ArrayMemory, size(dst)...)
+    return dst
+end
 
 #
 # Texture objects

test/base/texture.jl

@@ -82,14 +82,29 @@ end
     copyto!(texarr1D, a1D)
     tex1D = CuTexture(texarr1D)
     @test Array(fetch_all(tex1D)) == a1D
+    @test sizeof(texarr1D) == sizeof(a1D)
+    @test eltype(texarr1D) == Float32
+    h_arr_1D = zeros(Float32, length(a1D))
+    Base.unsafe_copyto!(pointer(h_arr_1D), pointer(texarr1D), length(h_arr_1D))
+    @test h_arr_1D == a1D
+    cu_arr_1D = CUDA.zeros(Float32, length(a1D))
+    Base.unsafe_copyto!(pointer(cu_arr_1D), pointer(texarr1D), length(cu_arr_1D))
+    @test Array(cu_arr_1D) == a1D
 
     texarr2D = CuTextureArray(a2D)
     tex2D = CuTexture(texarr2D)
     @test Array(fetch_all(tex2D)) == a2D
     texarr2D_2 = CuTextureArray(texarr2D)
     tex2D_2 = CuTexture(texarr2D_2)
     @test Array(fetch_all(tex2D_2)) == a2D
-
+    texarr2D_3 = CuTextureArray{Float32, 2}(texarr2D)
+    tex2D_2 = CuTexture(texarr2D_3)
+    @test Array(fetch_all(tex2D_2)) == a2D
+    copyto!(texarr2D_3, texarr2D_2)
+    tex2D_2 = CuTexture(texarr2D_3)
+    @test Array(fetch_all(tex2D_2)) == a2D
+    @test sizeof(texarr2D) == sizeof(a2D)
+    @test eltype(texarr2D) == Float32
 
     tex2D_dir = CuTexture(CuTextureArray(a2D))
     @test Array(fetch_all(tex2D_dir)) == a2D
@@ -100,6 +115,12 @@ end
     texarr3D_2 = CuTextureArray(texarr3D)
     tex3D_2 = CuTexture(texarr3D_2)
     @test Array(fetch_all(tex3D_2)) == a3D
+    texarr3D_3 = CuTextureArray{Float32, 3}(texarr3D)
+    copyto!(texarr2D_3, texarr2D_2)
+    tex3D_3 = CuTexture(texarr3D_3)
+    @test Array(fetch_all(tex3D_3)) == a3D
+    @test sizeof(texarr3D) == sizeof(a3D)
+    @test eltype(texarr3D) == Float32
 end
 
 @testset "CuTexture(::CuArray)" begin