Wait for the device cooperatively

LocalPreferences.toml

@@ -11,3 +11,8 @@
 # - "khronos": the Khronos SPIRV-LLVM-Translator, which additionally requires
 #   `SPIRV_LLVM_Translator_jll` (or `SPIRV_LLVM_Translator_unified_jll`) to be loaded
 #llvm_to_spirv_backend="llvm"
+
+# Whether to wait for the device cooperatively, yielding to other tasks instead of blocking
+# the calling thread in the OpenCL driver. Disable for bisecting issues, or to compare
+# against the blocking behavior.
+#nonblocking_synchronization=true

Project.toml

@@ -33,7 +33,7 @@ SPIRVIntrinsics = {path = "lib/intrinsics"}
 Adapt = "4"
 GPUArrays = "11.2.1"
 GPUCompiler = "2.9"
-GPUToolbox = "3.1"
+GPUToolbox = "3.3.2"
 KernelAbstractions = "0.9.38"
 LLVM = "9.6"
 LinearAlgebra = "1"

lib/cl/CL.jl

@@ -20,6 +20,7 @@ include("device.jl")
 include("context.jl")
 include("cmdqueue.jl")
 include("event.jl")
+include("synchronization.jl")
 include("memory.jl")
 include("program.jl")
 include("kernel.jl")

lib/cl/api.jl

@@ -1,4 +1,5 @@
 import OpenCL_jll
+using GPUToolbox: GPUToolbox, @gcsafe_ccall, cooperative_wait
 
 const libopencl = OpenCL_jll.libopencl
 
@@ -89,7 +90,7 @@ function ext_ccall(platform, ex)
     @gensym fptr
     esc(quote
         $fptr = $clGetExtensionFunctionAddressForPlatform($platform, $fn)
-        @ccall $(Expr(:($), fptr))($(argexprs...))::$ret
+        @gcsafe_ccall $(Expr(:($), fptr))($(argexprs...))::$ret
     end)
 end
 
@@ -183,7 +184,7 @@ const initialized = Ref{Bool}(false)
 
     withenv("OCL_ICD_FILENAMES"=>ocd_filenames) do
         num_platforms = Ref{Cuint}()
