(feat): auto-compact — timer-driven background capacity compaction

ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl

@@ -33,6 +33,9 @@ using CUDA
     # State management (checkpoint!, rewind!, reset!, empty!)
     include("state.jl")
 
+    # Auto-compact: service the timer-flagged request at `@with_pool :cuda` entry
+    include("auto_compact.jl")
+
     # Safety: poisoning, escape detection, borrow tracking
     include("debug.jl")
 

ext/AdaptiveArrayPoolsCUDAExt/auto_compact.jl

@@ -0,0 +1,25 @@
+# ==============================================================================
+# CUDA Auto-Compact: service a flagged CUDA pool at its `@with_pool :cuda` entry.
+# ==============================================================================
+#
+# Parity with the CPU `_maybe_auto_compact!(::AdaptiveArrayPool)` in src/auto_compact.jl
+# and the Metal mirror. The single global Timer (base module, CPU thread) sets this pool's
+# `@atomic _compact_requested` flag via the shared WeakRef registry; the owner task runs
+# the full `compact!` here, at the `_current_depth == 1` safepoint where nothing is borrowed.
+#
+# The scope-entry hook itself is already emitted by the SHARED `@with_pool` codegen (it
+# expands `@with_pool :cuda` through the same `_generate_block_inner`), so this concrete
+# method is all that's needed to make the hook act on CUDA pools — before it, the hook
+# resolved to the `::Any` no-op fallback. Reuses the base `_run_auto_compact!` (CUDA
+# `compact!` takes the same factor/shrink_to/min_bytes/active kwargs as CPU).
+#
+# `@inline` so the common case (flag clear) is a cheap inlined monotonic read + compare;
+# the cold `_run_auto_compact!` stays a non-inlined call. `:monotonic` matches CPU: the
+# flag is a one-way eventual-visibility signal with no cross-field ordering invariant.
+
+@inline function AdaptiveArrayPools._maybe_auto_compact!(pool::CuAdaptiveArrayPool)
+    if (@atomic :monotonic pool._compact_requested) && pool._current_depth == 1
+        AdaptiveArrayPools._run_auto_compact!(pool)
+    end
+    return nothing
+end

ext/AdaptiveArrayPoolsCUDAExt/state.jl

@@ -466,10 +466,15 @@ end
 # reference. All buffer ownership goes through `_update_cuda_wrapper_data!`
 # (unsafe_free! old + copy/retain new), so the old device buffer's refcount reaches
 # zero and CUDA frees it.
-function AdaptiveArrayPools._compact_slot!(tp::CuTypedPool{T}, slot::Int, target::Int, used::Int) where {T}
+# `copy_live` (parity with CPU `_compact_slot!`): copy the `used` live elements into the
+# new device buffer. TRUE for ACTIVE slots — still held, must survive the shrink. FALSE
+# for INACTIVE slots: their contents are dead (the next `acquire!` returns uninitialized
+# device memory the caller fills), so the GPU→GPU `copyto!` is pure waste; skipping it
+# makes inactive compaction allocate-only (no device copy / kernel launch).
+function AdaptiveArrayPools._compact_slot!(tp::CuTypedPool{T}, slot::Int, target::Int, used::Int, copy_live::Bool) where {T}
     v = @inbounds tp.vectors[slot]
     nv = CuVector{T}(undef, target)
-    used > 0 && copyto!(nv, 1, v, 1, used)
+    copy_live && copyto!(nv, 1, v, 1, used)
     _update_cuda_wrapper_data!(v, nv)           # backing takes nv's buffer (refcounted)
     setfield!(v, :dims, (target,))
     for wrappers in tp.arr_wrappers
@@ -487,12 +492,18 @@ end
 function AdaptiveArrayPools._maybe_compact_slot!(tp::CuTypedPool{T}, slot::Int, factor::Real, shrink_to::Real, min_bytes::Int) where {T}
     used = AdaptiveArrayPools._slot_used(tp, slot)
     used == 0 && return 0
