Simplify memcpy handling (#897)

Author: pxl-th

src/array.jl

@@ -194,8 +194,6 @@ Base.convert(::Type{T}, x::T) where T <: ROCArray = x
 
 ## memory operations
 
-# TODO rework, to pass pointers, instead of accessing .mem
-
 function Base.copyto!(
     dest::Array{T}, d_offset::Integer,
     source::ROCArray{T}, s_offset::Integer, amount::Integer;
@@ -204,16 +202,10 @@ function Base.copyto!(
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset + amount - 1)
     @boundscheck checkbounds(source, s_offset + amount - 1)
-
-    @debug "[gpu -> cpu] T=$T, shape=$(size(dest))"
     stm = stream()
-    Mem.download!(
-        pointer(dest, d_offset),
-        Mem.view(convert(Mem.AbstractAMDBuffer, source.buf[]),
-            (source.offset + s_offset - 1) * aligned_sizeof(T)),
-        amount * aligned_sizeof(T); stream=stm)
+    Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stm)
     async || synchronize(stm)
-    dest
+    return dest
 end
 
 function Base.copyto!(
@@ -223,13 +215,8 @@ function Base.copyto!(
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset + amount - 1)
     @boundscheck checkbounds(source, s_offset + amount - 1)
-
-    @debug "[cpu -> gpu] T=$T, shape=$(size(dest))"
-    Mem.upload!(
-        Mem.view(convert(Mem.AbstractAMDBuffer, dest.buf[]),
-            (dest.offset + d_offset - 1) * aligned_sizeof(T)),
-        pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream())
-    dest
+    Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream())
+    return dest
 end
 
 function Base.copyto!(
@@ -239,19 +226,14 @@ function Base.copyto!(
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset + amount - 1)
     @boundscheck checkbounds(source, s_offset + amount - 1)
-    Mem.transfer!(
-        Mem.view(convert(Mem.AbstractAMDBuffer, dest.buf[]),
-            (dest.offset + d_offset - 1) * aligned_sizeof(T)),
-        Mem.view(convert(Mem.AbstractAMDBuffer, source.buf[]),
-            (source.offset + s_offset - 1) * aligned_sizeof(T)),
-        amount * aligned_sizeof(T); stream=stream())
-    dest
+    Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream())
+    return dest
 end
 
 function Base.copy(X::ROCArray{T}) where T
     Xnew = ROCArray{T}(undef, size(X))
     copyto!(Xnew, 1, X, 1, length(X))
-    Xnew
+    return Xnew
 end
 
 function Base.unsafe_wrap(
@@ -344,21 +326,14 @@ function Base.resize!(A::ROCVector{T}, n::Integer) where T
     # if A.buf.host_ptr != C_NULL
     #     throw(ArgumentError("Cannot resize an unowned `ROCVector`"))
     # end
-
-    # TODO: add additional space to allow for quicker resizing
     n == length(A) && return A
 
     maxsize = n * aligned_sizeof(T)
     bufsize = Base.isbitsunion(T) ? (maxsize + n) : maxsize
     new_buf = Mem.HIPBuffer(bufsize; stream=stream())
 
     copy_size = min(length(A), n) * aligned_sizeof(T)
-    if copy_size > 0
-        Mem.transfer!(new_buf, convert(Mem.AbstractAMDBuffer, A.buf[]),
-            copy_size; stream=stream())
-    end
-
-    # Free old buffer.
+    copy_size > 0 && Mem.memcpy!(new_buf, pointer(A), copy_size; stream=stream())
     unsafe_free!(A)
 
     A.buf = DataRef(pool_free, Managed(new_buf))

src/runtime/memory/hip.jl

@@ -88,21 +88,16 @@ function free(buf::HIPBuffer; stream::HIP.HIPStream)
     return
 end
 
-function upload!(dst::HIPBuffer, src::Ptr, bytesize::Int; stream::HIP.HIPStream)
+function memcpy!(dst, src, bytesize::Int; stream::HIP.HIPStream)
     bytesize == 0 && return
-    HIP.hipMemcpyHtoDAsync(dst, src, bytesize, stream)
-    return
-end
-
-function download!(dst::Ptr, src::HIPBuffer, bytesize::Int; stream::HIP.HIPStream)
-    bytesize == 0 && return
-    HIP.hipMemcpyDtoHAsync(dst, src, bytesize, stream)
-    return
-end
-
-function transfer!(dst::HIPBuffer, src::HIPBuffer, bytesize::Int; stream::HIP.HIPStream)
-    bytesize == 0 && return
-    HIP.hipMemcpyDtoDAsync(dst, src, bytesize, stream)
+    dst_type = attributes(convert(Ptr{Cvoid}, dst)).type
+    src_type = attributes(convert(Ptr{Cvoid}, src)).type
+    kind = if src_type == HIP.hipMemoryTypeDevice
+        dst_type == HIP.hipMemoryTypeDevice ? HIP.hipMemcpyDeviceToDevice : HIP.hipMemcpyDeviceToHost
+    else
+        dst_type == HIP.hipMemoryTypeDevice ? HIP.hipMemcpyHostToDevice : HIP.hipMemcpyHostToHost
+    end
+    HIP.memcpy(dst, src, bytesize, kind, stream)
     return
 end
 
@@ -165,28 +160,6 @@ function view(buf::HostBuffer, bytesize::Int)
         buf.bytesize - bytesize, buf.own)
 end
 
-upload!(dst::HostBuffer, src::Ptr, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyHostToHost, stream)
-
-upload!(dst::HostBuffer, src::HIPBuffer, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyDeviceToHost, stream)
-
-download!(dst::Ptr, src::HostBuffer, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyHostToHost, stream)
-
-download!(dst::HIPBuffer, src::HostBuffer, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyHostToDevice, stream)
-
-transfer!(dst::HostBuffer, src::HostBuffer, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyHostToHost, stream)
-
-# download!(::Ptr, ::HIPBuffer)
-transfer!(dst::HostBuffer, src::HIPBuffer, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyDeviceToHost, stream)
-
-# upload!(::HIPBuffer, ::Ptr)
-transfer!(dst::HIPBuffer, src::HostBuffer, sz::Int; stream::HIP.HIPStream) =
-    HIP.memcpy(dst, src, sz, HIP.hipMemcpyHostToDevice, stream)
 
 Base.convert(::Type{Ptr{T}}, buf::HostBuffer) where T = convert(Ptr{T}, buf.ptr)