diff --git a/lib/level-zero/memory.jl b/lib/level-zero/memory.jl index 66cb3fa9..c74ba31e 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 8de49eb6..020ad9a1 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 b5b3a180..d33707cd 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 01773922..eb7f92ab 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 2e52da33..68ac8854 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 76048d11..c9de9f79 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