Harden command-queue lifetime to prevent banned-context faults on free

On the LTS NEO stack (25.18) freeing a buffer does not drain queues that
still have work in flight referencing it: the in-flight kernel then
faults and the context is banned, surfacing later as a
ZE_RESULT_ERROR_UNKNOWN at an unrelated op. global_queue is task-local,
so a test file's task can also die with work still queued.

  - Register every global queue in a per-(context,device) registry that
    holds the queue strongly, keyed by a weak reference to the owning
    task (src/context.jl). A WeakRef to the queue would be cleared in the
    same GC cycle that queues its finalizer, hiding it from release
    exactly when its in-flight work still references buffers being freed.
  - Before any BLOCKING_FREE, synchronize all queues that could
    reference the buffer (src/pool.jl, synchronize_all_queues).
  - Synchronize outside the registry lock with finalizers disabled, skip
    already-finalized queues, and retire queues of dead tasks.
  - In the queue finalizer, drain (unchecked, since a banned context
    returns an error) then destroy, and null the handle so a concurrent
    synchronize_all_queues skips it (lib/level-zero/cmdqueue.jl).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: michel2323

lib/level-zero/cmdqueue.jl

@@ -20,7 +20,17 @@ mutable struct ZeCommandQueue
         zeCommandQueueCreate(ctx, dev, desc_ref, handle_ref)
         obj = new(handle_ref[], ctx, dev, ordinal)
         finalizer(obj) do obj
+            # the queue may still have work in flight (nothing requires a task to
+            # synchronize before dying), and zeCommandQueueDestroy does not wait for
+            # it: on the LTS NEO stack the still-running work then faults as soon as
+            # a referenced allocation is freed, getting the context banned. drain the
+            # queue first; unchecked, as sync on a banned context returns an error.
+            unchecked_zeCommandQueueSynchronize(obj, typemax(UInt64))
             zeCommandQueueDestroy(obj)
+            # mark the queue as destroyed: it can still be weakly reachable (e.g. from
+            # the queue registry used by `synchronize_all_queues`), and synchronizing a
+            # destroyed handle crashes in the driver.
+            obj.handle = ze_command_queue_handle_t(C_NULL)
         end
         obj
     end

src/context.jl

@@ -224,8 +224,84 @@ function global_queue(ctx::ZeContext, dev::ZeDevice)
     # NOTE: dev purposefully does not default to context() or device() to stress that
     #       objects should track ownership, and not rely on implicit global state.
     get!(task_local_storage(), (:ZeCommandQueue, ctx, dev)) do
-        ZeCommandQueue(ctx, dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER)
+        queue = ZeCommandQueue(ctx, dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER)
+        # disable finalizers while mutating the registry: a GC-driven finalizer on this
+        # task could call back into `synchronize_all_queues` (the lock is reentrant) and
+        # observe/mutate the registry mid-update.
+        GC.enable_finalizers(false)
+        try
+            @lock queue_registry_lock begin
+                push!(get!(Vector{Tuple{WeakRef,ZeCommandQueue}}, queue_registry, (ctx, dev)),
+                      (WeakRef(current_task()), queue))
+            end
+        finally
+            GC.enable_finalizers(true)
+        end
+        queue
+    end
+end
+
+# Registry of all queues created through `global_queue`, across tasks. Buffers can be
+# freed from any task (GC finalizers), so `release` needs to be able to find the queues
+# that may still have work in flight referencing the buffer; queues themselves are
+# cached task-locally and would otherwise be unreachable from the finalizing task.
+#
+# Entries reference the queue *strongly*: the GC clears WeakRefs to a dead queue in the
+# same cycle that queues its finalizer, i.e., before the finalizer runs, so a WeakRef
+# would hide the queue from `release` exactly when its in-flight work still references
+# buffers about to be freed. The owning task is tracked weakly instead: queues are
+# task-local, so once their task is dead no new work can reach them, and the entry can
+# be dropped (allowing the queue to be finalized) after a final synchronize.
+const queue_registry_lock = ReentrantLock()
+const queue_registry = Dict{Tuple{ZeContext,ZeDevice},Vector{Tuple{WeakRef,ZeCommandQueue}}}()
+
+# synchronize all known queues that target the given context (and device, if specified),
+# i.e., all queues whose in-flight work could possibly reference an allocation that is
+# about to be freed.
+function synchronize_all_queues(ctx::ZeContext, dev::Union{ZeDevice,Nothing})
+    queues = ZeCommandQueue[]
+    stale = Tuple{WeakRef,ZeCommandQueue}[]
+    GC.enable_finalizers(false)
+    try
+        @lock queue_registry_lock begin
+            for ((qctx, qdev), entries) in queue_registry
+                qctx == ctx || continue
+                (dev === nothing || qdev == dev) || continue
+                for entry in entries
+                    (task, queue) = entry
+                    queue.handle == C_NULL && continue  # finalized, handle destroyed
+                    push!(queues, queue)
+                    # entries whose task was already dead at this point cannot
+                    # receive new work, so they are safe to retire after the sync
+                    if task.value === nothing || istaskdone(task.value::Task)
+                        push!(stale, entry)
+                    end
+                end
+            end
+        end
+        # synchronize outside the lock: this can block for as long as a kernel runs,
+        # and finalizers running concurrently also need to take the lock. Keep
+        # finalizers disabled so none of the collected queues can be destroyed
+        # between collection and synchronization.
+        for queue in queues
+            oneL0.synchronize(queue)
+        end
+        # retire drained queues of dead tasks, allowing them to be finalized (the
+        # finalizer synchronizes once more before destroying the queue, in case
+        # the queue is dropped through other means).
+        if !isempty(stale)
+            @lock queue_registry_lock begin
+                for ((qctx, qdev), entries) in queue_registry
+                    qctx == ctx || continue
+                    (dev === nothing || qdev == dev) || continue
+                    filter!(entry -> !any(s -> s === entry, stale), entries)
+                end
+            end
+        end
+    finally
+        GC.enable_finalizers(true)
     end
+    return
 end
 
 """

src/pool.jl

@@ -83,6 +83,15 @@ 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.
+    synchronize_all_queues(oneL0.context(buf), oneL0.device(buf))
+
     free(buf; policy=oneL0.ZE_DRIVER_MEMORY_FREE_POLICY_EXT_FLAG_BLOCKING_FREE)
 
     # TODO: queue-ordered free from non-finalizer tasks once we have