Merge pull request #62 from JuliaGPU/vc/multievent

add multievents

Author: vchuravy

Project.toml

@@ -1,7 +1,7 @@
 name = "KernelAbstractions"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 authors = ["Valentin Churavy <v.churavy@gmail.com>"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/KernelAbstractions.jl

@@ -2,7 +2,7 @@ module KernelAbstractions
 
 export @kernel
 export @Const, @localmem, @private, @uniform, @synchronize, @index, groupsize, @print
-export Device, GPU, CPU, CUDA, Event
+export Device, GPU, CPU, CUDA, Event, MultiEvent
 export async_copy!
 
 
@@ -64,6 +64,27 @@ macro Const end
 abstract type Event end
 import Base.wait
 
+struct NoneEvent <: Event end
+
+struct MultiEvent{T} <: Event
+    events::T
+    MultiEvent() = new{Tuple{}}(())
+    function MultiEvent(events::Tuple{Vararg{<:Event}})
+        evs = tuplejoin(map(flatten, events)...)
+        new{typeof(evs)}(evs)
+    end
+end
+MultiEvent(::Nothing) = MultiEvent()
+MultiEvent(ev::MultiEvent) = ev
+
+@inline tuplejoin(x) = x
+@inline tuplejoin(x, y) = (x..., y...)
+@inline tuplejoin(x, y, z...) = (x..., tuplejoin(y, z...)...)
+
+flatten(ev::MultiEvent) = tuplejoin(map(flatten, ev.events)...)
+flatten(ev::NoneEvent) = ()
+flatten(ev::Event) = (ev,)
+
 """
     async_copy!(::Device, dest::AbstractArray, src::AbstractArray; dependencies = nothing)
 

src/backends/cpu.jl

@@ -1,12 +1,24 @@
 struct CPUEvent <: Event
-    task::Union{Nothing, Core.Task}
+    task::Core.Task
 end
 
 function Event(::CPU)
-    return CPUEvent(nothing)
+    return NoneEvent()
+end
+
+wait(ev::Union{CPUEvent, NoneEvent, MultiEvent}, progress=nothing) = wait(CPU(), ev, progress)
+wait(::CPU, ev::NoneEvent, progress=nothing) = nothing
+
+function wait(cpu::CPU, ev::MultiEvent, progress=nothing)
+    dependencies = collect(ev.events)
+    cpudeps   = filter(d->d isa CPUEvent && d.task !== nothing, dependencies)
+    otherdeps = filter(d->!(d isa CPUEvent), dependencies)
+    Base.sync_end(map(e->e.task, cpudeps))
+    for event in otherdeps
+        wait(CPU(), event, progress)
+    end
 end
 
-wait(ev::CPUEvent, progress=nothing) = wait(CPU(), ev, progress)
 function wait(::CPU, ev::CPUEvent, progress=nothing)
     ev.task === nothing && return
 
@@ -18,23 +30,9 @@ function wait(::CPU, ev::CPUEvent, progress=nothing)
         end
     end
 end
-function __waitall(::CPU, dependencies, progress)
-    if dependencies isa Event
-        dependencies = (dependencies,)
-    end
-    if dependencies !== nothing
-        dependencies = collect(dependencies)
-        cpudeps   = filter(d->d isa CPUEvent && d.task !== nothing, dependencies)
-        otherdeps = filter(d->!(d isa CPUEvent), dependencies)
-        Base.sync_end(map(e->e.task, cpudeps))
-        for event in otherdeps
-            wait(CPU(), event, progress)
-        end
-    end
-end
 
 function async_copy!(::CPU, A, B; dependencies=nothing)
-    __waitall(CPU(), dependencies, yield)
+    wait(CPU(), MultiEvent(dependencies), yield)
     copyto!(A, B)
     return CPUEvent(nothing)
 end
@@ -65,7 +63,7 @@ end
 
 # Inference barriers
 function __run(obj, ndrange, iterspace, args, dependencies, ::Val{dynamic}) where dynamic
-    __waitall(CPU(), dependencies, yield)
+    wait(CPU(), MultiEvent(dependencies), yield)
     N = length(iterspace)
     Nthreads = Threads.nthreads()
     if Nthreads == 1

src/backends/cuda.jl

@@ -56,6 +56,7 @@ function Event(::CUDA)
 end
 
 wait(ev::CudaEvent, progress=nothing) = wait(CPU(), ev, progress)
+
 function wait(::CPU, ev::CudaEvent, progress=nothing)
     if progress === nothing
         CUDAdrv.synchronize(ev.event)
@@ -68,30 +69,24 @@ function wait(::CPU, ev::CudaEvent, progress=nothing)
 end
 
 # Use this to synchronize between computation using the CuDefaultStream
-function wait(::CUDA, ev::CudaEvent, progress=nothing)
-    CUDAdrv.wait(ev.event, CUDAdrv.CuDefaultStream())
-end
+wait(::CUDA, ev::CudaEvent, progress=nothing, stream=CUDAdrv.CuDefaultStream()) = CUDAdrv.wait(ev.event, stream)
+wait(::CUDA, ev::NoneEvent, progress=nothing, stream=nothing) = nothing
 
 # There is no efficient wait for CPU->GPU synchronization, so instead we
 # do a CPU wait, and therefore block anyone from submitting more work.
 # We maybe could do a spinning wait on the GPU and atomic flag to signal from the CPU,
 # but which stream would we target?
-wait(::CUDA, ev::CPUEvent,  progress=nothing) = wait(CPU(), ev, progress)
-
-function __waitall(::CUDA, dependencies, progress, stream)
-    if dependencies isa Event
-        dependencies = (dependencies,)
+wait(::CUDA, ev::CPUEvent, progress=nothing, stream=nothing) = wait(CPU(), ev, progress)
+
+function wait(::CUDA, ev::MultiEvent, progress=nothing, stream=CUDAdrv.CuDefaultStream())
+    dependencies = collect(ev.events)
+    cudadeps  = filter(d->d isa CudaEvent,    dependencies)
+    otherdeps = filter(d->!(d isa CudaEvent), dependencies)
+    for event in cudadeps
+        CUDAdrv.wait(event.event, stream)
     end
-    if dependencies !== nothing
-        dependencies = collect(dependencies)
-        cudadeps  = filter(d->d isa CudaEvent,    dependencies)
-        otherdeps = filter(d->!(d isa CudaEvent), dependencies)
-        for event in cudadeps
-            CUDAdrv.wait(event.event, stream)
-        end
-        for event in otherdeps
-            wait(CUDA(), event, progress)
-        end
+    for event in otherdeps
+        wait(CUDA(), event, progress)
     end
 end
 
@@ -119,7 +114,7 @@ function async_copy!(::CUDA, A, B; dependencies=nothing)
     B isa Array && __pin!(B)
 
     stream = next_stream()
-    __waitall(CUDA(), dependencies, yield, stream)
+    wait(CUDA(), MultiEvent(dependencies), yield, stream)
     event = CuEvent(CUDAdrv.EVENT_DISABLE_TIMING)
     GC.@preserve A B begin
         destptr = pointer(A)
@@ -145,12 +140,9 @@ function (obj::Kernel{CUDA})(args...; ndrange=nothing, dependencies=nothing, wor
     if workgroupsize isa Integer
         workgroupsize = (workgroupsize, )
     end
-    if dependencies isa Event
-        dependencies = (dependencies,)
-    end
 
     stream = next_stream()
-    __waitall(CUDA(), dependencies, yield, stream)
+    wait(CUDA(), MultiEvent(dependencies), yield, stream)
 
     if KernelAbstractions.workgroupsize(obj) <: DynamicSize && workgroupsize === nothing
         # TODO: allow for NDRange{1, DynamicSize, DynamicSize}(nothing, nothing)