(feat): compact! — in-place capacity compaction of pool buffers

ext/AdaptiveArrayPoolsCUDAExt/acquire.jl

@@ -96,13 +96,15 @@ CUDA.jl's resize! unnecessarily (especially after safety invalidation sets dims=
     idx = tp.n_active
     if idx > length(tp.vectors)
         push!(tp.vectors, allocate_vector(tp, total_len))
+        push!(tp.slot_extents, total_len)        # parallel to `vectors` (for compact!/_slot_used)
         _check_pool_growth(tp, idx)
     else
         # _resize_to_fit! handles all cases:
         # - n > capacity: resize! (GPU alloc)
         # - n != length: setfield!(:dims) — restores length after safety invalidation
         # - n == length: no-op (hot path)
         _resize_to_fit!(@inbounds(tp.vectors[idx]), total_len)
+        @inbounds tp.slot_extents[idx] = total_len   # current logical extent
     end
     return idx
 end
@@ -119,7 +121,10 @@ the wrapper points to a different array's memory via `setfield!(:data)`.
     idx = tp.n_active
     if idx > length(tp.vectors)
         push!(tp.vectors, CuVector{T}(undef, 0))
+        push!(tp.slot_extents, 0)        # placeholder slot: wrapper points elsewhere (reshape!)
         _check_pool_growth(tp, idx)
+    else
+        @inbounds tp.slot_extents[idx] = 0
     end
     return idx
 end

ext/AdaptiveArrayPoolsCUDAExt/state.jl

@@ -317,6 +317,7 @@ CUDA.reclaim()
 function Base.empty!(tp::CuTypedPool)
     empty!(tp.vectors)
     empty!(tp.arr_wrappers)
+    empty!(tp.slot_extents)
     tp.n_active = 0
     # Restore sentinel values
     empty!(tp._checkpoint_n_active)
@@ -429,3 +430,122 @@ end
         return AdaptiveArrayPools._trim_summary(counts, force_gc)
     end
 end
