From 6dd12633a64fdcd8e7cf221f5de3970abfd57cb4 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 11 Aug 2026 15:49:11 +0000 Subject: [PATCH] Submit work through per-task immediate command lists MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every kernel launch, copy, and fill used to create a fresh command list, submit it to the task's command queue, and drop the reference — leaving destruction of thousands of driver objects (lists, command buffers, heaps) to finalizer timing. Under launch storms that garbage is what pushes the driver into allocation failure, where NEO's error handling is at its worst (the scratch path aborts outright). It is also pure overhead: the per-dispatch list costs ~8x in submission latency. Replace the per-dispatch machinery with a per-task oneStream holding one in-order asynchronous immediate command list: appends submit directly, and the garbage source disappears entirely. Level Zero >= 1.9 is required; there is no fallback submission path. oneMKL work still needs a real command queue for SYCL interop, so each stream lazily creates a companion queue — a separate execution stream, which makes the previously implicit ordering between Julia kernels and oneMKL calls explicit: sycl_queue drains the immediate list before handing out the SYCL queue (Julia -> MKL), and a dirty flag makes the next Julia-side submission drain the companion queue (MKL -> Julia). FFT plans capture their queue at construction, so their _exec! methods apply the boundary themselves. The LTS drain-before-free machinery follows the shape change: the queue registry becomes a stream registry, draining both the immediate list and the companion queue before a buffer referenced by in-flight work is freed; immediate lists get the same bounded-drain finalizer as queues. The sync-each-submission workaround now host-synchronizes the list after each append. KA.priority! swaps the task's stream for one with the requested priority. The scratch hedge moves to the stream, and remains on the explicit-queue compatibility path (@oneapi queue=...), which still submits through a per-dispatch list. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WSFxSBXtckG3BVAT12wYkf --- docs/src/lts.md | 12 +- docs/src/troubleshooting.md | 13 ++ lib/level-zero/barrier.jl | 2 +- lib/level-zero/cmdlist.jl | 83 ++++++++++- lib/level-zero/cmdqueue.jl | 2 +- lib/level-zero/copy.jl | 8 +- lib/level-zero/event.jl | 6 +- lib/mkl/fft.jl | 9 ++ src/array.jl | 6 +- src/compiler/execution.jl | 33 +++-- src/context.jl | 275 +++++++++++++++++++++++++----------- src/memory.jl | 4 +- src/oneAPIKernels.jl | 37 +++-- src/pool.jl | 2 +- test/execution.jl | 62 +++++++- test/fft.jl | 12 ++ test/level-zero.jl | 52 +++++++ test/onemkl.jl | 39 +++++ 18 files changed, 516 insertions(+), 141 deletions(-) diff --git a/docs/src/lts.md b/docs/src/lts.md index 6251b10c..e2d4ba7d 100644 --- a/docs/src/lts.md +++ b/docs/src/lts.md @@ -93,9 +93,10 @@ references it. A garbage-collected free of a dead array whose last kernel has no then faults on the GPU, which bans the kernel context — after which *every* later submission fails with `ZE_RESULT_ERROR_UNKNOWN`. -On LTS, oneAPI.jl keeps a registry of the command queues in use and drains those that could -reference a buffer before freeing it. Command queues are likewise drained before being -destroyed; a queue still busy after 10 s is deliberately leaked, since destroying it would +On LTS, oneAPI.jl keeps a registry of the streams in use — each task's immediate command +list plus its companion command queue for oneMKL work — and drains those that could +reference a buffer before freeing it. Lists and queues are likewise drained before being +destroyed; one still busy after 10 s is deliberately leaked, since destroying it would trigger the very fault the drain prevents. The visible cost is that a GC-driven free can block until outstanding work completes. @@ -107,7 +108,10 @@ earlier, separately submitted command list. The result is a silent *dropped tail work-items of a kernel, or the last elements of a copy, never land. Synchronizing after every submission eliminates it, at roughly a 3× throughput cost. It is -off by default and enabled with: +off by default and enabled with the setting below. With the current immediate-command-list +submission path it host-synchronizes the stream after every append; the dropped-tail +failure was only ever observed on the earlier queue-submission path, so this workaround may +no longer be needed — it is kept until that is re-established under oversubscription. ```bash export ONEAPI_SYNC_EACH_SUBMISSION=1 diff --git a/docs/src/troubleshooting.md b/docs/src/troubleshooting.md index 7ff19357..22ffaf39 100644 --- a/docs/src/troubleshooting.md +++ b/docs/src/troubleshooting.md @@ -86,3 +86,16 @@ Enable debug mode in oneAPI.jl to use debug builds of underlying toolchains (if oneAPI.set_debug!(true) ``` + +### Scratch hedge + +Before the first submission of a kernel whose register spill exceeds the stream's +high-water mark, oneAPI.jl drains the stream and runs finalizers: the driver performs its +scratch-buffer allocation at that moment, and (at least through NEO 26.x) aborts the +process instead of erroring when it fails. Making the allocation happen at a clean moment +removes that failure mode for workloads under memory pressure. The drain costs one +synchronization per (stream, spill tier) — once per workload in practice. Disable it with: + +```bash +export ONEAPI_SCRATCH_HEDGE=0 +``` diff --git a/lib/level-zero/barrier.jl b/lib/level-zero/barrier.jl index 4f274102..36e2df8a 100644 --- a/lib/level-zero/barrier.jl +++ b/lib/level-zero/barrier.jl @@ -1,6 +1,6 @@ export append_barrier!, device_barrier -append_barrier!(list::ZeCommandList, signal_event=nothing, wait_events::ZeEvent...) = +append_barrier!(list::AbstractZeCommandList, signal_event=nothing, wait_events::ZeEvent...) = zeCommandListAppendBarrier(list, something(signal_event, C_NULL), length(wait_events), [wait_events...]) diff --git a/lib/level-zero/cmdlist.jl b/lib/level-zero/cmdlist.jl index e464c639..0e00f0d4 100644 --- a/lib/level-zero/cmdlist.jl +++ b/lib/level-zero/cmdlist.jl @@ -1,8 +1,10 @@ # list -export ZeCommandList, execute! +export ZeCommandList, ZeImmediateCommandList, AbstractZeCommandList, execute! -mutable struct ZeCommandList +abstract type AbstractZeCommandList end + +mutable struct ZeCommandList <: AbstractZeCommandList handle::ze_command_list_handle_t context::ZeContext @@ -22,10 +24,10 @@ mutable struct ZeCommandList end end -Base.unsafe_convert(::Type{ze_command_list_handle_t}, list::ZeCommandList) = list.handle +Base.unsafe_convert(::Type{ze_command_list_handle_t}, list::AbstractZeCommandList) = list.handle -Base.:(==)(a::ZeCommandList, b::ZeCommandList) = a.handle == b.handle -Base.hash(e::ZeCommandList, h::UInt) = hash(e.handle, h) +Base.:(==)(a::AbstractZeCommandList, b::AbstractZeCommandList) = a.handle == b.handle +Base.hash(e::AbstractZeCommandList, h::UInt) = hash(e.handle, h) Base.close(list::ZeCommandList) = zeCommandListClose(list) @@ -47,6 +49,62 @@ function ZeCommandList(f::Base.Callable, args...; kwargs...) return list end +""" + ZeImmediateCommandList(ctx::ZeContext, dev::ZeDevice, ordinal=1, index=1; + flags=0, mode=ZE_COMMAND_QUEUE_MODE_DEFAULT, + priority=ZE_COMMAND_QUEUE_PRIORITY_NORMAL) + +Create an immediate command list: appended commands are submitted to the device as they +are appended, with no separate close/execute step and no per-submission list object. +The descriptor is a command *queue* descriptor; pass +`flags=ZE_COMMAND_QUEUE_FLAG_IN_ORDER` and `mode=ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS` +for an asynchronous stream with in-order semantics (requires Level Zero >= 1.9). +Synchronize with [`synchronize`](@ref). +""" +mutable struct ZeImmediateCommandList <: AbstractZeCommandList + handle::ze_command_list_handle_t + + context::ZeContext + device::ZeDevice + ordinal::Int + + function ZeImmediateCommandList(ctx::ZeContext, dev::ZeDevice, ordinal=1, index=1; + flags=0, + mode::ze_command_queue_mode_t=ZE_COMMAND_QUEUE_MODE_DEFAULT, + priority::ze_command_queue_priority_t=ZE_COMMAND_QUEUE_PRIORITY_NORMAL) + desc_ref = Ref(ze_command_queue_desc_t(; + ordinal=ordinal-1, index=index-1, flags, mode, priority + )) + handle_ref = Ref{ze_command_list_handle_t}() + zeCommandListCreateImmediate(ctx, dev, desc_ref, handle_ref) + obj = new(handle_ref[], ctx, dev, ordinal) + finalizer(obj) do obj + # unlike a regular list, an immediate list can have work in flight at + # finalization on any stack, and destroying it then is illegal. Bounded + # unchecked drain, leaking the list on timeout — same rationale as the + # queue finalizer in cmdqueue.jl: an infinite wait on event-gated work + # whose event is never signaled would hang GC and process exit. + if unchecked_zeCommandListHostSynchronize(obj, FINALIZER_SYNC_TIMEOUT_NS) == RESULT_NOT_READY + @warn "Leaking an immediate command list still busy after $(FINALIZER_SYNC_TIMEOUT_NS ÷ 1_000_000_000)s to avoid blocking finalization" maxlog = 1 + return + end + zeCommandListDestroy(obj) + # mark destroyed: the stream registry can still reach this list, and + # synchronizing a destroyed handle crashes in the driver + obj.handle = ze_command_list_handle_t(C_NULL) + end + obj + end +end + +""" + synchronize(list::ZeImmediateCommandList, timeout=typemax(UInt64)) + +Block the host until all commands appended to the immediate command list have completed. +""" +synchronize(list::ZeImmediateCommandList, timeout::Number=typemax(UInt64)) = + zeCommandListHostSynchronize(list, timeout) + # Opt-in workaround for the Aurora LTS NEO stack (set ONEAPI_SYNC_EACH_SUBMISSION=1). # Under heavy multi-process oversubscription of a single tile, a whole-queue # `zeCommandQueueSynchronize` does not reliably retire the tail of an earlier, @@ -112,3 +170,18 @@ function execute!(f::Base.Callable, queue::ZeCommandQueue, fence=nothing; kwargs list = ZeCommandList(f, queue.context, queue.device, queue.ordinal; kwargs...) execute!(queue, [list], fence) end + +""" + execute!(list::ZeImmediateCommandList) do list + append_...!(list) + end + +Append operations to an immediate command list. Each append is submitted to the device +as it happens, so there is no separate close/execute step; the return value is that of +the do block. +""" +function execute!(f::Base.Callable, list::ZeImmediateCommandList) + ret = f(list) + sync_each_submission() && synchronize(list) + return ret +end diff --git a/lib/level-zero/cmdqueue.jl b/lib/level-zero/cmdqueue.jl index d32123a7..a308255a 100644 --- a/lib/level-zero/cmdqueue.jl +++ b/lib/level-zero/cmdqueue.jl @@ -50,7 +50,7 @@ mutable struct ZeCommandQueue zeCommandQueueDestroy(obj) if LTS[] # mark the queue as destroyed: it can still be weakly reachable (e.g. from - # the queue registry used by `synchronize_all_queues`), and synchronizing a + # the stream registry used by `synchronize_all_streams`), and synchronizing a # destroyed handle crashes in the driver. obj.handle = ze_command_queue_handle_t(C_NULL) end diff --git a/lib/level-zero/copy.jl b/lib/level-zero/copy.jl index d356eb57..5a63c9ea 100644 --- a/lib/level-zero/copy.jl +++ b/lib/level-zero/copy.jl @@ -2,22 +2,22 @@ export append_copy!, append_fill!, append_prefetch!, append_advise! -append_copy!(list::ZeCommandList, dst::Union{Ptr,ZePtr}, src::Union{Ptr,ZePtr}, +append_copy!(list::AbstractZeCommandList, dst::Union{Ptr,ZePtr}, src::Union{Ptr,ZePtr}, size::Integer, signal_event::Union{ZeEvent,Nothing}=nothing, wait_events::ZeEvent...) = zeCommandListAppendMemoryCopy(list, dst, src, size, something(signal_event, C_NULL), length(wait_events), [wait_events...]) -append_fill!(list::ZeCommandList, ptr::Union{Ptr,ZePtr}, pattern::Union{Ptr,ZePtr}, +append_fill!(list::AbstractZeCommandList, ptr::Union{Ptr,ZePtr}, pattern::Union{Ptr,ZePtr}, pattern_size::Integer, size::Integer, signal_event::Union{ZeEvent,Nothing}=nothing, wait_events::ZeEvent...) = zeCommandListAppendMemoryFill(list, ptr, pattern, pattern_size, size, something(signal_event, C_NULL), length(wait_events), [wait_events...]) -append_prefetch!(list::ZeCommandList, ptr::Union{Ptr,ZePtr}, size::Integer) = +append_prefetch!(list::AbstractZeCommandList, ptr::Union{Ptr,ZePtr}, size::Integer) = zeCommandListAppendMemoryPrefetch(list, ptr, size) -append_advise!(list::ZeCommandList, dev::ZeDevice, ptr::Union{Ptr,ZePtr}, size::Integer, +append_advise!(list::AbstractZeCommandList, dev::ZeDevice, ptr::Union{Ptr,ZePtr}, size::Integer, advise::ze_memory_advice_t) = zeCommandListAppendMemAdvise(list, dev, ptr, size, advise) diff --git a/lib/level-zero/event.jl b/lib/level-zero/event.jl index 7257828e..92460f2b 100644 --- a/lib/level-zero/event.jl +++ b/lib/level-zero/event.jl @@ -54,15 +54,15 @@ Base.:(==)(a::ZeEvent, b::ZeEvent) = a.handle == b.handle Base.hash(e::ZeEvent, h::UInt) = hash(e.handle, h) signal(event::ZeEvent) = zeEventHostSignal(event) -append_signal!(list::ZeCommandList, event::ZeEvent) = zeCommandListAppendSignalEvent(list, event) +append_signal!(list::AbstractZeCommandList, event::ZeEvent) = zeCommandListAppendSignalEvent(list, event) Base.wait(event::ZeEvent, timeout::Number=typemax(UInt64)) = zeEventHostSynchronize(event, timeout) -append_wait!(list::ZeCommandList, events::ZeEvent...) = +append_wait!(list::AbstractZeCommandList, events::ZeEvent...) = zeCommandListAppendWaitOnEvents(list, length(events), [events...]) Base.reset(event::ZeEvent) = zeEventHostReset(event) -append_reset!(list::ZeCommandList, event::ZeEvent) = zeCommandListAppendEventReset(list, event) +append_reset!(list::AbstractZeCommandList, event::ZeEvent) = zeCommandListAppendEventReset(list, event) function Base.isdone(event::ZeEvent) res = unchecked_zeEventQueryStatus(event) diff --git a/lib/mkl/fft.jl b/lib/mkl/fft.jl index 4429801b..c9802342 100644 --- a/lib/mkl/fft.jl +++ b/lib/mkl/fft.jl @@ -523,22 +523,31 @@ end # Execution helpers _rawptr(a::oneAPI.oneArray{T}) where T = reinterpret(Ptr{Cvoid}, pointer(a)) +# NOTE: plans capture their SYCL queue at construction and thus bypass `sycl_queue` — +# the usual Julia → MKL ordering boundary — at execution time, so every _exec! method +# must run `mkl_boundary!` itself before handing work to oneMKL. + function _exec!(p::cMKLFFTPlan{T,MKLFFT_FORWARD,true}, X::oneAPI.oneArray{T}) where T + oneAPI.mkl_boundary!() st = onemklDftComputeForward(p.handle, _rawptr(X)); st==0 || error("forward FFT failed ($st)"); X end function _exec!(p::cMKLFFTPlan{T,MKLFFT_INVERSE,true}, X::oneAPI.oneArray{T}) where T + oneAPI.mkl_boundary!() st = onemklDftComputeBackward(p.handle, _rawptr(X)); st==0 || error("inverse FFT failed ($st)"); X end function _exec!(p::cMKLFFTPlan{T,K,false}, X::oneAPI.oneArray{T}, Y::oneAPI.oneArray{T}) where {T,K} + oneAPI.mkl_boundary!() st = (K==MKLFFT_FORWARD ? onemklDftComputeForwardOutOfPlace : onemklDftComputeBackwardOutOfPlace)(p.handle, _rawptr(X), _rawptr(Y)); st==0 || error("FFT failed ($st)"); Y end # Real forward function _exec!(p::rMKLFFTPlan{T,MKLFFT_FORWARD,false}, X::oneAPI.oneArray{T}, Y::oneAPI.oneArray{Complex{T}}) where T + oneAPI.mkl_boundary!() st = onemklDftComputeForwardOutOfPlace(p.handle, _rawptr(X), _rawptr(Y)); st==0 || error("rfft failed ($st)"); Y end # Real inverse (complex -> real) function _exec!(p::rMKLFFTPlan{T,MKLFFT_INVERSE,false}, X::oneAPI.oneArray{T}, Y::oneAPI.oneArray{R}) where {R,T<:Complex{R}} + oneAPI.mkl_boundary!() st = onemklDftComputeBackwardOutOfPlace(p.handle, _rawptr(X), _rawptr(Y)); st==0 || error("brfft failed ($st)"); Y end diff --git a/src/array.jl b/src/array.jl index 6fdeaef6..e542c730 100644 --- a/src/array.jl +++ b/src/array.jl @@ -462,7 +462,7 @@ function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, end # copies to the host are synchronizing - synchronize(global_queue(context(src), device())) + synchronize(global_stream(context(src), device())) return dest end @@ -482,7 +482,7 @@ end function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::oneDenseArray{T,<:Any,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}, doffs, src::Array{T}, soffs, n) where T # maintain queue-ordered semantics - synchronize(global_queue(ctx, dev)) + synchronize(global_stream(ctx, dev)) if Base.isbitsunion(T) # copy selector bytes @@ -503,7 +503,7 @@ end function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::Array{T}, doffs, src::oneDenseArray{T,<:Any,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}, soffs, n) where T # maintain queue-ordered semantics - synchronize(global_queue(ctx, dev)) + synchronize(global_stream(ctx, dev)) if Base.isbitsunion(T) # copy selector bytes diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 1f04465e..f80e2ebb 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -29,7 +29,9 @@ launches the kernel on the GPU. ## Launch Keywords (runtime) - `groups`: Number of workgroups (required). Can be an integer or tuple. - `items`: Number of work-items per workgroup (required). Can be an integer or tuple. -- `queue::ZeCommandQueue=global_queue(...)`: Command queue to submit to. +- `queue=global_stream(...)`: Submission target — the task's `oneStream` by default. An + explicit `ZeCommandQueue` is also accepted and submits through a per-dispatch command + list. # Examples @@ -322,38 +324,53 @@ end const _kernel_instances = Dict{UInt, Any}() @inline function onecall(kernel::ZeKernel, tt, args...; groups::ZeDim=1, items::ZeDim=1, - queue::ZeCommandQueue=global_queue(context(), device())) + queue::Union{oneStream, ZeCommandQueue}=global_stream(context(), device())) for (i, arg) in enumerate(args) oneL0.arguments(kernel)[i] = arg end groupsize!(kernel, items) + launch!(queue, kernel, groups) +end + +@inline function launch!(s::oneStream, kernel::ZeKernel, groups::ZeDim) + mkl_wait!(s) - # NEO allocates a queue's scratch buffer at the first submission of a kernel whose + # NEO allocates a stream's scratch buffer at the first submission of a kernel whose # spill exceeds what is already allocated, and that allocation aborts the process on # failure (no null check). Cross each new spill high-water mark deliberately, at the # cleanest reachable moment, instead of at a GC-lottery-determined one. spill = oneL0.spill_mem_size(kernel) - spill > queue.scratch_hwm && scratch_hedge!(queue, spill) + spill > s.scratch_hwm && scratch_hedge!(s, spill) + + append_launch!(s.list, kernel, groups) + oneL0.sync_each_submission() && oneL0.synchronize(s.list) + return +end +# explicit-queue compatibility: `@oneapi queue=...` submits through a per-dispatch +# command list, as it always has +@inline function launch!(queue::ZeCommandQueue, kernel::ZeKernel, groups::ZeDim) + spill = oneL0.spill_mem_size(kernel) + spill > queue.scratch_hwm && scratch_hedge!(queue, spill) execute!(queue) do list append_launch!(list, kernel, groups) end end -# Slow path of the scratch hedge, firing once per (queue, spill tier): retire in-flight +# Slow path of the scratch hedge, firing once per (stream, spill tier): retire in-flight # work, flush deferred releases, and run finalizers so dead driver objects and arrays are # destroyed before NEO performs its null-check-free scratch allocation. Opt out with # ONEAPI_SCRATCH_HEDGE=0; the high-water mark is maintained regardless, so the toggle # only skips the drain. -@noinline function scratch_hedge!(queue::ZeCommandQueue, spill::Int) +@noinline function scratch_hedge!(target::Union{oneStream, ZeCommandQueue}, spill::Int) if oneL0.SCRATCH_HEDGE[] - oneL0.synchronize(queue) + oneL0.synchronize(target) oneL0._run_reclaim_callbacks() GC.gc(false) Threads.atomic_add!(oneL0.SCRATCH_HEDGE_COUNT, 1) end - queue.scratch_hwm = spill + target.scratch_hwm = spill return end diff --git a/src/context.jl b/src/context.jl index f625a40d..0a336b19 100644 --- a/src/context.jl +++ b/src/context.jl @@ -6,7 +6,8 @@ # XXX: rework this -- it doesn't work well when altering the state -export driver, driver!, device, device!, context, context!, global_queue, synchronize, is_integrated +export driver, driver!, device, device!, context, context!, global_stream, global_queue, + synchronize, is_integrated """ driver() -> ZeDriver @@ -194,100 +195,148 @@ function context!(ctx::ZeContext) task_local_storage(:ZeContext, ctx) end -""" - global_queue(ctx::ZeContext, dev::ZeDevice) -> ZeCommandQueue - -Get the global command queue for the given context and device. This queue is used as the -default queue for executing operations, guaranteeing expected semantics when using a device -on a Julia task. - -The queue is created with in-order execution flags, meaning commands are executed in the -order they are submitted. Queues are cached per task and (context, device) pair. +# The per-task submission target. `list` is where Julia-side work (kernels, copies, +# fills) is appended — an in-order asynchronous immediate command list, so appends are +# submitted to the device as they happen and no per-dispatch driver objects exist. +# `queue` exists only for SYCL/oneMKL interop, which requires a real command-queue +# handle, and is created on first use. Work on the two executes independently, so +# ordering at the oneMKL boundary is restored explicitly: `mkl_boundary!` (Julia → MKL) +# drains the list before handing out the SYCL queue, and `mkl_dirty` makes the next +# Julia-side submission drain the queue (MKL → Julia, see `mkl_wait!`). +mutable struct oneStream + const ctx::ZeContext + const dev::ZeDevice + const list::oneL0.ZeImmediateCommandList + queue::Union{Nothing, ZeCommandQueue} + mkl_dirty::Bool + # high-water mark of per-thread spill (bytes) among kernels submitted to this + # stream, maintained by the scratch hedge (see src/compiler/execution.jl) + scratch_hwm::Int + const priority::oneL0.ze_command_queue_priority_t +end -# Arguments -- `ctx::ZeContext`: The context for the command queue. -- `dev::ZeDevice`: The device for the command queue. +function create_stream(ctx::ZeContext, dev::ZeDevice, + priority::oneL0.ze_command_queue_priority_t = + oneL0.ZE_COMMAND_QUEUE_PRIORITY_NORMAL) + # In-order immediate command lists entered the spec in 1.9, but the reported API + # version is a floor, not a feature inventory: the Aurora LTS driver reports 1.6 + # while implementing them (they are DPC++'s production submission path on PVC). + # Probe by creating — a driver without support rejects the flag — since there is + # deliberately no fallback submission path. + list = try + oneL0.ZeImmediateCommandList(ctx, dev; + flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, + mode = oneL0.ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, + priority) + catch err + err isa oneL0.ZeError || rethrow() + error("oneAPI.jl requires driver support for in-order immediate command lists " * + "(Level Zero >= 1.9, or a driver implementing them regardless of its " * + "reported API version $(oneL0.api_version(ctx.driver))); creation failed " * + "with $(err.code)") + end + s = oneStream(ctx, dev, list, nothing, false, 0, priority) + return register_stream!(ctx, dev, s) +end -# Returns -- `ZeCommandQueue`: A cached command queue with in-order execution. +""" + global_stream(ctx::ZeContext, dev::ZeDevice) -> oneStream -# Examples -```julia -ctx = context() -dev = device() -queue = global_queue(ctx, dev) -``` +Get the stream all oneAPI.jl operations of the calling task target for the given context +and device: an in-order asynchronous immediate command list for kernels, copies and +fills, plus a lazily-created companion command queue for SYCL/oneMKL interop. Streams +are cached per task and (context, device) pair. -See also: `context`, `device`, `synchronize` +See also: [`global_queue`](@ref), [`synchronize`](@ref) """ -function global_queue(ctx::ZeContext, dev::ZeDevice) +function global_stream(ctx::ZeContext, dev::ZeDevice) # NOTE: dev purposefully does not default to context() or device() to stress that # objects should track ownership, and not rely on implicit global state. - get!(task_local_storage(), (:ZeCommandQueue, ctx, dev)) do - queue = ZeCommandQueue(ctx, dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER) - register_queue!(ctx, dev, queue) - end + get!(task_local_storage(), (:oneStream, ctx, dev)) do + create_stream(ctx, dev) + end::oneStream +end + +# the companion command queue of a stream, created on first use. Only the SYCL/oneMKL +# interop path needs one; pure-Julia workloads never create a queue. +function stream_queue(s::oneStream) + q = s.queue + q === nothing || return q + s.queue = ZeCommandQueue(s.ctx, s.dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, + priority = s.priority) + return s.queue::ZeCommandQueue end -# Register `queue` as a queue targeting (ctx, dev) so `synchronize_all_queues`/`release` -# can find and drain it before freeing buffers whose in-flight work it may still reference. -# EVERY queue that becomes a task's active queue must go through here — not just the one -# `global_queue` creates but also the replacement `KA.priority!` installs — or the -# unregistered queue's in-flight work can outlive a freed buffer (a use-after-free that -# faults and bans the context on the LTS NEO stack). Only the LTS stack maintains the -# registry; on the rolling stack this is a no-op. Returns `queue`. -function register_queue!(ctx::ZeContext, dev::ZeDevice, queue::ZeCommandQueue) - oneL0.LTS[] || return queue +""" + global_queue(ctx::ZeContext, dev::ZeDevice) -> ZeCommandQueue + +Get the calling task's companion command queue for the given context and device — the +queue oneMKL work is enqueued on through SYCL interop. Julia-side kernels, copies and +fills do not use it; they are appended to the task's [`global_stream`](@ref) immediate +command list instead. +""" +global_queue(ctx::ZeContext, dev::ZeDevice) = stream_queue(global_stream(ctx, dev)) + +# Register `stream` as a stream targeting (ctx, dev) so `synchronize_all_streams`/ +# `release` can find and drain it before freeing buffers whose in-flight work it may +# still reference. EVERY stream that becomes a task's active stream must go through here +# — not just the one `global_stream` creates but also the replacement `KA.priority!` +# installs — or the unregistered stream's in-flight work can outlive a freed buffer (a +# use-after-free that faults and bans the context on the LTS NEO stack). Only the LTS +# stack maintains the registry; on the rolling stack this is a no-op. Returns `stream`. +function register_stream!(ctx::ZeContext, dev::ZeDevice, stream::oneStream) + oneL0.LTS[] || return stream # disable finalizers while mutating the registry: a GC-driven finalizer on this - # task could call back into `synchronize_all_queues` (the lock is reentrant) and + # task could call back into `synchronize_all_streams` (the lock is reentrant) and # observe/mutate the registry mid-update. GC.enable_finalizers(false) try - @lock queue_registry_lock begin + @lock stream_registry_lock begin push!( - get!(Vector{Tuple{WeakRef, ZeCommandQueue}}, queue_registry, (ctx, dev)), - (WeakRef(current_task()), queue) + get!(Vector{Tuple{WeakRef, oneStream}}, stream_registry, (ctx, dev)), + (WeakRef(current_task()), stream) ) end finally GC.enable_finalizers(true) end - return queue + return stream end -# Registry of all queues created through `global_queue`, across tasks. Buffers can be -# freed from any task (GC finalizers), so `release` needs to be able to find the queues -# that may still have work in flight referencing the buffer; queues themselves are +# Registry of all streams created through `global_stream`, across tasks. Buffers can be +# freed from any task (GC finalizers), so `release` needs to be able to find the streams +# that may still have work in flight referencing the buffer; streams themselves are # cached task-locally and would otherwise be unreachable from the finalizing task. # -# Entries reference the queue *strongly*: the GC clears WeakRefs to a dead queue in the -# same cycle that queues its finalizer, i.e., before the finalizer runs, so a WeakRef -# would hide the queue from `release` exactly when its in-flight work still references -# buffers about to be freed. The owning task is tracked weakly instead: queues are -# task-local, so once their task is dead no new work can reach them, and the entry can -# be dropped (allowing the queue to be finalized) after a final synchronize. -const queue_registry_lock = ReentrantLock() -const queue_registry = Dict{Tuple{ZeContext, ZeDevice}, Vector{Tuple{WeakRef, ZeCommandQueue}}}() - -# synchronize all known queues that target the given context (and device, if specified), -# i.e., all queues whose in-flight work could possibly reference an allocation that is -# about to be freed. -function synchronize_all_queues(ctx::ZeContext, dev::Union{ZeDevice, Nothing}) - # only the LTS stack populates the queue registry (see `global_queue`); on the +# Entries reference the stream *strongly*: the GC clears WeakRefs to a dead stream in +# the same cycle that queues its members' finalizers, i.e., before they run, so a +# WeakRef would hide the stream from `release` exactly when its in-flight work still +# references buffers about to be freed. The owning task is tracked weakly instead: +# streams are task-local, so once their task is dead no new work can reach them, and +# the entry can be dropped (allowing list and queue to be finalized) after a final +# synchronize. +const stream_registry_lock = ReentrantLock() +const stream_registry = Dict{Tuple{ZeContext, ZeDevice}, Vector{Tuple{WeakRef, oneStream}}}() + +# synchronize all known streams that target the given context (and device, if +# specified), i.e., all streams whose in-flight work could possibly reference an +# allocation that is about to be freed. Drains both each stream's immediate command +# list and its companion queue (oneMKL work). +function synchronize_all_streams(ctx::ZeContext, dev::Union{ZeDevice, Nothing}) + # only the LTS stack populates the stream registry (see `global_stream`); on the # rolling stack this is a no-op and `release` frees directly. oneL0.LTS[] || return - queues = ZeCommandQueue[] - stale = Tuple{WeakRef, ZeCommandQueue}[] + streams = oneStream[] + stale = Tuple{WeakRef, oneStream}[] GC.enable_finalizers(false) try - @lock queue_registry_lock begin - for ((qctx, qdev), entries) in queue_registry - qctx == ctx || continue - (dev === nothing || qdev == dev) || continue + @lock stream_registry_lock begin + for ((sctx, sdev), entries) in stream_registry + sctx == ctx || continue + (dev === nothing || sdev == dev) || continue for entry in entries - (task, queue) = entry - queue.handle == C_NULL && continue # finalized, handle destroyed - push!(queues, queue) + (task, stream) = entry + push!(streams, stream) # entries whose task was already dead at this point cannot # receive new work, so they are safe to retire after the sync if task.value === nothing || istaskdone(task.value::Task) @@ -298,19 +347,22 @@ function synchronize_all_queues(ctx::ZeContext, dev::Union{ZeDevice, Nothing}) end # synchronize outside the lock: this can block for as long as a kernel runs, # and finalizers running concurrently also need to take the lock. Keep - # finalizers disabled so none of the collected queues can be destroyed - # between collection and synchronization. - for queue in queues - oneL0.synchronize(queue) + # finalizers disabled so no stream member can be destroyed between collection + # and synchronization; the null-handle checks are defense in depth against + # lists/queues finalized before their stream was registered stale. + for stream in streams + stream.list.handle == C_NULL || oneL0.synchronize(stream.list) + q = stream.queue + (q === nothing || q.handle == C_NULL) || oneL0.synchronize(q) end - # retire drained queues of dead tasks, allowing them to be finalized (the - # finalizer synchronizes once more before destroying the queue, in case - # the queue is dropped through other means). + # retire drained streams of dead tasks, allowing their list and queue to be + # finalized (the finalizers synchronize once more before destroying, in case + # the stream is dropped through other means). if !isempty(stale) - @lock queue_registry_lock begin - for ((qctx, qdev), entries) in queue_registry - qctx == ctx || continue - (dev === nothing || qdev == dev) || continue + @lock stream_registry_lock begin + for ((sctx, sdev), entries) in stream_registry + sctx == ctx || continue + (dev === nothing || sdev == dev) || continue filter!(entry -> !any(s -> s === entry, stale), entries) end end @@ -323,9 +375,11 @@ end """ synchronize() + synchronize(stream::oneStream) -Block the host thread until all operations on the global command queue for the current -context and device have completed. +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. This is useful for timing operations or ensuring that GPU work has finished before accessing results on the CPU. @@ -338,10 +392,59 @@ synchronize() # Wait for GPU computation to complete println("GPU work completed") ``` -See also: [`global_queue`](@ref), [`context`](@ref), [`device`](@ref) +See also: [`global_stream`](@ref), [`context`](@ref), [`device`](@ref) """ +function oneL0.synchronize(s::oneStream) + oneL0.synchronize(s.list) + q = s.queue + if q !== nothing + oneL0.synchronize(q) + s.mkl_dirty = false + end + return +end + function oneL0.synchronize() - oneL0.synchronize(global_queue(context(), device())) + oneL0.synchronize(global_stream(context(), device())) +end + +# Julia → MKL ordering: everything Julia appended to the task's immediate list must be +# visible before oneMKL work is enqueued on the companion queue, which is a separate +# stream from the driver's point of view. Runs on every `sycl_queue` access; oneMKL +# wrappers must evaluate `sycl_queue(...)` only after all device-side argument +# preparation (temporaries, conversions), which holds today because the queue is the +# first ccall argument and Julia evaluates arguments left to right. +function mkl_boundary!(s::oneStream = global_stream(context(), device())) + oneL0.synchronize(s.list) + s.mkl_dirty = true + return +end + +# MKL → Julia ordering: consumed at the head of every Julia-side submission. One Bool +# load on the fast path; only a preceding oneMKL call makes it synchronize. The queue +# can be absent with the flag set: an FFT plan executing on this task runs on the queue +# it captured at construction, which need not be this stream's companion queue (whose +# creation the flag does not force). +@inline function mkl_wait!(s::oneStream) + if s.mkl_dirty + q = s.queue + q === nothing || oneL0.synchronize(q) + s.mkl_dirty = false + end + return +end + +""" + execute!(stream::oneStream) do list + append_...!(list) + end + +Append operations to the stream's immediate command list, after waiting for any +outstanding oneMKL work. Appends are submitted to the device as they happen. +""" +@inline function oneL0.execute!(f::Base.Callable, s::oneStream) + mkl_wait!(s) + oneL0.execute!(f, s.list) end # re-export and augment parts of oneL0 to make driver and device selection easier @@ -394,10 +497,16 @@ function sycl_context(ctx=context(), dev=device()) end end +# Hands out the task's SYCL queue (wrapping the stream's companion command queue) for +# an imminent oneMKL call, which is why this is also the Julia → MKL ordering boundary: +# `mkl_boundary!` drains the immediate command list on EVERY access, so device work the +# wrappers prepared beforehand is complete before oneMKL work is enqueued. function sycl_queue(queue) + s = global_stream(queue.context, queue.device) + mkl_boundary!(s) get!(task_local_storage(), (:SYCLQueue, queue.context, queue.device)) do syclQueue(sycl_context(queue.context, queue.device), sycl_device(queue.device), - global_queue(queue.context, queue.device)) + stream_queue(s)) end end diff --git a/src/memory.jl b/src/memory.jl index 61b17106..4cb61970 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -25,7 +25,7 @@ function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dst::Union{Ptr{T},Ze src::Union{Ptr{T},ZePtr{T}}, N::Integer) where T bytes = N*sizeof(T) bytes==0 && return - execute!(global_queue(ctx, dev)) do list + execute!(global_stream(ctx, dev)) do list append_copy!(list, dst, src, bytes) end end @@ -54,7 +54,7 @@ function unsafe_fill!(ctx::ZeContext, dev::ZeDevice, ptr::Union{Ptr{T},ZePtr{T}} pattern::Union{Ptr{T},ZePtr{T}}, N::Integer) where T bytes = N*sizeof(T) bytes==0 && return - execute!(global_queue(ctx, dev)) do list + execute!(global_stream(ctx, dev)) do list append_fill!(list, ptr, pattern, sizeof(T), bytes) end end diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index 0a4b70c8..05e96285 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -250,27 +250,22 @@ function KA.priority!(::oneAPIBackend, prio::Symbol) ctx = oneAPI.context() dev = oneAPI.device() - # Update the cached queue - # We synchronize the current queue first to ensure safety - current_queue = oneAPI.global_queue(ctx, dev) - oneAPI.oneL0.synchronize(current_queue) - - # Replace the queue in task_local_storage - # The key used by global_queue is (:ZeCommandQueue, ctx, dev) - - new_queue = oneAPI.oneL0.ZeCommandQueue( - ctx, dev; - flags = oneAPI.oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, - priority = priority_enum - ) - - # Register the replacement queue so `synchronize_all_queues`/`release` can drain it - # before freeing a buffer whose in-flight work it references; otherwise all work after - # `priority!` runs on an unregistered queue and a freed buffer can be reused while its - # kernel is still running (use-after-free → banned context on the LTS NEO stack). - oneAPI.register_queue!(ctx, dev, new_queue) - - task_local_storage((:ZeCommandQueue, ctx, dev), new_queue) + # drain the task's current stream before swapping it out, so operations submitted + # to the new stream cannot overtake in-flight work on the old one + oneAPI.oneL0.synchronize(oneAPI.global_stream(ctx, dev)) + + # Replace the stream in task_local_storage. `create_stream` registers the + # replacement so `synchronize_all_streams`/`release` can drain it before freeing a + # buffer whose in-flight work it references; otherwise all work after `priority!` + # runs on an unregistered stream and a freed buffer can be reused while its kernel + # is still running (use-after-free → banned context on the LTS NEO stack). The old + # stream stays registered until its task dies, like replaced queues before it. + new_stream = oneAPI.create_stream(ctx, dev, priority_enum) + task_local_storage((:oneStream, ctx, dev), new_stream) + + # the cached SYCL queue wraps the old stream's companion queue; drop it so the next + # oneMKL call recreates it against the new stream (the old one was just drained) + delete!(task_local_storage(), (:SYCLQueue, ctx, dev)) return nothing end diff --git a/src/pool.jl b/src/pool.jl index d67f58c8..8d390234 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -119,7 +119,7 @@ function release(buf::oneL0.AbstractBuffer) # reference this buffer before freeing. (No-op on the rolling stack, which honors # BLOCKING_FREE.) if oneL0.LTS[] - synchronize_all_queues(oneL0.context(buf), oneL0.device(buf)) + synchronize_all_streams(oneL0.context(buf), oneL0.device(buf)) end free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_BLOCKING_FREE) diff --git a/test/execution.jl b/test/execution.jl index 8f11ae3e..1f44c128 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -657,8 +657,8 @@ end # queues are task-local, so a fresh task gets a fresh queue with a zero # high-water mark; observe there, assert on the test task observed = fetch(@async begin - q = global_queue(context(), device()) - hwm0 = q.scratch_hwm + s = global_stream(context(), device()) + hwm0 = s.scratch_hwm c0 = oneL0.SCRATCH_HEDGE_COUNT[] @oneapi items=64 groups=4 hedge_spill_kernel(out, a, Val(256), Val(2)) c1 = oneL0.SCRATCH_HEDGE_COUNT[] @@ -670,7 +670,7 @@ end @oneapi dummy() end c3 = oneL0.SCRATCH_HEDGE_COUNT[] - (hwm0, q.scratch_hwm, c1 - c0, c2 - c1, c3 - c2) + (hwm0, s.scratch_hwm, c1 - c0, c2 - c1, c3 - c2) end) @test observed == (0, spill, 1, 0, 0) @@ -679,10 +679,10 @@ end oneL0.SCRATCH_HEDGE[] = false try observed = fetch(@async begin - q = global_queue(context(), device()) + s = global_stream(context(), device()) c0 = oneL0.SCRATCH_HEDGE_COUNT[] @oneapi items=64 groups=4 hedge_spill_kernel(out, a, Val(256), Val(2)) - (oneL0.SCRATCH_HEDGE_COUNT[] - c0, q.scratch_hwm) + (oneL0.SCRATCH_HEDGE_COUNT[] - c0, s.scratch_hwm) end) @test observed == (0, spill) finally @@ -690,3 +690,55 @@ end end end end + +# ordering probe for the immediate submission stream +function stream_iota_kernel(a, off) + i = get_global_id() + @inbounds a[i] = i + off + return +end + +@testset "immediate submission stream" begin + n = 1024 + a = oneAPI.zeros(Int32, n) + host = zeros(Int32, n) + expected = zeros(Int32, n) + + # in-order: kernel write then synchronizing D2H readback, iterated + ok = true + for iter in 1:100 + @oneapi items=256 groups=4 stream_iota_kernel(a, Int32(iter)) + copyto!(host, a) # copies to the host synchronize the stream + expected .= Int32.(1:n) .+ Int32(iter) + ok &= host == expected + end + @test ok + + # a raw command queue is still accepted as an explicit submission target + queue = oneL0.ZeCommandQueue(context(), device(); + flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER) + k = @oneapi launch=false stream_iota_kernel(a, Int32(0)) + k(a, Int32(0); items=256, groups=4, queue) + oneL0.synchronize(queue) + @test Array(a) == Int32.(1:n) + + # the LTS sync-each-submission knob applies to immediate lists + oneL0.sync_each_submission(true) do + @oneapi items=256 groups=4 stream_iota_kernel(a, Int32(41)) + end + @test Array(a) == Int32.(1:n) .+ Int32(41) + + # concurrent tasks each get their own stream; alloc/free churn while launching + # exercises the cross-task drain-before-free path on the LTS stack + results = map(fetch, [@async begin + b = oneAPI.ones(Float32, 4096) + acc = true + for i in 1:50 + c = b .* Float32(i) + acc &= isapprox(sum(c), Float32(4096 * i); rtol=1e-4) + oneAPI.unsafe_free!(c) + end + acc + end for _ in 1:2]) + @test all(results) +end diff --git a/test/fft.jl b/test/fft.jl index d4419462..d8b37f85 100644 --- a/test/fft.jl +++ b/test/fft.jl @@ -102,4 +102,16 @@ end @test p2.queue == cached_handle cmp(p2 * dX2, fft(X2)) end + +@testset "stream interleave" begin + # FFT plans capture their SYCL queue at construction and execute through _exec!, + # which must apply the Julia → MKL ordering boundary itself. Broadcast → fft → + # broadcast with no intermediate synchronization. + X = gpu(rand(ComplexF32, 256)) + hX = Array(X) + X .= X .* 2f0 + Y = fft(X) + Z = abs.(Y) + cmp(Z, abs.(fft(hX .* 2f0))) +end end diff --git a/test/level-zero.jl b/test/level-zero.jl index 0fbc87db..bd070d8f 100644 --- a/test/level-zero.jl +++ b/test/level-zero.jl @@ -321,6 +321,58 @@ end end +@testset "immediate command list" begin + +# probe rather than version-gate: the reported API version is a floor, not a feature +# inventory (the Aurora LTS driver reports 1.6 while implementing these) +ilist = try + ZeImmediateCommandList(ctx, dev, group.ordinal; + flags=oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, + mode=oneL0.ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS) +catch err + err isa oneL0.ZeError || rethrow() + nothing +end +if ilist === nothing + @test_skip "driver rejects in-order immediate command lists" +else + # NOTE: no zeCommandListIsImmediate sanity check here — that entrypoint is spec 1.9 + # and missing from older loaders (Aurora LTS reports 1.6), where the ccall would + # fail on symbol lookup rather than with a ZeError. Immediacy is proven + # functionally: the appends below execute without any close/execute step. + + let src = rand(Int, 1024) + chk = zeros(Int, length(src)) + dst = device_alloc(ctx, dev, sizeof(src)) + + # appends submit immediately and execute in order; host-synchronize to observe + append_copy!(ilist, pointer(dst), pointer(src), sizeof(src)) + append_copy!(ilist, pointer(chk), pointer(dst), sizeof(src)) + synchronize(ilist) + @test chk == src + + free(dst) + end + + ret = execute!(ilist) do list + @test list isa ZeImmediateCommandList + 42 + end + @test ret == 42 + + # events and barriers work on immediate lists + pool = ZeEventPool(ctx, 1; flags=oneL0.ZE_EVENT_POOL_FLAG_HOST_VISIBLE) + event = pool[1] + append_barrier!(ilist, event) + wait(event, 10_000_000_000) # bounded: a hang here should fail, not stall the suite + @test Base.isdone(event) + + synchronize(ilist) +end + +end + + @testset "residency" begin diff --git a/test/onemkl.jl b/test/onemkl.jl index f0383962..0f743f69 100644 --- a/test/onemkl.jl +++ b/test/onemkl.jl @@ -1561,4 +1561,43 @@ end end end +@testset "stream interleave" begin + # Julia kernels run on the task's immediate command list, oneMKL work on the + # companion queue; the boundary in `sycl_queue`/`mkl_wait!` must order the two. + # Broadcast → gemm → broadcast, iterated with no intermediate synchronization. + T = Float32 + interleave_n = 64 + A0 = rand(T, interleave_n, interleave_n) + B0 = rand(T, interleave_n, interleave_n) + A = oneArray(A0) + B = oneArray(B0) + C = oneAPI.zeros(T, interleave_n, interleave_n) + hA = copy(A0) + hC = zeros(T, interleave_n, interleave_n) + for i in 1:100 + A .= A .+ T(1 / i) + mul!(C, A, B) + C .= C .+ T(1) + + hA .= hA .+ T(1 / i) + hC .= hA * B0 .+ T(1) + end + @test Array(C) ≈ hC rtol=1e-4 + + # the dirty flag makes the next Julia-side submission wait on oneMKL work + s = oneAPI.global_stream(oneAPI.context(), oneAPI.device()) + mul!(C, A, B) + @test s.mkl_dirty + oneAPI.oneL0.synchronize() + @test !s.mkl_dirty + + # scalar-result path (synchronizes inside the support library) between kernels + x = oneArray(rand(T, 1000)) + y = oneArray(rand(T, 1000)) + x .= x .+ T(1) + d = dot(x, y) + z = x .* y + @test d ≈ sum(Array(z)) rtol=1e-3 +end + end # oneMKL tests