diff --git a/Project.toml b/Project.toml index 88b57ba4..2a992247 100644 --- a/Project.toml +++ b/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" diff --git a/lib/level-zero/oneL0.jl b/lib/level-zero/oneL0.jl index 82e33eec..c36b9eee 100644 --- a/lib/level-zero/oneL0.jl +++ b/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") diff --git a/lib/level-zero/synchronization.jl b/lib/level-zero/synchronization.jl new file mode 100644 index 00000000..e3a3a4e2 --- /dev/null +++ b/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 diff --git a/lib/utils/APIUtils.jl b/lib/utils/APIUtils.jl index d8d30394..61ec0743 100644 --- a/lib/utils/APIUtils.jl +++ b/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 diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 20742c2d..f129919b 100644 --- a/src/compiler/execution.jl +++ b/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 diff --git a/src/context.jl b/src/context.jl index 409a9d7a..6c9e188f 100644 --- a/src/context.jl +++ b/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 diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index 7b90d2ca..ae693caa 100644 --- a/src/oneAPIKernels.jl +++ b/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 diff --git a/test/execution.jl b/test/execution.jl index 3efcdbf1..c32cb176 100644 --- a/test/execution.jl +++ b/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.