+
+# ==============================================================================
+# compact! for CuAdaptiveArrayPool (parity with CPU / Metal)
+# ==============================================================================
+#
+# Shrinks an over-allocated device buffer in place. Like Metal (and unlike CPU),
+# CUDA has no uncached views — acquire_view! == acquire! → cached CuArray wrapper —
+# so a slot's live references are exactly its cached wrappers; compaction swaps the
+# backing buffer's DataRef and re-syncs those wrappers via the same refcount helper
+# used on grow (`_update_cuda_wrapper_data!`). `slot_extents` (set in
+# `_cuda_claim_slot!`) gives the live extent, matching the CPU `_slot_used` model.
+# Device capacity uses `_aligned_sizeof(T)` (CUDA buffer alignment), like `_resize_to_fit!`.
+
+# Device capacity (elements) of a CUDA backing buffer.
+@inline AdaptiveArrayPools._slot_capacity(v::CuVector{T}) where {T} =
+    Int(getfield(v, :maxsize) ÷ _aligned_sizeof(T))
+
+# Current logical extent of a slot (records both acquire!/acquire_view! since CUDA
+# routes both through `_cuda_claim_slot!`). `::Int` keeps the gate allocation-free.
+@inline function AdaptiveArrayPools._slot_used(tp::CuTypedPool, slot::Int)::Int
+    ext = tp.slot_extents
+    return slot <= length(ext) ? (@inbounds ext[slot]) : 0
+end
+
+# Keep `slot_extents` parallel to `vectors` when `trim!` truncates a CUDA pool.
+@inline function AdaptiveArrayPools._trim_slot_extents!(tp::CuTypedPool, keep::Int)
+    length(tp.slot_extents) > keep && resize!(tp.slot_extents, keep)
+    return nothing
+end
+
+# Shrink one slot's device buffer to capacity `target` in place: allocate a smaller
+# CuVector, copy the live `used` elements, swap the backing's DataRef (keeping the
+# backing object's identity), re-sync cached wrappers, and drop the temporary's own
+# 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}
+    v = @inbounds tp.vectors[slot]
+    nv = CuVector{T}(undef, target)
+    used > 0 && 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
+        wrappers === nothing && continue
+        slot <= length(wrappers) || continue
+        w = @inbounds wrappers[slot]
+        w === nothing && continue
+        _update_cuda_wrapper_data!(w, v)        # held wrapper re-points to new buffer (dims kept)
+    end
+    unsafe_free!(getfield(nv, :data))           # drop the temporary's own reference
+    return nothing
+end
+
+# Gate + compact one CUDA slot. Returns reclaimed device bytes (0 if skipped).
+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
+    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)
+    return reclaim
+end
+
+# Guard + compact for one element type on CUDA (mirrors `_trim_one_counts!`): never
+# creates a pool for a never-used fallback type.
+@inline function AdaptiveArrayPools._compact_one_counts!(pool::CuAdaptiveArrayPool, ::Type{T}, factor::Real, shrink_to::Real, min_bytes::Int, active::Bool)::NTuple{2, Int} where {T}
+    (!(T <: _CUDA_FIXED_TYPES) && !haskey(pool.others, T)) && return (0, 0)
+    return AdaptiveArrayPools._compact_counts!(AdaptiveArrayPools.get_typed_pool!(pool, T), factor, shrink_to, min_bytes, active)
+end
+
+function AdaptiveArrayPools.compact!(
+        pool::CuAdaptiveArrayPool;
+        factor::Real = 10, shrink_to::Real = 1.5, min_bytes::Int = 2^20,
+        active::Bool = true, force_gc::Bool = false,
+    )
+    # Mirror the CUDA `trim!`: Ref accumulators keep the `foreach_fixed_slot` closure
+    # box-free over the CUDA fixed-slot set, so the returned summary stays concrete.
+    slots = Ref(0)
+    bytes = Ref(0)
+    AdaptiveArrayPools.foreach_fixed_slot(pool) do tp
+        c = AdaptiveArrayPools._compact_counts!(tp, factor, shrink_to, min_bytes, active)
+        slots[] += c[1]
+        bytes[] += c[2]
+    end
+    oc = AdaptiveArrayPools._compact_others_counts!(values(pool.others), factor, shrink_to, min_bytes, active)
+    force_gc && (GC.gc(); CUDA.reclaim())
+    return AdaptiveArrayPools._compact_summary((slots[] + oc[1], bytes[] + oc[2]), force_gc)
+end
+
+@inline function AdaptiveArrayPools.compact!(
+        pool::CuAdaptiveArrayPool, ::Type{T};
+        factor::Real = 10, shrink_to::Real = 1.5, min_bytes::Int = 2^20,
+        active::Bool = true, force_gc::Bool = false,
+    ) where {T}
+    counts = AdaptiveArrayPools._compact_one_counts!(pool, T, factor, shrink_to, min_bytes, active)
+    force_gc && (GC.gc(); CUDA.reclaim())
+    return AdaptiveArrayPools._compact_summary(counts, force_gc)
+end
+
+@generated function AdaptiveArrayPools.compact!(
+        pool::CuAdaptiveArrayPool, types::Type...;
+        factor::Real = 10, shrink_to::Real = 1.5, min_bytes::Int = 2^20,
+        active::Bool = true, force_gc::Bool = false,
+    )
+    n = length(types)
+    syms = [Symbol(:c, i) for i in 1:n]
+    assigns = [:($(syms[i]) = AdaptiveArrayPools._compact_one_counts!(pool, types[$i], factor, shrink_to, min_bytes, active)) for i in 1:n]
+    counts_expr = n == 0 ? :((0, 0)) :
+        Expr(:tuple, (Expr(:call, :+, (:($(s)[$j]) for s in syms)...) for j in 1:2)...)
+    return quote
+        $(assigns...)
+        counts = $counts_expr
+        force_gc && (GC.gc(); CUDA.reclaim())
+        return AdaptiveArrayPools._compact_summary(counts, force_gc)
+    end
+end

