Merge pull request #116 from JuliaGPU/jps/mark-wait

Add mark/wait synchronization system

Author: jpsamaroo

Project.toml

@@ -1,7 +1,7 @@
 name = "AMDGPU"
 uuid = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 authors = ["Julian P Samaroo <jpsamaroo@jpsamaroo.me>"]
-version = "0.2.5"
+version = "0.2.6"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/AMDGPU.jl

@@ -48,6 +48,7 @@ if get(ENV, "AMDGPUNATIVE_OPENCL", "") != ""
 end
 =#
 include("runtime.jl")
+include("sync.jl")
 
 # Device sources must load _before_ the compiler infrastructure
 # because of generated functions.

src/array.jl

@@ -53,17 +53,18 @@ end
 # Host abstractions
 #
 
-
 mutable struct ROCArray{T,N} <: AbstractGPUArray{T,N}
     buf::Mem.Buffer
     own::Bool
 
     dims::Dims{N}
     offset::Int
 
+    syncstate::SyncState
+
     function ROCArray{T,N}(buf::Mem.Buffer, dims::Dims{N}; offset::Integer=0, own::Bool=true) where {T,N}
         @assert isbitstype(T) "ROCArray only supports bits types"
-        xs = new{T,N}(buf, own, dims, offset)
+        xs = new{T,N}(buf, own, dims, offset, SyncState())
         if own
             hsaref!()
             Mem.retain(buf)
@@ -79,6 +80,21 @@ function unsafe_free!(xs::ROCArray)
     return
 end
 
+wait!(x::ROCArray) = wait!(x.syncstate)
+mark!(x::ROCArray, s) = mark!(x.syncstate, s)
+wait!(xs::Vector{<:ROCArray}) = foreach(wait!, xs)
+mark!(xs::Vector{<:ROCArray}, s) = foreach(x->mark!(x,s), xs)
+wait!(xs::NTuple{N,<:ROCArray} where N) = foreach(wait!, xs)
+mark!(xs::NTuple{N,<:ROCArray} where N, s) = foreach(x->mark!(x,s), xs)
+function Adapt.adapt_storage(::WaitAdaptor, x::ROCArray)
+    wait!(x.syncstate)
+    x
+end
+function Adapt.adapt_storage(ma::MarkAdaptor, x::ROCArray)
+    mark!(x.syncstate, ma.s)
+    x
+end
+
 ## aliases
 
 const ROCVector{T} = ROCArray{T,1}
@@ -154,6 +170,7 @@ function Base.copyto!(dest::Array{T}, d_offset::Integer,
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset+amount-1)
     @boundscheck checkbounds(source, s_offset+amount-1)
+    wait!(source)
     Mem.download!(pointer(dest, d_offset),
                   Mem.view(source.buf, (s_offset-1)*sizeof(T)),
                   amount*sizeof(T))
@@ -165,6 +182,7 @@ function Base.copyto!(dest::ROCArray{T}, d_offset::Integer,
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset+amount-1)
     @boundscheck checkbounds(source, s_offset+amount-1)
+    wait!(dest)
     Mem.upload!(Mem.view(dest.buf, (d_offset-1)*sizeof(T)),
                 pointer(source, s_offset),
                 amount*sizeof(T))
@@ -176,6 +194,8 @@ function Base.copyto!(dest::ROCArray{T}, d_offset::Integer,
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset+amount-1)
     @boundscheck checkbounds(source, s_offset+amount-1)
+    wait!(dest)
+    wait!(source)
     Mem.transfer!(Mem.view(dest.buf, (d_offset-1)*sizeof(T)),
                   Mem.view(source.buf, (s_offset-1)*sizeof(T)),
                   amount*sizeof(T))

src/blas/rocBLAS.jl

@@ -1,6 +1,7 @@
 module rocBLAS
 
 using ..AMDGPU
+import AMDGPU: wait!, mark!
 
 using LinearAlgebra