Don't wait for the GPU when freeing memory

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
 

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()

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

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

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))

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