Make synchronization cooperative, using GPUToolbox

Project.toml

@@ -40,7 +40,7 @@ CEnum = "0.4, 0.5"
 ExprTools = "0.1"
 GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
-GPUToolbox = "3.1"
+GPUToolbox = "3.3.1"
 KernelAbstractions = "0.9.39"
 LLVM = "6, 7, 8, 9"
 NEO_jll = "=26.18.38308"

lib/level-zero/oneL0.jl

@@ -124,6 +124,7 @@ end
 include("context.jl")
 include("cmdqueue.jl")
 include("cmdlist.jl")
+include("synchronization.jl")
 include("fence.jl")
 include("event.jl")
 include("barrier.jl")

lib/level-zero/synchronization.jl

@@ -0,0 +1,51 @@
+# cooperative synchronization
+#
+# `zeCommandListHostSynchronize` and `zeCommandQueueSynchronize` block the calling thread
+# until the work has completed, so no other task can run on it in the meantime. Instead, wait
+# using GPUToolbox's `cooperative_wait`: first poll, which keeps the latency of short
+# operations low, and then block in the driver on a separate thread, while the calling task
+# yields.
+
+using GPUToolbox: cooperative_wait
+
+export nonblocking_synchronize
+
+const SyncObject = Union{ZeImmediateCommandList, ZeCommandQueue}
+
+# with a zero timeout, a synchronization is a query
+function check_done(res::ze_result_t)
+    if res == RESULT_NOT_READY
+        return false
+    elseif res == RESULT_SUCCESS
+        return true
+    else
+        throw_api_error(res)
+    end
+end
+Base.isdone(list::ZeImmediateCommandList) =
+    check_done(unchecked_zeCommandListHostSynchronize(list, 0))
+Base.isdone(queue::ZeCommandQueue) =
+    check_done(unchecked_zeCommandQueueSynchronize(queue, 0))
+
+# the blocking synchronization, marked GC-safe so that it doesn't keep the GC from running
+gcsafe_synchronize(list::ZeImmediateCommandList) =
+    @gcsafe_ccall libze_loader.zeCommandListHostSynchronize(
+        list::ze_command_list_handle_t, typemax(UInt64)::UInt64)::ze_result_t
+gcsafe_synchronize(queue::ZeCommandQueue) =
+    @gcsafe_ccall libze_loader.zeCommandQueueSynchronize(
+        queue::ze_command_queue_handle_t, typemax(UInt64)::UInt64)::ze_result_t
+
+"""
+    nonblocking_synchronize(list_or_queue)
+
+Wait for the work on an immediate command list or command queue to complete, like
+[`synchronize`](@ref), but without blocking the calling thread: other tasks keep running
+while this one waits.
+"""
+function nonblocking_synchronize(obj::SyncObject)
+    # when polling found the work to be done, synchronize again to check for errors
+    res = @something(cooperative_wait(gcsafe_synchronize, obj; isdone=Base.isdone),
+                     gcsafe_synchronize(obj))
+    res == RESULT_SUCCESS || throw_api_error(res)
+    return
+end

lib/utils/APIUtils.jl

@@ -1,8 +1,8 @@
 module APIUtils
 
 # helpers that facilitate working with C APIs
-using GPUToolbox: @checked, @debug_ccall
-export @checked, @debug_ccall
+using GPUToolbox: @checked, @debug_ccall, @gcsafe_ccall
+export @checked, @debug_ccall, @gcsafe_ccall
 include("enum.jl")
 
 end

src/compiler/execution.jl

@@ -360,7 +360,8 @@ end
     spill > s.scratch_hwm && scratch_hedge!(s, spill)
 
     append_launch!(s.list, kernel, groups)
-    oneL0.sync_each_submission() && oneL0.synchronize(s.list)
+    # wait cooperatively, as `synchronize` does, or the workaround blocks the thread
+    oneL0.sync_each_submission() && oneL0.nonblocking_synchronize(s.list)
     return
 end
 

src/context.jl

@@ -378,12 +378,15 @@ function synchronize_all_streams(ctx::ZeContext, dev::Union{ZeDevice, Nothing})
 end
 
 """
-    synchronize()
-    synchronize(stream::oneStream)
+    synchronize(; blocking=false)
+    synchronize(stream::oneStream; blocking=false)
 
-Block the host thread until all operations on the calling task's stream for the current
-context and device have completed: work appended to the immediate command list as well
-as oneMKL work on the companion queue.
+Block the calling task until all operations on its stream for the current context and
+device have completed: work appended to the immediate command list as well as oneMKL work
+on the companion queue.
+
+Unless `blocking` is set, other tasks keep running while waiting: the host thread is only
+blocked in the driver when the work is already done.
 
 This is useful for timing operations or ensuring that GPU work has finished before
 accessing results on the CPU.
@@ -398,18 +401,19 @@ println("GPU work completed")
 
 See also: [`global_stream`](@ref), [`context`](@ref), [`device`](@ref)
 """
-function oneL0.synchronize(s::oneStream)
-    oneL0.synchronize(s.list)
+function oneL0.synchronize(s::oneStream; blocking::Bool=false)
+    sync = blocking ? oneL0.synchronize : oneL0.nonblocking_synchronize
+    sync(s.list)
     q = s.queue
     if q !== nothing