+    # RUNTIME_CHECK (S=1) invalidation poisons a released slot AND shrinks its logical length
+    # to 0 (device buffer retained) so an escaped view reads NaN/typemax. Compacting it would
+    # swap in fresh device storage and UNDO the poison — skip it (compacts once re-acquired).
+    # Parity with CPU `_maybe_compact_slot!`; `length == 0 && used > 0` ⟺ poisoned.
+    @inbounds(length(tp.vectors[slot])) == 0 && return 0
     cap = AdaptiveArrayPools._slot_capacity(@inbounds tp.vectors[slot])
     cap >= factor * used || return 0
     target = max(used, ceil(Int, min(Float64(cap), shrink_to * used)))   # clamp to [used, cap]
     reclaim = (cap - target) * _aligned_sizeof(T)
     reclaim >= min_bytes || return 0
-    AdaptiveArrayPools._compact_slot!(tp, slot, target, used)
+    # Copy live data only for ACTIVE slots (1:n_active); inactive slots hold dead data.
+    AdaptiveArrayPools._compact_slot!(tp, slot, target, used, slot <= tp.n_active)
     return reclaim
 end
 

ext/AdaptiveArrayPoolsCUDAExt/task_local_pool.jl

@@ -36,6 +36,10 @@ Uses `Dict{Int, CuAdaptiveArrayPool{RUNTIME_CHECK}}` in task-local storage, keye
     if pool === nothing
         pool = CuAdaptiveArrayPool()  # Uses RUNTIME_CHECK for initial S
         pools[dev_id] = pool
+        # Register for auto-compaction (gated by the compile-time const; mirrors CPU
+        # get_task_local_pool). One entry per (task, device) pool. DCE'd when the
+        # Preference is off, so no registry traffic on the zero-overhead build.
+        AdaptiveArrayPools.AUTO_COMPACT && AdaptiveArrayPools.register_auto_compact!(pool)
     end
 
     return pool::CuAdaptiveArrayPool{RUNTIME_CHECK}

ext/AdaptiveArrayPoolsCUDAExt/types.jl

@@ -125,6 +125,11 @@ mutable struct CuAdaptiveArrayPool{S} <: AbstractArrayPool
     _pending_callsite::String
     _pending_return_site::String
     _borrow_log::Union{Nothing, IdDict{Any, String}}
+
+    # Auto-compact request flag (parity with CPU AdaptiveArrayPool): set by the base
+    # module's background Timer sweep (a different thread), read + reset by the owner task
+    # at the `@with_pool :cuda` entry safepoint. Atomic for the cross-thread handoff.
+    @atomic _compact_requested::Bool
 end
 
 function CuAdaptiveArrayPool{S}() where {S}
@@ -145,7 +150,8 @@ function CuAdaptiveArrayPool{S}() where {S}
         CUDA.deviceid(dev),
         "",             # _pending_callsite
         "",             # _pending_return_site
-        nothing         # _borrow_log: lazily created when S >= 1
+        nothing,        # _borrow_log: lazily created when S >= 1
+        false           # _compact_requested: no pending auto-compact request
     )
 end
 
@@ -206,7 +212,8 @@ function _transfer_cuda_pool(::Val{V}, old::CuAdaptiveArrayPool) where {V}
         old.device_id,
         "",       # _pending_callsite: reset
         "",       # _pending_return_site: reset
-        V >= 1 ? IdDict{Any, String}() : nothing  # _borrow_log
+        V >= 1 ? IdDict{Any, String}() : nothing,  # _borrow_log
+        false                                       # _compact_requested: reset on migration
     )
 end
 

ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl

@@ -27,6 +27,7 @@ using Metal
     include("acquire.jl")
     include("task_local_pool.jl")
     include("state.jl")
+    include("auto_compact.jl")
     include("debug.jl")
     include("utils.jl")
     include("macros.jl")

ext/AdaptiveArrayPoolsMetalExt/auto_compact.jl