ext/AdaptiveArrayPoolsCUDAExt/types.jl

@@ -36,6 +36,11 @@ mutable struct CuTypedPool{T} <: AbstractTypedPool{T, CuVector{T}}
     # --- N-D Wrapper Cache (setfield!-based reuse, matches CPU TypedPool) ---
     arr_wrappers::Vector{Union{Nothing, Vector{Any}}}  # index=N (dimensionality), value=per-slot CuArray{T,N}
 
+    # --- Per-slot current logical extent (parallel to `vectors`, matches CPU TypedPool) ---
+    # Size of each slot's most recent `_cuda_claim_slot!`; `compact!` reads it
+    # (`_slot_used`) to know how much of an over-allocated device buffer is in use.
+    slot_extents::Vector{Int}
+
     # --- State Management (1-based sentinel pattern) ---
     n_active::Int
     _checkpoint_n_active::Vector{Int}
@@ -46,6 +51,7 @@ function CuTypedPool{T}() where {T}
     return CuTypedPool{T}(
         CuVector{T}[],                   # vectors
         Union{Nothing, Vector{Any}}[],   # arr_wrappers (indexed by N)
+        Int[],                           # slot_extents (parallel to vectors)
         0, [0], [0]                      # State (1-based sentinel)
     )
 end

ext/AdaptiveArrayPoolsMetalExt/acquire.jl

@@ -80,9 +80,11 @@ Metal.jl's resize! unnecessarily (especially after safety invalidation sets dims
     idx = tp.n_active
     if idx > length(tp.vectors)
         push!(tp.vectors, allocate_vector(tp, total_len))
+        push!(tp.slot_extents, total_len)        # parallel to `vectors` (for compact!/_slot_used)
         _check_pool_growth(tp, idx)
     else
         _resize_to_fit!(@inbounds(tp.vectors[idx]), total_len)
+        @inbounds tp.slot_extents[idx] = total_len   # current logical extent
     end
     return idx
 end
@@ -99,7 +101,10 @@ the wrapper points to a different array's memory via `setfield!(:data)`.
     idx = tp.n_active
     if idx > length(tp.vectors)
         push!(tp.vectors, MtlArray{T, 1, S}(undef, 0))
+        push!(tp.slot_extents, 0)        # placeholder slot: wrapper points elsewhere (reshape!)
         _check_pool_growth(tp, idx)
+    else
+        @inbounds tp.slot_extents[idx] = 0
     end
     return idx
 end

ext/AdaptiveArrayPoolsMetalExt/state.jl

@@ -299,6 +299,7 @@ Actual VRAM release depends on GC + Metal.jl's memory pool.
 function Base.empty!(tp::MetalTypedPool)
     empty!(tp.vectors)
     empty!(tp.arr_wrappers)
+    empty!(tp.slot_extents)
     tp.n_active = 0
     # Restore sentinel values
     empty!(tp._checkpoint_n_active)
@@ -403,3 +404,122 @@ end
         return AdaptiveArrayPools._trim_summary(counts, force_gc)
     end
 end
