diff --git a/lib/level-zero/cmdqueue.jl b/lib/level-zero/cmdqueue.jl index 75ddb63b..d32123a7 100644 --- a/lib/level-zero/cmdqueue.jl +++ b/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 diff --git a/lib/level-zero/module.jl b/lib/level-zero/module.jl index 0c9ddbb1..6abdd836 100644 --- a/lib/level-zero/module.jl +++ b/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 diff --git a/lib/level-zero/oneL0.jl b/lib/level-zero/oneL0.jl index d24ba412..82e33eec 100644 --- a/lib/level-zero/oneL0.jl +++ b/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 diff --git a/lib/level-zero/utils.jl b/lib/level-zero/utils.jl index bf0153b8..8de49eb6 100644 --- a/lib/level-zero/utils.jl +++ b/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. diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 327cf8c7..1f04465e 100644 --- a/src/compiler/execution.jl +++ b/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 diff --git a/test/execution.jl b/test/execution.jl index 1e2e1797..8f11ae3e 100644 --- a/test/execution.jl +++ b/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