@@ -0,0 +1,25 @@
+# ==============================================================================
+# Metal Auto-Compact: service a flagged Metal pool at its `@with_pool :metal` entry.
+# ==============================================================================
+#
+# Parity with the CPU `_maybe_auto_compact!(::AdaptiveArrayPool)` in src/auto_compact.jl.
+# The single global Timer (base module, CPU thread) sets this pool's `@atomic
+# _compact_requested` flag via the shared WeakRef registry; the owner task runs the full
+# `compact!` here, at the `_current_depth == 1` safepoint where nothing is borrowed.
+#
+# The scope-entry hook itself is already emitted by the SHARED `@with_pool` codegen
+# (it expands `@with_pool :metal` through the same `_generate_block_inner`), so this
+# concrete method is all that's needed to make the hook act on Metal pools — before it,
+# the hook resolved to the `::Any` no-op fallback. Reuses the base `_run_auto_compact!`
+# (Metal `compact!` takes the same factor/shrink_to/min_bytes/active kwargs as CPU).
+#
+# `@inline` so the common case (flag clear) is a cheap inlined monotonic read + compare;
+# the cold `_run_auto_compact!` stays a non-inlined call. `:monotonic` matches CPU: the
+# flag is a one-way eventual-visibility signal with no cross-field ordering invariant.
+
+@inline function AdaptiveArrayPools._maybe_auto_compact!(pool::MetalAdaptiveArrayPool)
+    if (@atomic :monotonic pool._compact_requested) && pool._current_depth == 1
+        AdaptiveArrayPools._run_auto_compact!(pool)
+    end
+    return nothing
+end

ext/AdaptiveArrayPoolsMetalExt/state.jl

@@ -439,10 +439,15 @@ end
 # reference. All buffer ownership goes through `_update_metal_wrapper_data!`
 # (unsafe_free! old + copy/retain new), so the old device buffer's refcount reaches
 # zero and Metal frees it.
-function AdaptiveArrayPools._compact_slot!(tp::MetalTypedPool{T, S}, slot::Int, target::Int, used::Int) where {T, S}
+# `copy_live` (parity with CPU `_compact_slot!`): copy the `used` live elements into the
+# new device buffer. TRUE for ACTIVE slots — still held, must survive the shrink. FALSE
+# for INACTIVE slots: their contents are dead (the next `acquire!` returns uninitialized
+# device memory the caller fills), so the GPU→GPU `copyto!` is pure waste; skipping it
+# makes inactive compaction allocate-only (no device copy / kernel launch).
+function AdaptiveArrayPools._compact_slot!(tp::MetalTypedPool{T, S}, slot::Int, target::Int, used::Int, copy_live::Bool) where {T, S}
     v = @inbounds tp.vectors[slot]
     nv = MtlArray{T, 1, S}(undef, target)
-    used > 0 && copyto!(nv, 1, v, 1, used)
+    copy_live && copyto!(nv, 1, v, 1, used)
     _update_metal_wrapper_data!(v, nv)          # backing takes nv's buffer (refcounted)
     setfield!(v, :dims, (target,))
     for wrappers in tp.arr_wrappers
@@ -460,12 +465,18 @@ end
 function AdaptiveArrayPools._maybe_compact_slot!(tp::MetalTypedPool{T, S}, slot::Int, factor::Real, shrink_to::Real, min_bytes::Int) where {T, S}
     used = AdaptiveArrayPools._slot_used(tp, slot)
     used == 0 && return 0
+    # RUNTIME_CHECK (R=1) invalidation poisons a released slot AND shrinks its logical length
+    # to 0 (device buffer retained) so an escaped view reads NaN/typemax. Compacting it would
+    # swap in fresh device storage and UNDO the poison — skip it (compacts once re-acquired).
+    # Parity with CPU `_maybe_compact_slot!`; `length == 0 && used > 0` ⟺ poisoned.
+    @inbounds(length(tp.vectors[slot])) == 0 && return 0
     cap = AdaptiveArrayPools._slot_capacity(@inbounds tp.vectors[slot])
     cap >= factor * used || return 0
     target = max(used, ceil(Int, min(Float64(cap), shrink_to * used)))   # clamp to [used, cap]
     reclaim = (cap - target) * sizeof(T)
     reclaim >= min_bytes || return 0
