Flush deferred sparse releases during memory reclaim

Author: michel2323

lib/level-zero/utils.jl

@@ -1,5 +1,23 @@
 isdebug(group) = Base.CoreLogging.current_logger_for_env(Base.CoreLogging.Debug, group, oneL0) !== nothing
 
+# 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.
+const _reclaim_callbacks = Function[]
+
+function register_reclaim_callback!(f::Function)
+    push!(_reclaim_callbacks, f)
+end
+
+function _run_reclaim_callbacks()
+    for cb in _reclaim_callbacks
+        try
+            cb()
+        catch
+        end
+    end
+end
+
 function retry_reclaim(f, isfailed)
     ret = f()
 
@@ -11,6 +29,12 @@ function retry_reclaim(f, isfailed)
                 GC.gc(false)
             elseif phase == 2
                 GC.gc(true)
+            elseif phase == 3
+                # After GC, finalizers may have deferred resource releases (e.g., MKL
+                # sparse handles).  Flush them now, then GC again to free the memory
+                # those releases made available.
+                _run_reclaim_callbacks()
+                GC.gc(true)
             else
                 break
             end

lib/mkl/oneMKL.jl

@@ -31,6 +31,11 @@ include("linalg.jl")
 include("interfaces.jl")
 include("fft.jl")
 
+# Register deferred sparse handle flush as a memory reclaim callback so that OOM
+# recovery (retry_reclaim) and proactive GC (_maybe_gc) can free MKL internal buffers
+# associated with sparse matrix handles that were deferred from finalizer threads.
+oneL0.register_reclaim_callback!(flush_deferred_sparse_releases)
+
 function version()
     major = Ref{Int64}()
     minor = Ref{Int64}()

src/pool.jl

@@ -18,8 +18,11 @@ function _maybe_gc(dev, bytes)
     allocated <= 0 && return
     total_mem = _get_total_mem(dev)
     if allocated + bytes > total_mem * 0.8
+        # Flush deferred resource releases (e.g., MKL sparse handles) from previous GC
+        # cycles first — these are safe to release now because they were deferred earlier.
+        # Do this BEFORE GC to avoid racing with new finalizers.
+        oneL0._run_reclaim_callbacks()
         # Full GC to collect old-generation objects whose finalizers free GPU memory.
-        # GC.gc(false) only does minor collection which won't reclaim promoted objects.
         GC.gc(true)
     elseif allocated + bytes > total_mem * 0.4
         GC.gc(false)