Guard NEO's scratch allocation with a spill-triggered drain - #609
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #609 +/- ##
==========================================
+ Coverage 80.48% 80.97% +0.48%
==========================================
Files 50 50
Lines 3490 3505 +15
==========================================
+ Hits 2809 2838 +29
+ Misses 681 667 -14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
NEO allocates a queue's scratch buffer at the first submission of a kernel whose spill exceeds what is already allocated, and that allocation has no error path: allocateGraphicsMemoryWithProperties is called without a null check, and programSurfaceState aborts the process (UNRECOVERABLE_IF) when it failed. Whether it fails depends on how much dead-but-unfinalized driver memory happens to be live at that one instant — a GC lottery. Remove the lottery: track a per-queue high-water mark of kernel spill sizes, and before the first submission that crosses it, retire in-flight work, flush deferred releases, and run finalizers, so the fatal allocation happens at the cleanest reachable moment. Fires once per (queue, scratch tier); costs one cached-Int compare on the non-spilling fast path. Opt out with ONEAPI_SCRATCH_HEDGE=0. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WSFxSBXtckG3BVAT12wYkf
michel2323
force-pushed
the
scratch-hedge
branch
from
August 20, 2026 14:35
1ee7863 to
2d2a1ca
Compare
michel2323
marked this pull request as ready for review
August 20, 2026 15:03
Contributor
|
Your PR requires formatting changes to meet the project's style guidelines. Click here to view the suggested changes.diff --git a/test/execution.jl b/test/execution.jl
index 8f11ae3..303c8c7 100644
--- a/test/execution.jl
+++ b/test/execution.jl
@@ -635,7 +635,7 @@ function hedge_spill_kernel(out, a, ::Val{N}, ::Val{R}) where {N, R}
i = get_global_id()
acc = hedge_rounds(ntuple(j -> a[i] + Float32(j), Val(N)), Val(R))
s = 0.0f0
- @inbounds for j = 1:N
+ @inbounds for j in 1:N
s += acc[j]
end
@inbounds out[i] = s
@@ -647,7 +647,7 @@ end
out = oneAPI.zeros(Float32, 256)
synchronize()
- k = @oneapi launch=false hedge_spill_kernel(out, a, Val(256), Val(2))
+ k = @oneapi launch = false hedge_spill_kernel(out, a, Val(256), Val(2))
spill = oneL0.spill_mem_size(k.fun)
if spill == 0
@@ -656,34 +656,38 @@ end
else
# 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
- c0 = oneL0.SCRATCH_HEDGE_COUNT[]
- @oneapi items=64 groups=4 hedge_spill_kernel(out, a, Val(256), Val(2))
- c1 = oneL0.SCRATCH_HEDGE_COUNT[]
- # same spill tier: must not fire again
- @oneapi items=64 groups=4 hedge_spill_kernel(out, a, Val(256), Val(2))
- c2 = oneL0.SCRATCH_HEDGE_COUNT[]
- # a storm of no-spill kernels must not fire the hedge either
- for _ in 1:32
- @oneapi dummy()
+ observed = fetch(
+ @async begin
+ q = global_queue(context(), device())
+ hwm0 = q.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[]
+ # same spill tier: must not fire again
+ @oneapi items = 64 groups = 4 hedge_spill_kernel(out, a, Val(256), Val(2))
+ c2 = oneL0.SCRATCH_HEDGE_COUNT[]
+ # a storm of no-spill kernels must not fire the hedge either
+ for _ in 1:32
+ @oneapi dummy()
+ end
+ c3 = oneL0.SCRATCH_HEDGE_COUNT[]
+ (hwm0, q.scratch_hwm, c1 - c0, c2 - c1, c3 - c2)
end
- c3 = oneL0.SCRATCH_HEDGE_COUNT[]
- (hwm0, q.scratch_hwm, c1 - c0, c2 - c1, c3 - c2)
- end)
+ )
@test observed == (0, spill, 1, 0, 0)
# with the hedge disabled the high-water mark is still maintained
old = oneL0.SCRATCH_HEDGE[]
oneL0.SCRATCH_HEDGE[] = false
try
- observed = fetch(@async begin
- q = global_queue(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)
- end)
+ observed = fetch(
+ @async begin
+ q = global_queue(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)
+ end
+ )
@test observed == (0, spill)
finally
oneL0.SCRATCH_HEDGE[] = old |
michel2323
referenced
this pull request
Aug 21, 2026
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
NEO allocates a queue's scratch buffer at the first submission of a kernel whose spill exceeds what is already allocated, and that allocation has no error path:
prepareScratchAllocation()callsallocateGraphicsMemoryWithProperties()without a null check, andprogramSurfaceState()aborts the process (UNRECOVERABLE_IF,scratch_space_controller_xehp_and_later.cpp:78, byte-identical from the 25.18 LTS driver through current master) when it failed. Whether it fails depends on how much dead-but-unfinalized driver memory happens to be live at that one instant — a GC lottery. We have logged repeated aborts of exactly this shape on Aurora (PVC, LTS NEO 25.18) under launch storms from real workloads (ExaModels), where the process dies insidezeCommandQueueExecuteCommandListswith no Julia-side recourse.This PR removes the lottery instead of the (driver-side) bug:
ZeKernelcaches itsspillMemSize(seeded byproperties(), one Int load on the launch path)._run_reclaim_callbacks), and runsGC.gc(false)— so the one allocation the driver cannot recover from happens at the cleanest reachable moment.The hedge fires once per (queue, scratch tier): one synchronize per workload in practice, and a cached-Int compare on the non-spilling fast path. Opt out with
ONEAPI_SCRATCH_HEDGE=0(the high-water mark is maintained regardless).Tested with a deliberately spilling kernel (256 live accumulators): the hedge fires exactly once, never for no-spill kernels, and the knob-off path still tracks the mark. Suite green on Aurora LTS (PVC) and the kernel spills on both PVC and DG2.
This is deliberately independent of the immediate-command-list rework (#610): it guards a different resource (scratch) against a different garbage source (USM churn), and remains useful on current drivers — the missing null check is in NEO's shared code, not LTS-specific.