From f6be2c5e4e2751d7c46e139cf79dfad04419f0e7 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 19 Oct 2026 12:15:25 +0200 Subject: [PATCH] Don't wait for the GPU when freeing memory Kernels access memory indirectly, so a blocking free waits for all work submitted so far. Memory is mostly freed from finalizers, so that stalled whichever thread collected garbage for as long as unrelated kernels ran (~860 ms with a 1 s kernel), and deadlocked when that work waited for a task on the thread. On the rolling stack, free with the DEFER_FREE policy, which has the driver release the memory once the GPU is done with it. NEO only checks that when memory is freed again, so reclaiming memory now frees a small allocation after collecting garbage. That doesn't wait for the GPU either: reclaiming can happen inside any driver call, with locks held. The LTS stack doesn't honor the free policies and keeps synchronizing. Part of #717. --- lib/level-zero/memory.jl | 30 ++++++++++++++++++++++++++++++ lib/level-zero/utils.jl | 2 ++ src/compiler/execution.jl | 1 + src/pool.jl | 32 ++++++++++++++++++-------------- test/execution.jl | 9 +++++++++ test/level-zero.jl | 25 +++++++++++++++++++++++++ 6 files changed, 85 insertions(+), 14 deletions(-) diff --git a/lib/level-zero/memory.jl b/lib/level-zero/memory.jl index 66cb3fa9c..c74ba31ed 100644 --- a/lib/level-zero/memory.jl +++ b/lib/level-zero/memory.jl @@ -25,9 +25,39 @@ function free(buf::AbstractBuffer; policy=nothing) else desc_ref = Ref(ze_memory_free_ext_desc_t(; freePolicy=policy)) zeMemFreeExt(ctx, desc_ref, buf) + if policy == ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_DEFER_FREE + @atomic deferred_frees.context = WeakRef(ctx) + end end end +# Memory freed with the DEFER_FREE policy is released by the driver once the GPU is done +# with it, but NEO only checks that when memory is freed again, so remember a context in +# which that happened. +mutable struct DeferredFrees + @atomic context::WeakRef +end +const deferred_frees = DeferredFrees(WeakRef()) + +""" + sweep_deferred_frees() + +Have the driver release memory that was freed with the `DEFER_FREE` policy and that the +GPU is done with, by freeing memory again. Best effort: this does not wait for the GPU. +""" +function sweep_deferred_frees() + ctx = (@atomic deferred_frees.context).value + ctx === nothing && return + # unchecked, as running out of memory here should not reclaim memory again + ptr_ref = Ref{Ptr{Cvoid}}() + res = unchecked_zeMemAllocHost(ctx, Ref(ze_host_mem_alloc_desc_t()), 64, 1, ptr_ref) + res == RESULT_SUCCESS || return + desc_ref = Ref(ze_memory_free_ext_desc_t(; + freePolicy=ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_DEFER_FREE)) + unchecked_zeMemFreeExt(ctx, desc_ref, ptr_ref[]) + return +end + ## device buffer diff --git a/lib/level-zero/utils.jl b/lib/level-zero/utils.jl index 8de49eb69..020ad9a1e 100644 --- a/lib/level-zero/utils.jl +++ b/lib/level-zero/utils.jl @@ -48,6 +48,8 @@ function retry_reclaim(f, isfailed) else break end + # finalizers may only have deferred freeing memory + sweep_deferred_frees() phase += 1 ret = f() diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index b5b3a180b..d33707cdb 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -454,6 +454,7 @@ end end oneL0._run_reclaim_callbacks() GC.gc(false) + oneL0.sweep_deferred_frees() Threads.atomic_add!(oneL0.SCRATCH_HEDGE_COUNT, 1) end target.scratch_hwm = spill diff --git a/src/pool.jl b/src/pool.jl index 017739221..eb7f92ab5 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -59,6 +59,8 @@ function _maybe_gc(dev, bytes) end # the memory of collected arrays only becomes available once their slabs are empty trim_pools!() + # and the driver may only have deferred freeing it (see `release`) + oneL0.sweep_deferred_frees() return end @@ -288,7 +290,7 @@ function trim_pools!() end empty end - # waits for outstanding work on the memory (see `release`) + # on the LTS stack, this waits for outstanding work on the memory (see `release`) foreach(slab -> release(slab.buf), empty) return end @@ -313,22 +315,24 @@ function release(buf::oneL0.AbstractBuffer) # evict(ctx, dev, buf) #end - # NEO (at least the 25.18 LTS release) does not honor the BLOCKING_FREE/DEFER_FREE - # policies of zeMemFreeExt: it advertises ZE_extension_memory_free_policies but - # unmaps the allocation immediately, even with work in flight that references it. - # That turns a GC-driven free of a dead array whose last kernel/copy hasn't retired - # into a GPU pagefault, which gets the kernel context banned and makes every later - # submission fail with ZE_RESULT_ERROR_UNKNOWN. Synchronize the queues that could - # reference this buffer before freeing. (No-op on the rolling stack, which honors - # BLOCKING_FREE.) if oneL0.LTS[] + # NEO (at least the 25.18 LTS release) does not honor the BLOCKING_FREE/DEFER_FREE + # policies of zeMemFreeExt: it advertises ZE_extension_memory_free_policies but + # unmaps the allocation immediately, even with work in flight that references it. + # That turns a GC-driven free of a dead array whose last kernel/copy hasn't retired + # into a GPU pagefault, which gets the kernel context banned and makes every later + # submission fail with ZE_RESULT_ERROR_UNKNOWN. Synchronize the queues that could + # reference this buffer before freeing. synchronize_all_streams(oneL0.context(buf), oneL0.device(buf)) + free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_BLOCKING_FREE) + else + # with kernels accessing memory indirectly, a blocking free waits for all work + # submitted so far, which stalls the thread and deadlocks when that work waits for + # a task on it (this is usually called from a finalizer). have the driver free the + # memory once that work has completed instead. NEO doesn't report the free policies + # it supports, so this relies on the rolling stack implementing DEFER_FREE. + free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_DEFER_FREE) end - free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_BLOCKING_FREE) - - # TODO: queue-ordered free from non-finalizer tasks once we have - # `zeMemFreeAsync(ptr, queue)` - return end diff --git a/test/execution.jl b/test/execution.jl index 2e52da333..68ac8854b 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -1138,6 +1138,15 @@ end end @test Array(a) == fill(Int32(1), 64) + # freeing memory, e.g. from a finalizer, doesn't wait for the GPU + if !oneL0.LTS[] + b = oneArray{Int32}(undef, 2^20) # too large to be sub-allocated + @test gated() do + oneAPI.unsafe_free!(b) + synchronize() + end + end + @test gated() do oneL0.sync_each_submission(false) do @oneapi items=64 fill_kernel(a, Int32(2)) diff --git a/test/level-zero.jl b/test/level-zero.jl index 76048d119..c9de9f791 100644 --- a/test/level-zero.jl +++ b/test/level-zero.jl @@ -292,6 +292,31 @@ ptr = convert(Ptr{Cvoid}, buf) @test lookup_alloc(ctx, ptr) isa typeof(buf) free(buf) +# memory whose freeing was deferred is released once the GPU is done with it, and memory +# is freed again. this submits work that waits on an event signaled afterwards, which the +# ONEAPI_SYNC_EACH_SUBMISSION workaround would deadlock. +if !oneL0.LTS[] + queue = ZeCommandQueue(ctx, dev) + event = ZeEventPool(ctx, 1)[1] + pattern = host_alloc(ctx, 1) + buf = device_alloc(ctx, dev, 1024) + oneL0.sync_each_submission(false) do + execute!(queue) do list + append_wait!(list, event) + append_fill!(list, pointer(buf), convert(Ptr{Cvoid}, pattern), 1, 1024) + end + end + free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_DEFER_FREE) + oneL0.sweep_deferred_frees() # doesn't wait for the GPU + @test properties(buf).type == oneL0.ZE_MEMORY_TYPE_DEVICE + signal(event) + synchronize(queue) + @test properties(buf).type == oneL0.ZE_MEMORY_TYPE_DEVICE + oneL0.sweep_deferred_frees() + @test properties(buf).type == oneL0.ZE_MEMORY_TYPE_UNKNOWN + free(pattern) +end + end