-    AdaptiveArrayPools._compact_slot!(tp, slot, target, used)
+    # Copy live data only for ACTIVE slots (1:n_active); inactive slots hold dead data.
+    AdaptiveArrayPools._compact_slot!(tp, slot, target, used, slot <= tp.n_active)
     return reclaim
 end
 

ext/AdaptiveArrayPoolsMetalExt/task_local_pool.jl

@@ -37,6 +37,10 @@ Uses `Dict{UInt64, MetalAdaptiveArrayPool{RUNTIME_CHECK, METAL_STORAGE}}` in tas
     if pool === nothing
         pool = MetalAdaptiveArrayPool()  # Uses RUNTIME_CHECK for initial R
         pools[dev_key] = pool
+        # Register for auto-compaction (gated by the compile-time const; mirrors CPU
+        # get_task_local_pool). One entry per (task, device) pool. DCE'd when the
+        # Preference is off, so no registry traffic on the zero-overhead build.
+        AdaptiveArrayPools.AUTO_COMPACT && AdaptiveArrayPools.register_auto_compact!(pool)
     end
 
     return pool::MetalAdaptiveArrayPool{RUNTIME_CHECK, METAL_STORAGE}

ext/AdaptiveArrayPoolsMetalExt/types.jl

@@ -121,6 +121,11 @@ mutable struct MetalAdaptiveArrayPool{R, S} <: AbstractArrayPool
     _pending_callsite::String
     _pending_return_site::String
     _borrow_log::Union{Nothing, IdDict{Any, String}}
+
+    # Auto-compact request flag (parity with CPU AdaptiveArrayPool): set by the base
+    # module's background Timer sweep (a different thread), read + reset by the owner task
+    # at the `@with_pool :metal` entry safepoint. Atomic for the cross-thread handoff.
+    @atomic _compact_requested::Bool
 end
 
 function MetalAdaptiveArrayPool{R, S}() where {R, S}
@@ -139,6 +144,7 @@ function MetalAdaptiveArrayPool{R, S}() where {R, S}
         "",             # _pending_callsite
         "",             # _pending_return_site
         nothing,        # _borrow_log: lazily created when R >= 1
+        false,          # _compact_requested: no pending auto-compact request
     )
 end
 

src/AdaptiveArrayPools.jl

@@ -17,6 +17,7 @@ export @with_pool, @maybe_with_pool, @safe_with_pool, @safe_maybe_with_pool
 export STATIC_POOLING, MAYBE_POOLING, RUNTIME_CHECK
 export PoolEscapeError, EscapePoint, PoolMutationError, MutationPoint
 export checkpoint!, rewind!, reset!, trim!, compact!
+export enable_auto_compact!, disable_auto_compact!, auto_compact_enabled, AUTO_COMPACT
 export get_task_local_cuda_pool, get_task_local_cuda_pools  # CUDA (stubs, overridden by extension)
 export get_task_local_metal_pool, get_task_local_metal_pools  # Metal (stubs, overridden by extension)
 
@@ -33,6 +34,7 @@ export DisabledPool, DISABLED_CPU, pooling_enabled  # Disabled pool support
     include("bitarray.jl")
     include("convenience.jl")
     include("state.jl")
+    include("auto_compact.jl")
     include("task_local_pool.jl")
     include("debug.jl")
     include("macros.jl")
@@ -44,6 +46,7 @@ else
     include("legacy/bitarray.jl")
     include("convenience.jl")
     include("legacy/state.jl")
+    include("legacy/auto_compact.jl")
     include("task_local_pool.jl")
     include("debug.jl")
     include("macros.jl")