+
+# ==============================================================================
+# compact! for MetalAdaptiveArrayPool (parity with CPU)
+# ==============================================================================
+#
+# Shrinks an over-allocated device buffer in place. Unlike CPU, Metal has no
+# uncached views (`acquire_view!` == `acquire!` → cached `MtlArray` wrapper), so a
+# slot's live references are exactly its cached wrappers; compaction swaps the
+# backing buffer's `DataRef` and re-syncs those wrappers via the same refcount
+# helper used on grow (`_update_metal_wrapper_data!`). `slot_extents` (set in
+# `_metal_claim_slot!`) gives the live extent, matching the CPU `_slot_used` model.
+
+# Device capacity (elements) of a Metal backing buffer.
+@inline AdaptiveArrayPools._slot_capacity(v::MtlArray{T}) where {T} =
+    Int(getfield(v, :maxsize) ÷ sizeof(T))
+
+# Current logical extent of a slot (records both acquire!/acquire_view! since Metal
+# routes both through `_metal_claim_slot!`). `::Int` keeps the gate allocation-free.
+@inline function AdaptiveArrayPools._slot_used(tp::MetalTypedPool, slot::Int)::Int
+    ext = tp.slot_extents
+    return slot <= length(ext) ? (@inbounds ext[slot]) : 0
+end
+
+# Keep `slot_extents` parallel to `vectors` when `trim!` truncates a Metal pool.
+@inline function AdaptiveArrayPools._trim_slot_extents!(tp::MetalTypedPool, keep::Int)
+    length(tp.slot_extents) > keep && resize!(tp.slot_extents, keep)
+    return nothing
+end
+
+# Shrink one slot's device buffer to capacity `target` in place: allocate a smaller
+# MtlArray, copy the live `used` elements, swap the backing's `DataRef` (keeping the
+# backing object's identity), re-sync cached wrappers, and drop the temporary's own
+# 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}
+    v = @inbounds tp.vectors[slot]
+    nv = MtlArray{T, 1, S}(undef, target)
+    used > 0 && 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
+        wrappers === nothing && continue
+        slot <= length(wrappers) || continue
+        w = @inbounds wrappers[slot]
+        w === nothing && continue
+        _update_metal_wrapper_data!(w, v)       # held wrapper re-points to new buffer (dims kept)
+    end
+    unsafe_free!(getfield(nv, :data))           # drop the temporary's own reference
+    return nothing
+end
+
+# Gate + compact one Metal slot. Returns reclaimed device bytes (0 if skipped).
+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
+    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)
+    return reclaim
+end
+
+# Guard + compact for one element type on Metal (mirrors `_trim_one_counts!`): never
+# creates a pool for a never-used fallback type.
+@inline function AdaptiveArrayPools._compact_one_counts!(pool::MetalAdaptiveArrayPool, ::Type{T}, factor::Real, shrink_to::Real, min_bytes::Int, active::Bool)::NTuple{2, Int} where {T}
+    (!(T <: _METAL_FIXED_TYPES) && !haskey(pool.others, T)) && return (0, 0)
+    return AdaptiveArrayPools._compact_counts!(AdaptiveArrayPools.get_typed_pool!(pool, T), factor, shrink_to, min_bytes, active)
+end
+
+function AdaptiveArrayPools.compact!(
+        pool::MetalAdaptiveArrayPool;
+        factor::Real = 10, shrink_to::Real = 1.5, min_bytes::Int = 2^20,
+        active::Bool = true, force_gc::Bool = false,
+    )
+    # Mirror the Metal `trim!`: Ref accumulators keep the `foreach_fixed_slot`
+    # closure box-free over the Metal-specific fixed-slot set, so the returned
+    # summary stays a concrete NamedTuple.
+    slots = Ref(0)
+    bytes = Ref(0)
+    AdaptiveArrayPools.foreach_fixed_slot(pool) do tp
+        c = AdaptiveArrayPools._compact_counts!(tp, factor, shrink_to, min_bytes, active)
+        slots[] += c[1]
+        bytes[] += c[2]
+    end
+    oc = AdaptiveArrayPools._compact_others_counts!(values(pool.others), factor, shrink_to, min_bytes, active)
+    force_gc && GC.gc()
+    return AdaptiveArrayPools._compact_summary((slots[] + oc[1], bytes[] + oc[2]), force_gc)
+end
+
+@inline function AdaptiveArrayPools.compact!(
+        pool::MetalAdaptiveArrayPool, ::Type{T};
+        factor::Real = 10, shrink_to::Real = 1.5, min_bytes::Int = 2^20,
+        active::Bool = true, force_gc::Bool = false,
+    ) where {T}
+    counts = AdaptiveArrayPools._compact_one_counts!(pool, T, factor, shrink_to, min_bytes, active)
+    force_gc && GC.gc()
+    return AdaptiveArrayPools._compact_summary(counts, force_gc)
+end
+
+@generated function AdaptiveArrayPools.compact!(
+        pool::MetalAdaptiveArrayPool, types::Type...;
+        factor::Real = 10, shrink_to::Real = 1.5, min_bytes::Int = 2^20,
+        active::Bool = true, force_gc::Bool = false,
+    )
+    n = length(types)
+    syms = [Symbol(:c, i) for i in 1:n]
+    assigns = [:($(syms[i]) = AdaptiveArrayPools._compact_one_counts!(pool, types[$i], factor, shrink_to, min_bytes, active)) for i in 1:n]
+    counts_expr = n == 0 ? :((0, 0)) :
+        Expr(:tuple, (Expr(:call, :+, (:($(s)[$j]) for s in syms)...) for j in 1:2)...)
+    return quote
+        $(assigns...)
+        counts = $counts_expr
+        force_gc && GC.gc()
+        return AdaptiveArrayPools._compact_summary(counts, force_gc)
+    end
+end