-        oneL0.synchronize(q)
+        sync(q)
         s.mkl_dirty = false
     end
     return
 end
 
-function oneL0.synchronize()
-    oneL0.synchronize(global_stream(context(), device()))
+function oneL0.synchronize(; blocking::Bool=false)
+    oneL0.synchronize(global_stream(context(), device()); blocking)
 end
 
 # Julia → MKL ordering: everything Julia appended to the task's immediate list must be

src/oneAPIKernels.jl

@@ -26,7 +26,6 @@ oneAPIBackend(; prefer_blocks = false, always_inline = false) = oneAPIBackend(pr
 @inline KA.ones(::oneAPIBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where {T} = fill!(oneArray{T, length(dims), unified ? oneAPI.oneL0.SharedBuffer : oneAPI.oneL0.DeviceBuffer}(undef, dims), one(T))
 
 KA.get_backend(::oneArray) = oneAPIBackend()
-# TODO should be non-blocking
 KA.synchronize(::oneAPIBackend) = oneAPI.oneL0.synchronize()
 KA.supports_float64(::oneAPIBackend) = false  # TODO: Check if this is device dependent
 KA.supports_unified(::oneAPIBackend) = true

test/execution.jl

@@ -743,6 +743,109 @@ end
     @test all(results)
 end
 
+function fill_kernel(a, x)
+    @inbounds a[get_global_id()] = x
+    return
+end
+
+@testset "cooperative synchronize" begin
+    # keep the work on the task's stream waiting for an event until another task signals
+    # it, which makes these tests independent of timing: synchronizing can only return after
+    # that task got to run. if it cannot (e.g., because synchronizing blocks the thread), a
+    # watchdog on another thread signals the event after a while instead, and records that.
+    pool = oneL0.ZeEventPool(context(), 1, device();
+                             flags=oneL0.ZE_EVENT_POOL_FLAG_HOST_VISIBLE)
+    gate = pool[1]
+    timeout = UInt64(60_000_000_000)    # ns
+
+    # run `f` with the task's stream waiting for the gate, which another task on the same
+    # thread opens, but only after it got to run many more times than the polling at the
+    # start of a synchronization yields. returns whether `f` only returned after the gate
+    # had been opened, and the watchdog did not have to.
+    function gated(f)
+        list = oneAPI.global_stream(context(), device()).list
+        reset(gate)
+        # while the gate is closed, nothing on the host may wait for the GPU to become
+        # idle, as freeing memory does on some stacks. so avoid running finalizers, by
+        # collecting beforehand and not collecting while the gate is closed.
+        GC.gc(true)
+        gc_enabled = GC.enable(false)
+        watching = Threads.Atomic{Bool}(false)
+        opened = Threads.Atomic{Bool}(false)
+        in_time, timed_out = false, true
+        local watchdog, opener
+        try
+            # the watchdog waits on one of `cooperative_wait`'s worker threads, which keep
+            # running when this thread is blocked
+            watchdog = @async oneL0.cooperative_wait(gate; spin=false) do gate
+                watching[] = true
+                res = oneL0.@gcsafe_ccall oneL0.libze_loader.zeEventHostSynchronize(
+                    gate::oneL0.ze_event_handle_t, timeout::UInt64)::oneL0.ze_result_t
+                res == oneL0.RESULT_NOT_READY || return false
+                oneL0.signal(gate)
+                return true
+            end
+            while !watching[]
+                yield()
+            end
+
+            oneL0.append_wait!(list, gate)
+            opener = @async begin
+                for _ in 1:10_000
+                    yield()
+                end
+                opened[] = true
+                oneL0.signal(gate)
+            end
+            f()
+            in_time = opened[]
+        finally
+            # also when `f` failed, as the gate is reused
+            oneL0.signal(gate)
+            @isdefined(opener) && wait(opener)
+            @isdefined(watchdog) && (timed_out = something(fetch(watchdog)))
+            GC.enable(gc_enabled)
+        end
+        synchronize()
+        return in_time && !timed_out
+    end
+
+    # compile everything beforehand, as that could wait for the GPU to become idle
+    a = oneArray{Int32}(undef, 64)
+    @oneapi items=64 fill_kernel(a, Int32(0))
+    synchronize()
+    @test gated(synchronize)
+
+    @test gated() do
+        oneL0.sync_each_submission(false) do
+            @oneapi items=64 fill_kernel(a, Int32(1))
+        end
+        synchronize()
+    end
+    @test Array(a) == fill(Int32(1), 64)
+
+    @test gated() do
+        oneL0.sync_each_submission(false) do
+            @oneapi items=64 fill_kernel(a, Int32(2))
+        end
+        synchronize(oneAPI.global_stream(context(), device()))
+    end
+    @test Array(a) == fill(Int32(2), 64)
+
+    # synchronizing after every launch, as on the LTS stack
+    @test gated() do
+        oneL0.sync_each_submission(true) do
+            @oneapi items=64 fill_kernel(a, Int32(3))
+        end
+    end
+    @test Array(a) == fill(Int32(3), 64)
+
+    # blocking synchronization is still available
+    @oneapi items=64 fill_kernel(a, Int32(4))
+    @test synchronize(; blocking=true) === nothing
+    @test Array(a) == fill(Int32(4), 64)
+end
+
 ############################################################################################
 
 # Keep allocation consumers at top level so kernels do not capture test state.