-        @ccall libopencl.clGetPlatformIDs(
+        @gcsafe_ccall libopencl.clGetPlatformIDs(
             0::cl_uint, C_NULL::Ptr{cl_platform_id},
             num_platforms::Ptr{cl_uint})::cl_int
 

lib/cl/event.jl

@@ -28,27 +28,36 @@ end
 
 NannyEvent(evt::Event, obj; retain=false) = NannyEvent(evt.id, obj; retain)
 
-macro return_event(evt)
+# commands that should block are submitted without blocking in the driver, and waited for
+# cooperatively instead (see `wait_events`). that wait cannot be interrupted, as the caller
+# may release the memory the command accesses once we return.
+macro return_event(evt, blocking=false)
     quote
         evt = $(esc(evt))
-        try
-            return Event(evt, retain=false)
+        evt = try
+            Event(evt, retain=false)
         catch err
+            $(esc(blocking)) && unchecked_clWaitForEvents(1, Ref(evt))
             clReleaseEvent(evt)
             throw(err)
         end
+        $(esc(blocking)) && wait_events([evt]; cancellable=false)
+        return evt
     end
 end
 
-macro return_nanny_event(evt, obj)
+macro return_nanny_event(evt, obj, blocking=false)
     quote
         evt = $(esc(evt))
-        try
-            return NannyEvent(evt, $(esc(obj)))
+        evt = try
+            NannyEvent(evt, $(esc(obj)))
         catch err
+            $(esc(blocking)) && unchecked_clWaitForEvents(1, Ref(evt))
             clReleaseEvent(evt)
             throw(err)
         end
+        $(esc(blocking)) && wait_events([evt]; cancellable=false)
+        return evt
     end
 end
 
@@ -143,26 +152,12 @@ function add_callback(evt::AbstractEvent, callback::Function)
     end
 end
 
-function Base.wait(evt::AbstractEvent)
-    evt_id = [evt.id]
-    clWaitForEvents(cl_uint(1), evt_id)
-    return evt
-end
-
-function Base.wait(evts::Vector{AbstractEvent})
-    isempty(evts) && return evts
-    evt_ids = [pointer(evt) for evt in evts]
-    GC.@preserve evts begin
-        clWaitForEvents(length(evt_ids), evt_ids)
-    end
-    return evts
-end
-
-function enqueue_marker_with_wait_list(wait_for::Vector{AbstractEvent})
+function enqueue_marker_with_wait_list(wait_for::Vector{AbstractEvent};
+                                       queue::CmdQueue=queue())
     n_wait_events = cl_uint(length(wait_for))
     wait_evt_ids = [evt.id for evt in wait_for]
     ret_evt = Ref{cl_event}()
-    clEnqueueMarkerWithWaitList(queue(), n_wait_events,
+    clEnqueueMarkerWithWaitList(queue, n_wait_events,
                                 isempty(wait_evt_ids) ? C_NULL : wait_evt_ids,
                                 ret_evt)
     @return_event ret_evt[]

lib/cl/memory/buffer.jl

@@ -89,9 +89,9 @@ function enqueue_read(dst::Ptr, src::Buffer, src_off::Int, nbytes::Int;
     evt_ids = isempty(wait_for) ? C_NULL : [pointer(evt) for evt in wait_for]
     GC.@preserve wait_for begin
         ret_evt = Ref{cl_event}()
-        clEnqueueReadBuffer(queue(), src, blocking, src_off, nbytes, dst,
+        clEnqueueReadBuffer(queue(), src, false, src_off, nbytes, dst,
                             n_evts, evt_ids, ret_evt)
-        @return_nanny_event(ret_evt[], dst)
+        @return_nanny_event(ret_evt[], dst, blocking)
     end
 end
 enqueue_read(dst::Ptr, src::Buffer, nbytes; kwargs...) =
@@ -104,9 +104,9 @@ function enqueue_write(dst::Buffer, dst_off::Int, src::Ptr, nbytes::Int;
     evt_ids = isempty(wait_for) ? C_NULL : [pointer(evt) for evt in wait_for]
     GC.@preserve wait_for begin
         ret_evt = Ref{cl_event}()
-        clEnqueueWriteBuffer(queue(), dst, blocking, dst_off, nbytes, src,
+        clEnqueueWriteBuffer(queue(), dst, false, dst_off, nbytes, src,
                              n_evts, evt_ids, ret_evt)
-        @return_nanny_event(ret_evt[], dst)
+        @return_nanny_event(ret_evt[], dst, blocking)
     end
 end
 enqueue_write(dst::Buffer, src::Ptr, nbytes; kwargs...) =
@@ -122,7 +122,7 @@ function enqueue_copy(dst::Buffer, dst_off::Int, src::Buffer, src_off::Int,
         ret_evt = Ref{cl_event}()
         clEnqueueCopyBuffer(queue(), src, dst, src_off, dst_off, nbytes,
                             n_evts, evt_ids, ret_evt)
-        @return_event ret_evt[]
+        @return_event ret_evt[] blocking
     end
 end
 enqueue_copy(dst::Buffer, src::Buffer, N; kwargs...) =

lib/cl/memory/svm.jl

@@ -86,8 +86,8 @@ function enqueue_svm_copy(
     evt_ids = isempty(wait_for) ? C_NULL : [pointer(evt) for evt in wait_for]
     return GC.@preserve wait_for begin
         ret_evt = Ref{cl_event}()
-        clEnqueueSVMMemcpy(queue, blocking, dst, src, nbytes, n_evts, evt_ids, ret_evt)
-        @return_event ret_evt[]
+        clEnqueueSVMMemcpy(queue, false, dst, src, nbytes, n_evts, evt_ids, ret_evt)
+        @return_event ret_evt[] blocking
     end
 end
 

lib/cl/memory/usm.jl

@@ -208,8 +208,8 @@ function enqueue_usm_copy(
     return GC.@preserve wait_for begin
         com = (CL_NULL, C_NULL)
         ret_evt = (dst in com || src in com) ? C_NULL : Ref{cl_event}()
-        clEnqueueMemcpyINTEL(queue, blocking, dst, src, nbytes, n_evts, evt_ids, ret_evt)
-        @return_event ret_evt[]
+        clEnqueueMemcpyINTEL(queue, false, dst, src, nbytes, n_evts, evt_ids, ret_evt)
+        @return_event ret_evt[] blocking
     end
 end
 

lib/cl/pointer.jl

@@ -138,6 +138,10 @@ function Base.unsafe_convert(::Type{PtrOrCLPtr{T}}, val) where {T}
     return Base.bitcast(PtrOrCLPtr{T}, ptr)
 end
 
+# `@gcsafe_ccall` converts arguments itself before passing them to `ccall`, which converts
+# them again, so make that second (identity) conversion unambiguous
+Base.unsafe_convert(::Type{PtrOrCLPtr{T}}, x::PtrOrCLPtr{T}) where {T} = x
+
 
 #
 # Device reference objects
@@ -207,6 +211,8 @@ Base.unsafe_convert(::Type{RefOrCLRef{T}}, x::Ref{T}) where {T} =
     Base.bitcast(RefOrCLRef{T}, Base.unsafe_convert(Ptr{T}, x))
 Base.unsafe_convert(::Type{RefOrCLRef{T}}, x) where {T} =
     Base.bitcast(RefOrCLRef{T}, Base.unsafe_convert(Ptr{T}, x))
+# see `PtrOrCLPtr`
+Base.unsafe_convert(::Type{RefOrCLRef{T}}, x::RefOrCLRef{T}) where {T} = x
 
 # support conversion from GPU ref
 Base.unsafe_convert(::Type{RefOrCLRef{T}}, x::CLRefs{T}) where {T} =

lib/cl/synchronization.jl

@@ -0,0 +1,136 @@
+# cooperative synchronization
+#
+# Waiting for the device should not block the calling thread in the OpenCL driver, but
+# yield to the Julia scheduler so that other tasks can run in the meantime. GPUToolbox's
+# `cooperative_wait` implements this, in one of two ways:
+# - for GPUs, by polling the status of commands for a while, and then handing a blocking
+#   wait to a worker thread;
+# - for devices that execute on the host's CPU cores, by polling only briefly, and then
+#   having the driver notify us when commands complete. waking a worker, or polling for
+#   longer, would compete with the commands for those cores, slowing them down
+#   considerably. (drivers for GPUs may deliver these notifications late, NVIDIA's by
+#   about 20 ms.)
+
+# whether to wait for the device cooperatively. disable to block in the OpenCL driver
+# instead, e.g., for bisecting issues or comparing against the blocking behavior.
+const nonblocking_synchronization =
+    @load_preference("nonblocking_synchronization", true)
+
+event_ids(evts) = cl_event[pointer(evt) for evt in evts]
+
+# commands have completed when their execution status is `CL_COMPLETE` (zero), or negative
+# when they were terminated abnormally
+function iscomplete(evt::AbstractEvent)
+    status = Ref{Cint}()
+    clGetEventInfo(evt, CL_EVENT_COMMAND_EXECUTION_STATUS, sizeof(Cint), status, C_NULL)
+    return status[] <= CL_COMPLETE
+end
+
+# commands only need to start executing once the queue they were submitted to has been
+# flushed. waiting in the driver flushes implicitly, but polling does not. returns the
+# queue, or `C_NULL` for user events, which do not belong to one.
+function flush_queue(evt::AbstractEvent)
+    queue = Ref{cl_command_queue}()
+    clGetEventInfo(evt, CL_EVENT_COMMAND_QUEUE, sizeof(cl_command_queue), queue, C_NULL)
+    queue[] == C_NULL || clFlush(queue[])
+    return queue[]
+end
+
+# whether commands on `queue` execute on the host's CPU cores
+const cpu_devices = Dict{cl_device_id, Bool}()
+const cpu_devices_lock = ReentrantLock()
+function executes_on_cpu(queue::cl_command_queue)
+    queue == C_NULL && return false
+    device = Ref{cl_device_id}()
+    clGetCommandQueueInfo(queue, CL_QUEUE_DEVICE, sizeof(cl_device_id), device, C_NULL)
+    return @lock cpu_devices_lock get!(cpu_devices, device[]) do
+        type = Ref{cl_device_type}()
+        clGetDeviceInfo(device[], CL_DEVICE_TYPE, sizeof(cl_device_type), type, C_NULL)
+        type[] & CL_DEVICE_TYPE_CPU != 0
+    end
+end
+
+# driver notification that a command has completed
+function notify_completion(::cl_event, ::Cint, payload::Ptr{Cvoid})
+    GPUToolbox.signal_completion(payload)
+    return
+end
+subscribe_completion(evt, payload) =
+    clSetEventCallback(evt, CL_COMPLETE,
+                       @cfunction(notify_completion, Cvoid, (cl_event, Cint, Ptr{Cvoid})),
+                       payload)
+
+# block until the events have completed, without checking for errors
+blocking_wait(evts::Vector{<:AbstractEvent}) =
+    GC.@preserve evts unchecked_clWaitForEvents(length(evts), event_ids(evts))
+blocking_wait(evt::AbstractEvent) =
+    GC.@preserve evt unchecked_clWaitForEvents(1, Ref(pointer(evt)))
+
+# wait for events to complete, throwing a `CLError` if a command was terminated abnormally.
+#
+# by default, the wait can be interrupted, in which case the commands may still be
+# executing. waits that cannot be interrupted only return (or throw) once the commands have
+# completed, which is required when they access memory that the caller releases upon
+# returning.
+function wait_events(evts::Vector{<:AbstractEvent}; cancellable::Bool=true)
+    if nonblocking_synchronization
+        try
+            on_cpu = true
+            for evt in evts
+                on_cpu &= executes_on_cpu(flush_queue(evt))
+            end
+            if on_cpu
+                for evt in evts
+                    cooperative_wait(blocking_wait, evt; subscribe=subscribe_completion,
+                                     isdone=iscomplete, spin=10e-6, cancellable)
+                end
+            else
+                cooperative_wait(blocking_wait, evts; isdone=evts -> all(iscomplete, evts),
+                                 cancellable)
+            end
+        catch
+            # waits that cannot be interrupted only throw once the commands have completed
+            # (or waiting failed, in which case this blocks), but host memory still needs to
+            # be synchronized before the caller unwinds
+            cancellable || blocking_wait(evts)
+            rethrow()
+        end
+    end
+
+    # synchronize host memory and report errors (without blocking, if we waited above)
+    GC.@preserve evts clWaitForEvents(length(evts), event_ids(evts))
+    return
+end
+
+"""
+    wait(evt::cl.AbstractEvent)
+    wait(evts::Vector{cl.AbstractEvent})
+
+Wait for the commands associated with the event(s) to complete, throwing a `CLError` if a
+command was terminated abnormally.
+
+While waiting, the current task yields to others. Waiting is interruptible, in which case
+the commands keep executing.
+"""
+function Base.wait(evt::AbstractEvent)
+    wait_events(AbstractEvent[evt])
+    return evt
+end
+
+function Base.wait(evts::Vector{AbstractEvent})
+    isempty(evts) || wait_events(evts)
+    return evts
+end
+
+# wait for the commands submitted to `queue` so far to complete. this does not check for
+# device-side exceptions, as `finish` does.
+function wait_idle(queue::CmdQueue; cancellable::Bool=true)
+    if nonblocking_synchronization
+        # there is no way to query whether a queue is idle, so wait for a marker instead
+        marker = enqueue_marker_with_wait_list(AbstractEvent[]; queue)
+        wait_events([marker]; cancellable)
+    else
+        clFinish(queue)
+    end
+    return
+end

res/wrap.jl

@@ -77,9 +77,14 @@ function rewriter!(ctx, options)
             target_expr = call_expr.args[1].args[1]
             fn = String(target_expr.args[2].value)
 
-            # replace `@ccall` with `@ext_ccall` when dealing with an extension function
+            # replace `@ccall` with `@ext_ccall` when dealing with an extension function,
+            # and with `@gcsafe_ccall` otherwise: API calls may block (waiting for the
+            # device, compiling, ...), which should not prevent other threads from
+            # collecting garbage.
             if basename(get_filename(node.cursor)) == "cl_ext.h"
                 expr.args[2].args[1].args[1] = Symbol("@ext_ccall")
+            else
+                expr.args[2].args[1].args[1] = Symbol("@gcsafe_ccall")
             end
 
             # rewrite pointer return types

src/OpenCLKernels.jl

@@ -42,7 +42,6 @@ end
 KA.supports_unified(::OpenCLBackend) = cl.default_memory_backend(cl.device(); unified=true) !== nothing
 
 KA.get_backend(::CLArray) = OpenCLBackend()
-# TODO should be non-blocking
 KA.synchronize(::OpenCLBackend) = OpenCL.synchronize()
 KA.supports_float64(::OpenCLBackend) = in("cl_khr_fp64", cl.device().extensions)