ext/AdaptiveArrayPoolsMetalExt/types.jl

@@ -34,6 +34,11 @@ mutable struct MetalTypedPool{T, S} <: AbstractTypedPool{T, MtlArray{T, 1, S}}
     # --- N-D Wrapper Cache (setfield!-based reuse, matches CPU TypedPool) ---
     arr_wrappers::Vector{Union{Nothing, Vector{Any}}}  # index=N (dimensionality), value=per-slot MtlArray{T,N,S}
 
+    # --- Per-slot current logical extent (parallel to `vectors`, matches CPU TypedPool) ---
+    # Size of each slot's most recent `_metal_claim_slot!`; `compact!` reads it
+    # (`_slot_used`) to know how much of an over-allocated device buffer is in use.
+    slot_extents::Vector{Int}
+
     # --- State Management (1-based sentinel pattern) ---
     n_active::Int
     _checkpoint_n_active::Vector{Int}
@@ -44,6 +49,7 @@ function MetalTypedPool{T, S}() where {T, S}
     return MetalTypedPool{T, S}(
         MtlArray{T, 1, S}[],                # vectors
         Union{Nothing, Vector{Any}}[],       # arr_wrappers (indexed by N)
+        Int[],                               # slot_extents (parallel to vectors)
         0, [0], [0],                         # State (1-based sentinel)
     )
 end

src/AdaptiveArrayPools.jl

@@ -16,7 +16,7 @@ export Bit  # Sentinel type for BitArray (use with acquire!, trues!, falses!)
 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!
+export checkpoint!, rewind!, reset!, trim!, 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)
 

src/acquire.jl

@@ -67,12 +67,14 @@ Returns the slot index. This is the shared primitive for all acquisition paths
     idx = tp.n_active
     if idx > length(tp.vectors)
         push!(tp.vectors, allocate_vector(tp, n))
+        push!(tp.slot_extents, n)        # parallel to `vectors`; records this slot's extent
         _check_pool_growth(tp, idx)
     else
         @inbounds vec = tp.vectors[idx]
         if length(vec) < n
             resize!(vec, n)
         end
+        @inbounds tp.slot_extents[idx] = n   # current logical extent (for compact!/_slot_used)
     end
     return idx
 end
@@ -89,7 +91,10 @@ the wrapper points to a different array's memory via `setfield!(:ref)`.
     idx = tp.n_active
     if idx > length(tp.vectors)
         push!(tp.vectors, Vector{T}(undef, 0))
+        push!(tp.slot_extents, 0)        # placeholder slot: backing points elsewhere (reshape!)
         _check_pool_growth(tp, idx)
+    else
+        @inbounds tp.slot_extents[idx] = 0
     end
     return idx
 end