From 460005e9681f318692ef93513b810524d1ba578e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 15:21:52 +0200 Subject: [PATCH 1/2] Serialize submissions to a stream's immediate list Kernel launches and `execute!` appended to the stream's immediate list directly, and other code reached into it as well. Give the stream a lock that every append holds, and rename the list so that code can't append to it without going through `launch!` or `execute!`. This prepares for batching launches, which needs all submissions to a stream to be ordered with the batch. --- src/compiler/execution.jl | 4 ++-- src/context.jl | 14 ++++++++------ src/memory.jl | 2 +- test/execution.jl | 5 +++-- 4 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 614538c8..fd71e7cc 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -403,9 +403,9 @@ end spill = oneL0.spill_mem_size(kernel) spill > s.scratch_hwm && scratch_hedge!(s, spill) - append_launch!(s.list, kernel, groups) + @lock s.lock append_launch!(s.immediate, kernel, groups) # wait cooperatively, as `synchronize` does, or the workaround blocks the thread - oneL0.sync_each_submission() && oneL0.nonblocking_synchronize(s.list) + oneL0.sync_each_submission() && oneL0.nonblocking_synchronize(s.immediate) return end diff --git a/src/context.jl b/src/context.jl index 1ad4ae34..6af614b7 100644 --- a/src/context.jl +++ b/src/context.jl @@ -210,7 +210,9 @@ end mutable struct oneStream const ctx::ZeContext const dev::ZeDevice - const list::oneL0.ZeImmediateCommandList + # only append to this through `execute!` or `launch!`, which hold `lock` + const immediate::oneL0.ZeImmediateCommandList + const lock::ReentrantLock # atomic because other tasks read it when waiting for work on memory this stream owns Base.@atomic queue::Union{Nothing, ZeCommandQueue} mkl_dirty::Bool @@ -240,7 +242,7 @@ function create_stream(ctx::ZeContext, dev::ZeDevice, "reported API version $(oneL0.api_version(ctx.driver))); creation failed " * "with $(err.code)") end - s = oneStream(ctx, dev, list, nothing, false, 0, priority) + s = oneStream(ctx, dev, list, ReentrantLock(), nothing, false, 0, priority) return register_stream!(ctx, dev, s) end @@ -357,7 +359,7 @@ function synchronize_all_streams(ctx::ZeContext, dev::Union{ZeDevice, Nothing}) # 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) + stream.immediate.handle == C_NULL || oneL0.synchronize(stream.immediate) q = Base.@atomic stream.queue (q === nothing || q.handle == C_NULL) || oneL0.synchronize(q) end @@ -405,7 +407,7 @@ See also: [`global_stream`](@ref), [`context`](@ref), [`device`](@ref) """ function oneL0.synchronize(s::oneStream; blocking::Bool=false) sync = blocking ? oneL0.synchronize : oneL0.nonblocking_synchronize - sync(s.list) + sync(s.immediate) q = Base.@atomic s.queue if q !== nothing sync(q) @@ -425,7 +427,7 @@ end # 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) + oneL0.synchronize(s.immediate) s.mkl_dirty = true return end @@ -454,7 +456,7 @@ 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) + @lock s.lock oneL0.execute!(f, s.immediate) end # re-export and augment parts of oneL0 to make driver and device selection easier diff --git a/src/memory.jl b/src/memory.jl index 8208026c..72a0a159 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -85,7 +85,7 @@ Base.sizeof(managed::Managed) = sizeof(managed.mem) # wait for all work on a stream, possibly from another task. unlike `synchronize`, this # leaves `mkl_dirty` alone, which only the stream's own task may modify. function drain(s::oneStream) - oneL0.nonblocking_synchronize(s.list) + oneL0.nonblocking_synchronize(s.immediate) q = Base.@atomic s.queue q === nothing || oneL0.nonblocking_synchronize(q) return diff --git a/test/execution.jl b/test/execution.jl index 44b21b9c..6eaea160 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -862,7 +862,6 @@ end # 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 @@ -888,7 +887,9 @@ end yield() end - oneL0.append_wait!(list, gate) + oneL0.execute!(oneAPI.global_stream(context(), device())) do list + oneL0.append_wait!(list, gate) + end opener = @async begin for _ in 1:10_000 yield() From a4842277d414cb09b8e65067e27b4d0de2ad2c11 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 15:34:03 +0200 Subject: [PATCH 2/2] Submit kernel launches in batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every append to an immediate command list is a submission, which costs a system call: on an Iris Xe ~9 µs of host time per kernel launch, more than the rest of the launch path combined. Kernels launched on a stream are now recorded in an in-order regular command list, which is submitted to the stream's immediate list with zeCommandListImmediateAppendCommandListsExp once it holds 32 launches, or before anything else needs the stream's work: other operations on the stream, synchronization, host accesses to memory the stream uses (through `drain`), and oneMKL calls. Recording costs ~1.5 µs per launch, and a submission also pays the driver's per-allocation residency cost once per batch instead of once per launch. The first launch after synchronizing is submitted directly, so a lone kernel starts right away. `oneAPI.flush!` submits recorded launches explicitly. Each batch keeps the memory its kernels use alive until it has executed, so freeing an array can't race with a batch that hasn't been submitted yet. `synchronize()` also submits the batches of tasks that finished without doing so. A failed submission is reported right away, and again by the stream's next synchronization. Drivers that can't submit regular lists to an immediate one, which is an experimental extension, keep submitting every launch directly, as does ONEAPI_COMMAND_BATCHING=0. --- docs/src/api/context.md | 16 +- src/compiler/execution.jl | 68 ++++++--- src/context.jl | 301 +++++++++++++++++++++++++++++++++++++- src/memory.jl | 11 +- src/oneAPI.jl | 1 + src/pool.jl | 1 + test/execution.jl | 204 +++++++++++++++++++++++++- 7 files changed, 572 insertions(+), 30 deletions(-) diff --git a/docs/src/api/context.md b/docs/src/api/context.md index dc666455..f4676c06 100644 --- a/docs/src/api/context.md +++ b/docs/src/api/context.md @@ -80,7 +80,21 @@ do not use it, so synchronizing this queue does not wait for them; use `synchron Block the host thread until all operations on the global stream for the current context and device have completed: both the immediate command list and, if one has been created, the companion command queue. `synchronize(stream::oneStream)` does the same for an -explicit stream. +explicit stream. `synchronize()` also submits and waits for the work of tasks that +finished without synchronizing. + +### Command batching + +Submitting work to the GPU is expensive, so kernels launched on a stream are recorded and +submitted in batches of up to 32. A batch is submitted when it is full, and before +anything needs the stream's work: synchronizing, copies and other operations on the +stream, accessing an array's memory from the host or from another task, and oneMKL calls. +The first kernel after synchronizing is submitted right away. + +As a result, the GPU may not start executing a kernel until one of those happens. Call +`oneAPI.flush!()` (or `oneAPI.flush!(stream)`) to submit the recorded kernels, e.g., before +doing work on the host that should overlap with them. Set `ONEAPI_COMMAND_BATCHING=0` to +submit every kernel when it is launched. ## Example Workflow diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index fd71e7cc..b5b3a180 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -154,7 +154,7 @@ end # known yet. instead, the allocations are collected, if requested, so that the launch can # take ownership of them right before submitting the kernel. struct KernelAdaptor - managed::Union{Nothing, Vector{Managed}} + refs::Union{Nothing, Vector{DataRef}} end KernelAdaptor() = KernelAdaptor(nothing) @@ -163,7 +163,7 @@ Adapt.adapt_storage(to::KernelAdaptor, p::ZePtr{T}) where {T} = reinterpret(Ptr{ # convert oneAPI host arrays to device arrays function Adapt.adapt_storage(to::KernelAdaptor, xs::oneArray{T,N}) where {T,N} - to.managed === nothing || push!(to.managed, xs.data[]) + to.refs === nothing || push!(to.refs, xs.data) ptr = reinterpret(LLVMPtr{T,AS.CrossWorkgroup}, device_pointer(xs)) oneDeviceArray{T,N,AS.CrossWorkgroup}(size(xs), ptr, xs.maxsize - xs.offset) end @@ -206,7 +206,7 @@ input object `x` as-is. Do not add methods to this function, but instead extend the underlying Adapt.jl package and register methods for the the `oneAPI.KernelAdaptor` type. """ -kernel_convert(arg, managed=nothing) = adapt(KernelAdaptor(managed), arg) +kernel_convert(arg, refs=nothing) = adapt(KernelAdaptor(refs), arg) ## abstract kernel functionality @@ -374,36 +374,55 @@ end end @inline function onecall(kernel::ZeKernel, tt, args::Tuple; groups::ZeDim=1, items::ZeDim=1, - queue::Union{oneStream, ZeCommandQueue}=global_stream(context(), device())) + queue::Union{oneStream, ZeCommandQueue}=global_stream(context(), device()), + roots=nothing) Base.@lock kernel begin # unrolled, as iterating a heterogeneous tuple would dispatch at run time ntuple(Val(length(args))) do i oneL0.arguments(kernel)[i] = args[i] return end - launch_kernel(kernel, groups, items, queue) + launch_kernel(kernel, groups, items, queue, roots) end end # kept out of `onecall`, which is compiled for every kernel signature, as none of this # depends on it @noinline function launch_kernel(kernel::ZeKernel, groups::ZeDim, items::ZeDim, - queue::Union{oneStream, ZeCommandQueue}) + queue::Union{oneStream, ZeCommandQueue}, roots) groupsize!(kernel, items) - launch!(queue, kernel, groups) + launch!(queue, kernel, groups, roots) end -@inline function launch!(s::oneStream, kernel::ZeKernel, groups::ZeDim) +# `roots` are the references to the memory the kernel uses, which a batch keeps alive +@inline function launch!(s::oneStream, kernel::ZeKernel, groups::ZeDim, roots=nothing) mkl_wait!(s) # 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. + # cleanest reachable moment, instead of at a GC-lottery-determined one, and submit + # the kernel right away. spill = oneL0.spill_mem_size(kernel) - spill > s.scratch_hwm && scratch_hedge!(s, spill) - - @lock s.lock append_launch!(s.immediate, kernel, groups) + direct = spill > s.scratch_hwm + direct && scratch_hedge!(s, spill) + + @lock s.lock begin + check_error(s) + direct |= !s.batching || s.idle || oneL0.sync_each_submission() + batch = direct ? nothing : open_batch!(s) + if batch === nothing + flush_batch!(s) + append_launch!(s.immediate, kernel, groups) + s.idle = false + else + append_launch!(batch.list, kernel, groups) + s.recorder === current_task() || register_batching!(s) + root!(batch, roots) + batch.launches += 1 + batch.launches == BATCH_SIZE && flush_batch!(s) + end + end # wait cooperatively, as `synchronize` does, or the workaround blocks the thread oneL0.sync_each_submission() && oneL0.nonblocking_synchronize(s.immediate) return @@ -411,7 +430,7 @@ 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) +@inline function launch!(queue::ZeCommandQueue, kernel::ZeKernel, groups::ZeDim, roots=nothing) spill = oneL0.spill_mem_size(kernel) spill > queue.scratch_hwm && scratch_hedge!(queue, spill) execute!(queue) do list @@ -426,7 +445,13 @@ end # only skips the drain. @noinline function scratch_hedge!(target::Union{oneStream, ZeCommandQueue}, spill::Int) if oneL0.SCRATCH_HEDGE[] - oneL0.synchronize(target) + # (without releasing memory batches kept alive: the caller holds locks) + if target isa oneStream + err = synchronize_stream!(target) + err === nothing || throw(err) + else + oneL0.synchronize(target) + end oneL0._run_reclaim_callbacks() GC.gc(false) Threads.atomic_add!(oneL0.SCRATCH_HEDGE_COUNT, 1) @@ -450,12 +475,17 @@ end @inline function launch_managed(kernel::HostKernel, source, args::Tuple; queue::Union{oneStream, ZeCommandQueue}=global_stream(context(), device()), kwargs...) - managed = Managed[] - source === nothing || kernel_convert(source, managed) - converted = map(arg -> kernel_convert(arg, managed), args) - GC.@preserve source args with_ownership(queue, managed) do - call(kernel, converted; queue, kwargs...) + refs = DataRef[] + source === nothing || kernel_convert(source, refs) + converted = map(arg -> kernel_convert(arg, refs), args) + managed = Managed[ref[] for ref in refs] + ret = GC.@preserve source args with_ownership(queue, managed) do + call(kernel, converted; queue, roots=refs, kwargs...) end + # memory that batches kept alive and turned out to be unused, also while waiting for + # another stream above. this is the first point without locks held. + release_pending!() + return ret end diff --git a/src/context.jl b/src/context.jl index 6af614b7..abbf3a40 100644 --- a/src/context.jl +++ b/src/context.jl @@ -199,9 +199,63 @@ function context!(ctx::ZeContext) task_local_storage(:ZeContext, ctx) end -# 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. +## command batching +# +# Each submission to an immediate command list is a system call, costing ~9 µs per kernel +# launch on i915. Kernel launches are therefore recorded in a regular command list, a +# batch, which is submitted to the stream's immediate list once it holds `BATCH_SIZE` +# launches, or before anything else needs the stream's work: other submissions, +# synchronization, host accesses to memory the stream uses (see `drain`), and oneMKL +# calls. The first launch after the stream was synchronized is submitted directly, so that +# a lone kernel doesn't wait for one of those. The memory the kernels of a batch use is +# kept alive until the batch has executed. + +# initialized from ONEAPI_COMMAND_BATCHING in `__init__` +const BATCHING = Ref{Bool}(true) + +const BATCH_SIZE = 32 +const MAX_BATCHES = 8 + +mutable struct Batch + const list::oneL0.ZeCommandList + const event::oneL0.ZeEvent + # copies of the references to the memory the batch's kernels use, by allocation + const roots::Vector{DataRef} + const rooted::Base.IdSet{Any} + launches::Int +end + +# whether the driver can submit regular command lists to an immediate one, which is an +# experimental extension, and executes in-order regular lists +const batching_support = Dict{Tuple{ZeContext, ZeDevice}, Bool}() +const batching_support_lock = ReentrantLock() +function batching_supported(ctx::ZeContext, dev::ZeDevice) + @lock batching_support_lock get!(batching_support, (ctx, dev)) do + try + list = oneL0.ZeCommandList(ctx, dev; flags = oneL0.ZE_COMMAND_LIST_FLAG_IN_ORDER) + oneL0.append_barrier!(list) + close(list) + immediate = oneL0.ZeImmediateCommandList(ctx, dev; + flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, + mode = oneL0.ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS) + GC.@preserve list begin + oneL0.zeCommandListImmediateAppendCommandListsExp(immediate, 1, + Ref(list.handle), + C_NULL, 0, C_NULL) + oneL0.synchronize(immediate) + end + true + catch err + # the driver rejects the flag or the function, or the loader doesn't have it + err isa Union{oneL0.ZeError, ErrorException} || rethrow() + false + end + end +end + +# The per-task submission target. `immediate` is where Julia-side work (kernels, copies, +# fills) is submitted: an in-order asynchronous immediate command list. Kernel launches are +# mostly recorded in batches first, see above. # `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) @@ -213,6 +267,21 @@ mutable struct oneStream # only append to this through `execute!` or `launch!`, which hold `lock` const immediate::oneL0.ZeImmediateCommandList const lock::ReentrantLock + + # command batching (see above), protected by `lock` + const batching::Bool + batch::Union{Nothing, Batch} # the batch being recorded + const submitted::Vector{Batch} # in submission order + const spare::Vector{Batch} + nbatches::Int + events::Union{Nothing, oneL0.ZeEventPool} + idle::Bool # submit the next launch directly + recorder::Union{Nothing, Task} # the last task to record a launch, if registered + const retired::Vector{DataRef} # roots of executed batches, still to be released + const failed::Vector{Batch} # possibly submitted, so never released + # a failed submission, reported by the next synchronization + error::Union{Nothing, Exception} + # atomic because other tasks read it when waiting for work on memory this stream owns Base.@atomic queue::Union{Nothing, ZeCommandQueue} mkl_dirty::Bool @@ -242,10 +311,195 @@ function create_stream(ctx::ZeContext, dev::ZeDevice, "reported API version $(oneL0.api_version(ctx.driver))); creation failed " * "with $(err.code)") end - s = oneStream(ctx, dev, list, ReentrantLock(), nothing, false, 0, priority) + batching = BATCHING[] && batching_supported(ctx, dev) + s = oneStream(ctx, dev, list, ReentrantLock(), + batching, nothing, Batch[], Batch[], 0, nothing, true, nothing, DataRef[], + Batch[], nothing, nothing, false, 0, priority) return register_stream!(ctx, dev, s) end +# The functions below expect the caller to hold `s.lock`. + +check_error(s::oneStream) = (s.error === nothing || throw(s.error); nothing) + +function new_batch!(s::oneStream) + if s.events === nothing + s.events = oneL0.ZeEventPool(s.ctx, MAX_BATCHES, s.dev; + flags = oneL0.ZE_EVENT_POOL_FLAG_HOST_VISIBLE) + end + s.nbatches += 1 + list = oneL0.ZeCommandList(s.ctx, s.dev; flags = oneL0.ZE_COMMAND_LIST_FLAG_IN_ORDER) + event = oneL0.ZeEvent(s.events, s.nbatches; + signal = oneL0.ZE_EVENT_SCOPE_FLAG_HOST, + wait = oneL0.ZE_EVENT_SCOPE_FLAG_HOST) + return Batch(list, event, DataRef[], Base.IdSet{Any}(), 0) +end + +# recycle the batches that have executed. their roots are released by `release_retired!`, +# not here, as that can free memory, which waits for the GPU on some stacks. +function retire!(s::oneStream) + while !isempty(s.submitted) && Base.isdone(first(s.submitted).event) + b = popfirst!(s.submitted) + isempty(b.roots) || queue_release!(s) + append!(s.retired, b.roots) + empty!(b.roots) + empty!(b.rooted) + reset(b.list) + reset(b.event) + b.launches = 0 + push!(s.spare, b) + end + return +end + +# the batch to record a launch in, or `nothing` if it should be submitted directly +function open_batch!(s::oneStream) + b = s.batch + b === nothing || return b + retire!(s) + if isempty(s.spare) + # with all batches executing, submit directly rather than waiting for one: the + # caller holds the kernel's lock, which another task may need to make progress + s.nbatches < MAX_BATCHES || return nothing + push!(s.spare, new_batch!(s)) + end + b = pop!(s.spare) + s.batch = b + return b +end + +function flush_batch!(s::oneStream) + b = s.batch + b === nothing && return + s.batch = nothing + try + close(b.list) + oneL0.zeCommandListImmediateAppendCommandListsExp(s.immediate, 1, Ref(b.list.handle), + b.event, 0, C_NULL) + catch err + # the batch may have been submitted, so keep it and the memory it uses alive + push!(s.failed, b) + s.error = err + rethrow() + end + push!(s.submitted, b) + return +end + +# keep `roots`, the references to the memory a launch uses, alive with the batch +function root!(b::Batch, roots) + roots === nothing && return + for ref in roots + managed = ref[] + managed in b.rooted && continue + push!(b.rooted, managed) + push!(b.roots, copy(ref)) + end + return +end + +""" + flush!(stream::oneStream = global_stream(context(), device())) + +Submit the kernel launches the stream has recorded but not yet submitted to the device. +Launches are submitted in batches, so call this before doing significant work on the host +that should overlap with the GPU executing those kernels. Synchronizing flushes too. +""" +function flush!(s::oneStream = global_stream(context(), device())) + @lock s.lock begin + check_error(s) + flush_batch!(s) + end + return +end + +# streams with roots to release, which happens at `release_pending!` +const pending_release = Base.IdSet{oneStream}() +const pending_release_lock = ReentrantLock() +const npending_release = Threads.Atomic{Int}(0) + +function queue_release!(s::oneStream) + @lock pending_release_lock push!(pending_release, s) + npending_release[] = 1 + return +end + +# release the memory of executed batches. the caller must not hold any lock, as freeing +# memory can wait for the GPU on some stacks. +function release_pending!() + npending_release[] == 0 && return + streams = @lock pending_release_lock begin + npending_release[] = 0 + streams = collect(pending_release) + empty!(pending_release) + streams + end + foreach(release_retired!, streams) + return +end + +function release_retired!(s::oneStream) + roots = @lock s.lock begin + isempty(s.retired) && return + roots = copy(s.retired) + empty!(s.retired) + roots + end + foreach(unsafe_free!, roots) + return +end + +# streams that recorded launches since they were last synchronized, so that `synchronize` +# can submit what tasks recorded before finishing. the stream is referenced strongly, so +# that it doesn't get destroyed with work recorded. a stream's lock is taken before this +# one, and the stream's `recorder` is the task last to record on it. +const batching_streams = Base.IdSet{oneStream}() +const batching_streams_lock = ReentrantLock() + +function register_batching!(s::oneStream) + @lock batching_streams_lock push!(batching_streams, s) + s.recorder = current_task() + return +end + +# forget a stream once nothing it recorded is still to execute. streams with batches that +# failed to submit, which may still execute, are kept around. +function maybe_unregister_batching!(s::oneStream) + if s.recorder !== nothing && s.batch === nothing && isempty(s.submitted) && + isempty(s.failed) + @lock batching_streams_lock delete!(batching_streams, s) + s.recorder = nothing + end + return +end + +# wait for the streams on which tasks that have finished were the last to record. other +# tasks' streams are left alone: their tasks synchronize them. +function adopt_orphans!(ctx::ZeContext, dev::ZeDevice) + streams = @lock batching_streams_lock collect(batching_streams) + for s in streams + (s.ctx == ctx && s.dev == dev) || continue + orphaned = @lock s.lock begin + task = s.recorder + task !== nothing && istaskdone(task) + end + orphaned || continue + err = nothing + try + drain(s) + catch e + err = e + end + @lock s.lock begin + # reported here, as the task that would have isn't around anymore + err === nothing || (s.error = nothing) + maybe_unregister_batching!(s) + end + err === nothing || throw(err) + end + return +end + """ global_stream(ctx::ZeContext, dev::ZeDevice) -> oneStream @@ -406,18 +660,43 @@ println("GPU work completed") See also: [`global_stream`](@ref), [`context`](@ref), [`device`](@ref) """ function oneL0.synchronize(s::oneStream; blocking::Bool=false) + err = synchronize_stream!(s; blocking) + release_pending!() + err === nothing || throw(err) + return +end + +# synchronize a stream, returning the error of a failed submission +function synchronize_stream!(s::oneStream; blocking::Bool=false) sync = blocking ? oneL0.synchronize : oneL0.nonblocking_synchronize + # a failed submission is recorded in `s.error`, and reported below + @lock s.lock try + flush_batch!(s) + catch + end sync(s.immediate) q = Base.@atomic s.queue if q !== nothing sync(q) s.mkl_dirty = false end - return + err = @lock s.lock begin + retire!(s) + s.idle = true + maybe_unregister_batching!(s) + err = s.error + s.error = nothing + err + end + return err end function oneL0.synchronize(; blocking::Bool=false) - oneL0.synchronize(global_stream(context(), device()); blocking) + ctx = context() + dev = device() + oneL0.synchronize(global_stream(ctx, dev); blocking) + adopt_orphans!(ctx, dev) + release_pending!() end # Julia → MKL ordering: everything Julia appended to the task's immediate list must be @@ -427,6 +706,10 @@ end # 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())) + @lock s.lock begin + check_error(s) + flush_batch!(s) + end oneL0.synchronize(s.immediate) s.mkl_dirty = true return @@ -456,7 +739,11 @@ outstanding oneMKL work. Appends are submitted to the device as they happen. """ @inline function oneL0.execute!(f::Base.Callable, s::oneStream) mkl_wait!(s) - @lock s.lock oneL0.execute!(f, s.immediate) + @lock s.lock begin + check_error(s) + flush_batch!(s) + oneL0.execute!(f, s.immediate) + end end # re-export and augment parts of oneL0 to make driver and device selection easier diff --git a/src/memory.jl b/src/memory.jl index 72a0a159..4afeabdf 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -85,9 +85,17 @@ Base.sizeof(managed::Managed) = sizeof(managed.mem) # wait for all work on a stream, possibly from another task. unlike `synchronize`, this # leaves `mkl_dirty` alone, which only the stream's own task may modify. function drain(s::oneStream) + @lock s.lock begin + check_error(s) + flush_batch!(s) + end oneL0.nonblocking_synchronize(s.immediate) q = Base.@atomic s.queue q === nothing || oneL0.nonblocking_synchronize(q) + @lock s.lock begin + retire!(s) + s.idle = true + end return end drain(q::ZeCommandQueue) = oneL0.nonblocking_synchronize(q) @@ -125,7 +133,7 @@ end function with_managed_locks(f, managed::Vector{Managed}) locked = sort!(unique!(managed); by=objectid) nlocked = 0 - try + ret = try for m in locked lock(m.lock) nlocked += 1 @@ -136,6 +144,7 @@ function with_managed_locks(f, managed::Vector{Managed}) unlock(locked[i].lock) end end + return ret end # submit an operation using the memory of `managed` to `owner`, as a single transaction diff --git a/src/oneAPI.jl b/src/oneAPI.jl index 15c13d8f..da72d9f9 100644 --- a/src/oneAPI.jl +++ b/src/oneAPI.jl @@ -139,6 +139,7 @@ function __init__() # slabs are process-local empty!(pools) POOL[] = oneL0.parse_env_bool("ONEAPI_MEMORY_POOL", true) + BATCHING[] = oneL0.parse_env_bool("ONEAPI_COMMAND_BATCHING", true) @static if Sys.iswindows() @warn """oneAPI.jl support for native Windows is experimental and incomplete. diff --git a/src/pool.jl b/src/pool.jl index 1e125997..01773922 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -44,6 +44,7 @@ function _maybe_gc(dev, bytes) allocated = _allocated_bytes[] allocated <= 0 && return total_mem = _memory_limit(dev) + allocated + bytes > total_mem * 0.4 && release_pending!() if allocated + bytes > total_mem * 0.8 # Flush deferred resource releases (e.g., MKL sparse handles) from previous GC # cycles first — these are safe to release now because they were deferred earlier. diff --git a/test/execution.jl b/test/execution.jl index 6eaea160..28a674bd 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -1,4 +1,5 @@ import Adapt +import LinearAlgebra using StaticArrays @@ -847,6 +848,202 @@ function fill_kernel(a, x) return end +# a step of a recurrence that doesn't commute, so its result depends on the order of launches +function recurrence_kernel(a, i) + @inbounds a[1] = 3 * a[1] + i + return +end +function recurrence(n) + x = 0 + for i in 1:n + x = 3 * x + i + end + return x +end + +function ptr_increment_kernel(ptr) + unsafe_store!(ptr, unsafe_load(ptr) + 1) + return +end + +@testset "command batching" begin + n = 100 # spans several batches + batching = oneAPI.global_stream(context(), device()).batching && + !oneL0.sync_each_submission() + + # dependent launches execute in order, across batches + a = oneArray(Int64[0]) + for i in 1:n + @oneapi recurrence_kernel(a, i) + end + @test Array(a)[1] == recurrence(n) + + # interleaved with operations that are submitted directly + a = oneArray(Int64[0]) + for i in 1:n + @oneapi recurrence_kernel(a, i) + i % 10 == 0 && copyto!(a, a) + i % 25 == 0 && copyto!(oneArray{Int64}(undef, 1), [i]) + end + @test Array(a)[1] == recurrence(n) + + # host accesses wait for launches that haven't been submitted yet + a = oneArray{Int64, 1, oneL0.SharedBuffer}([0]) + for i in 1:n + @oneapi recurrence_kernel(a, i) + end + @test a[1] == recurrence(n) + + # also when they were recorded by another task + a = oneArray(Int64[0]) + wait(@async begin + for i in 1:n + @oneapi recurrence_kernel(a, i) + end + end) + @test Array(a)[1] == recurrence(n) + + # synchronizing submits what tasks recorded before finishing, also when the memory it + # uses isn't tracked, and when another task records on the same stream later + a = oneArray(Int64[0]) + ptr = reinterpret(Core.LLVMPtr{Int64, oneAPI.AS.CrossWorkgroup}, pointer(a)) + stream = fetch(@async begin + for _ in 1:3 + @oneapi ptr_increment_kernel(ptr) + end + oneAPI.global_stream(context(), device()) + end) + synchronize() + @test Array(a)[1] == 3 + wait(@async begin + for _ in 1:3 + @oneapi queue=stream ptr_increment_kernel(ptr) + end + end) + synchronize() + @test Array(a)[1] == 6 + + # including on the stream of a task that is still running + streams = Channel{Any}(1) + done = Channel{Nothing}(1) + owner = @async begin + # with a batch recorded, which the other task records into + for _ in 1:2 + @oneapi ptr_increment_kernel(ptr) + end + put!(streams, oneAPI.global_stream(context(), device())) + take!(done) + end + stream = take!(streams) + batching && @test stream.recorder === owner + wait(@async begin + for _ in 1:3 + @oneapi queue=stream ptr_increment_kernel(ptr) + end + end) + synchronize() + @test Array(a)[1] == 11 + put!(done, nothing) + wait(owner) + + # memory used by recorded launches stays alive until they have executed + a = oneArray(Int64[0]) + GC.enable(false) + try + @oneapi recurrence_kernel(a, 1) # submitted directly + bytes = 1 << 23 # not sub-allocated + before = oneAPI._allocated_bytes[] + b = oneArray{Int64}(undef, bytes ÷ sizeof(Int64)) + @oneapi items=256 groups=4 fill_kernel(b, 1) + oneAPI.unsafe_free!(b) + batching && @test oneAPI._allocated_bytes[] == before + bytes + synchronize() + @test oneAPI._allocated_bytes[] == before + finally + GC.enable(true) + end + + # recorded launches are submitted on request + a = oneArray{Int64, 1, oneL0.SharedBuffer}([0]) + synchronize() + for i in 1:3 + @oneapi recurrence_kernel(a, i) + end + oneAPI.flush!() + # wait for the stream without waiting for the work on `a`, which would submit it, and + # read the memory directly + oneL0.synchronize(oneAPI.global_stream(context(), device()).immediate) + @test unsafe_load(reinterpret(Ptr{Int64}, oneAPI.device_pointer(a))) == recurrence(3) + + # with all batches executing, launches are submitted directly instead of waiting for + # one, which may not finish before this task does something (here, open the gate) + if batching + pool = oneL0.ZeEventPool(context(), 1, device(); flags=oneL0.ZE_EVENT_POOL_FLAG_HOST_VISIBLE) + gate = pool[1] + a = oneArray(Int64[0]) + # (only a few more: many submissions waiting for the gate make the driver stall, and + # eventually give up on the device) + m = oneAPI.MAX_BATCHES * oneAPI.BATCH_SIZE + 10 + GC.enable(false) # freeing memory may wait for the stream on some stacks + try + oneL0.execute!(oneAPI.global_stream(context(), device())) do list + oneL0.append_wait!(list, gate) + end + for i in 1:m + @oneapi recurrence_kernel(a, i) + end + oneL0.signal(gate) + finally + GC.enable(true) + end + @test Array(a)[1] == recurrence(m) + end + + # work recorded before using oneMKL is executed first, and work after it waits for it + x = oneArray(ones(Float32, 1024)) + y = oneArray(zeros(Float32, 1024)) + for _ in 1:40 + x .+= 1 + end + LinearAlgebra.axpy!(2f0, x, y) + for _ in 1:40 + y .+= 1 + end + @test Array(y) == fill(2f0 * 41 + 40, 1024) + + # a failed submission is reported until the stream is synchronized + @test fetch(@async begin + a = oneArray(Int64[0]) + stream = oneAPI.global_stream(context(), device()) + err = ErrorException("submission failed") + stream.error = err + caught(f) = try f(); nothing catch e; e end + results = [caught(() -> @oneapi recurrence_kernel(a, 1)), + caught(() -> copyto!(a, [1])), + caught(() -> oneAPI.flush!(stream)), + caught(() -> oneAPI.drain(stream)), + caught(() -> synchronize())] + # also when crossing a new spill tier, which synchronizes first + stream.error = err + stream.scratch_hwm = -1 + push!(results, caught(() -> @oneapi recurrence_kernel(a, 1))) + stream.error = nothing + all(r -> r === err, results) && caught(() -> synchronize()) === nothing + end) + + # and reported once by whoever synchronizes when the task that recorded didn't + err = ErrorException("submission failed") + a = oneArray(Int64[0]) + wait(@async begin + for i in 1:2 + @oneapi recurrence_kernel(a, i) + end + oneAPI.global_stream(context(), device()).error = err + end) + @test try synchronize(); nothing catch e; e end === (batching ? err : nothing) + @test try synchronize(); nothing catch e; e end === nothing +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 @@ -887,8 +1084,11 @@ end yield() end - oneL0.execute!(oneAPI.global_stream(context(), device())) do list - oneL0.append_wait!(list, gate) + # (without synchronizing after the submission, which would wait for the gate) + oneL0.sync_each_submission(false) do + oneL0.execute!(oneAPI.global_stream(context(), device())) do list + oneL0.append_wait!(list, gate) + end end opener = @async begin for _ in 1:10_000