Guard NEO's scratch allocation with a spill-triggered drain

lib/level-zero/cmdqueue.jl

@@ -15,6 +15,10 @@ mutable struct ZeCommandQueue
     device::ZeDevice
     ordinal::Int
 
+    # high-water mark of per-thread spill (bytes) among kernels submitted to this queue,
+    # maintained by the scratch hedge (see `scratch_hedge!` in src/compiler/execution.jl)
+    scratch_hwm::Int
+
     function ZeCommandQueue(ctx::ZeContext, dev::ZeDevice, ordinal=1, index=1;
                             flags=0,
                             mode::ze_command_queue_mode_t=ZE_COMMAND_QUEUE_MODE_DEFAULT,
@@ -24,7 +28,7 @@ mutable struct ZeCommandQueue
         ))
         handle_ref = Ref{ze_command_queue_handle_t}()
         zeCommandQueueCreate(ctx, dev, desc_ref, handle_ref)
-        obj = new(handle_ref[], ctx, dev, ordinal)
+        obj = new(handle_ref[], ctx, dev, ordinal, 0)
         finalizer(obj) do obj
             if LTS[]
                 # the queue may still have work in flight (nothing requires a task to

lib/level-zero/module.jl

@@ -80,13 +80,17 @@ mutable struct ZeKernel
     mod::ZeModule
     handle::ze_kernel_handle_t
 
+    # cached spillMemSize in bytes/thread, seeded by `properties`; -1 while unqueried.
+    # Read on every launch by the scratch hedge, so it must not cost an API call.
+    spill::Int
+
     function ZeKernel(mod, name)
         GC.@preserve name begin
             desc_ref = Ref(ze_kernel_desc_t(; pKernelName=pointer(name)))
             handle_ref = Ref{ze_kernel_handle_t}()
             zeKernelCreate(mod, desc_ref, handle_ref)
         end
-        obj = new(mod, handle_ref[])
+        obj = new(mod, handle_ref[], -1)
 
         finalizer(obj) do obj
             zeKernelDestroy(obj)
@@ -258,6 +262,7 @@ function properties(kernel::ZeKernel)
     end
 
     props = props_ref[]
+    kernel.spill = Int(props.spillMemSize)
     return (
         numKernelArgs=Int(props.numKernelArgs),
         requiredGroupSize=ZeDim3(props.requiredGroupSizeX,
@@ -278,6 +283,13 @@ function properties(kernel::ZeKernel)
     )
 end
 
+# Cached access to a kernel's spill (scratch) size for the launch path: one Int load
+# after the first query. `properties` seeds the cache.
+function spill_mem_size(kernel::ZeKernel)
+    s = kernel.spill
+    return s >= 0 ? s : Int(properties(kernel).spillMemSize)
+end
+
 
 ## execution
 

lib/level-zero/oneL0.jl

@@ -234,6 +234,7 @@ function __init__()
 
     validation_layer[] = parse(Bool, get(ENV, "ZE_ENABLE_VALIDATION_LAYER", "false"))
     parameter_validation[] = parse(Bool, get(ENV, "ZE_ENABLE_PARAMETER_VALIDATION", "false"))
+    SCRATCH_HEDGE[] = parse_env_bool("ONEAPI_SCRATCH_HEDGE", true)
     return sync_each_submission!(parse_env_bool("ONEAPI_SYNC_EACH_SUBMISSION", false))
 end
 

lib/level-zero/utils.jl

@@ -1,5 +1,14 @@
 isdebug(group) = Base.CoreLogging.current_logger_for_env(Base.CoreLogging.Debug, group, oneL0) !== nothing
 
+# Opt-out knob for the scratch hedge (`scratch_hedge!` in src/compiler/execution.jl):
+# NEO's scratch-buffer allocation has no error path — on failure the process aborts —
+# so the hedge drains garbage right before the first submission that triggers it.
+# Initialized from ONEAPI_SCRATCH_HEDGE in `__init__`; default on.
+const SCRATCH_HEDGE = Ref{Bool}(true)
+
+# number of times the hedge actually drained; observable for tests
+const SCRATCH_HEDGE_COUNT = Threads.Atomic{Int}(0)
+
 # Registered callbacks invoked during memory reclamation (e.g., flushing deferred MKL
 # sparse handle releases).  Extensions like oneMKL can register cleanup functions here
 # so they run when Level Zero reports OOM or when proactive GC fires.

src/compiler/execution.jl

@@ -328,11 +328,35 @@ const _kernel_instances = Dict{UInt, Any}()
     end
 
     groupsize!(kernel, items)
+
+    # NEO allocates a queue'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)
+
     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
+# 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)
+    if oneL0.SCRATCH_HEDGE[]
+        oneL0.synchronize(queue)
+        oneL0._run_reclaim_callbacks()
+        GC.gc(false)
+        Threads.atomic_add!(oneL0.SCRATCH_HEDGE_COUNT, 1)
+    end
+    queue.scratch_hwm = spill
+    return
+end
+
 function (kernel::HostKernel)(args...; kwargs...)
     call(kernel, map(kernel_convert, args)...; kwargs...)
 end

test/execution.jl

@@ -621,3 +621,72 @@ end
 end
 
 ############################################################################################
+
+# A chain of N live accumulators forces the register allocator to spill once N exceeds
+# the register file (256 spills on both PVC and DG2); the reversed index in the update
+# keeps every element live across the chain, and everything is unrolled through Val so
+# no value is boxed. Top-level definitions: inside the testset the self-recursion of
+# `hedge_rounds` becomes a boxed closure capture, which is not a bitstype.
+hedge_step(acc::NTuple{N, T}) where {N, T} =
+    ntuple(j -> muladd(acc[j], T(1.0000001), acc[N - j + 1]), Val(N))
+hedge_rounds(acc, ::Val{0}) = acc
+hedge_rounds(acc, ::Val{R}) where {R} = hedge_rounds(hedge_step(acc), Val(R - 1))
+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
+        s += acc[j]
+    end
+    @inbounds out[i] = s
+    return
+end
+
+@testset "scratch hedge" begin
+    a = oneAPI.ones(Float32, 256)
+    out = oneAPI.zeros(Float32, 256)
+    synchronize()
+
+    k = @oneapi launch=false hedge_spill_kernel(out, a, Val(256), Val(2))
+    spill = oneL0.spill_mem_size(k.fun)
+
+    if spill == 0
+        # the hedge only acts on spilling kernels; nothing to observe on this device
+        @test_skip "kernel did not spill on this device/compiler"
+    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()
+            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)
+            @test observed == (0, spill)
+        finally
+            oneL0.SCRATCH_HEDGE[] = old
+        end
+    end
+end