(feat): trim! — manual reclamation of inactive pool slots

docs/src/basics/api-essentials.md

@@ -94,18 +94,49 @@ This pattern works because pool arrays share underlying storage with the pool.
 
 ### `reset!(pool)`
 
-Releases all memory held by the pool. Useful for long-running processes:
+Resets all active-slot counters to zero but **keeps** every allocated buffer for
+reuse. Use it to recover a clean pool state (e.g. after code acquired without a
+matching `rewind!`) without paying reallocation costs:
 
 ```julia
-# After processing a large batch
-@with_pool pool begin
-    # ... large computation ...
-end
-
-# Optionally release memory if pool grew too large
+# Mark all slots inactive, but keep buffers warm for the next batch
 reset!(get_task_local_pool())
 ```
 
+### `trim!(pool; force_gc=false)`
+
+Releases the pool's references to **all inactive** retained buffers (those
+already freed by `rewind!`/scope exit), while preserving the **active** arrays
+still in use. It does not selectively keep "warm" buffers — at top level, after a
+normal `rewind!` (where `n_active == 0`), every reusable buffer is released. The
+only guarantee is that arrays still active on the call stack are left untouched.
+
+```julia
+# After a scope rewinds, drop retained-but-unused buffers
+summary = trim!(get_task_local_pool())
+@show summary.slots_released summary.estimated_bytes_released
+
+# To release everything reusable, mark top-level slots inactive first:
+reset!(pool); trim!(pool)
+
+# force_gc=true additionally runs GC.gc() (no guarantee of immediate OS/VRAM return)
+trim!(pool; force_gc = true)
+```
+
+`estimated_bytes_released` is capacity-aware (it counts buffers even when runtime
+checks shrank them to length 0). It estimates detached storage; it is **not** a
+guarantee that RSS/VRAM dropped immediately — that depends on the GC.
+
+!!! note "Julia version"
+    Actual reclamation requires Julia 1.12+. On older Julia, `trim!` is a defined
+    no-op (returns a zero summary, warns once) so dependent packages stay portable
+    across the full supported Julia range.
+
+### `empty!(pool)`
+
+Drops **all** buffers (active and inactive) and resets the pool to its initial
+state. The most aggressive reclamation; invalidates any arrays still in use.
+
 ### `pooling_enabled(pool)`
 
 Check if pooling is active (returns `false` for `DisabledPool`):
@@ -131,7 +162,9 @@ end
 | `ones!(pool, [T,] dims...)` | `Array{T,N}` | 0 bytes | One-initialized |
 | `similar!(pool, A)` | `Array{T,N}` | 0 bytes | Match existing array |
 | `reshape!(pool, A, dims...)` | Reshaped array | 0 bytes (1.12+) | Reshape sharing memory |
-| `reset!(pool)` | `nothing` | - | Release all memory |
+| `reset!(pool)` | `pool` | - | Reset counters, keep buffers for reuse |
+| `trim!(pool; force_gc=false)` | summary | - | Release inactive buffers, keep active |
+| `empty!(pool)` | `pool` | - | Release all memory, reset state |
 | `pooling_enabled(pool)` | `Bool` | - | Check pool status |
 
 ## See Also

docs/src/reference/api.md

@@ -23,7 +23,9 @@
 | `rewind!(pool)` | Restores the pool to the last checkpoint, freeing all arrays acquired since then. |
 | `pool_stats(pool)` | Prints detailed statistics about pool usage. |
 | `get_task_local_pool()` | Returns the task-local pool instance. |
-| `empty!(pool)` | Clears all internal storage, releasing all memory. |
+| `reset!(pool)` | Resets active-slot counters to 0 but **keeps** all buffers for reuse. |
+| `trim!(pool; force_gc=false)` | Releases **inactive** retained buffers, keeping active arrays. Returns `(; slots_released, wrappers_released, estimated_bytes_released, gc_triggered)`. Works on CPU, CUDA, and Metal pools. Reclaims on Julia 1.12+; defined no-op (zero summary + one-time warning) on older Julia. |
+| `empty!(pool)` | Clears all internal storage, releasing **all** memory. |
 
 ### Convenience Functions
 

ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl

@@ -27,6 +27,10 @@ using AdaptiveArrayPools: allocate_vector, get_typed_pool!
 @inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{ComplexF64}) = p.complexf64
 @inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Bool}) = p.bool
 
+# Fixed-slot element types (dedicated struct fields above). Used by trim!(pool, T)
+# to recognize always-present types without creating an `others` entry.
+const _CUDA_FIXED_TYPES = Union{Float32, Float64, Float16, Int32, Int64, ComplexF32, ComplexF64, Bool}
+
 # Slow path: rare types via IdDict (with checkpoint correction!)
 @inline function AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{T}) where {T}
     return get!(p.others, T) do

ext/AdaptiveArrayPoolsCUDAExt/state.jl

@@ -352,3 +352,80 @@ function Base.empty!(pool::CuAdaptiveArrayPool)
 
     return pool
 end
+
+# ==============================================================================
+# trim! for CuAdaptiveArrayPool (parity with CPU/Metal)
+# ==============================================================================
+
+# CUDA backing storage lives on the GPU. Use `maxsize` (allocated device-buffer
+# bytes), NOT `sizeof`: in runtime-check mode (S>=1), rewind shrinks a released
+# slot's logical dims to 0 via `_resize_to_fit!` while PRESERVING `maxsize`, so
+# `sizeof` (logical length) would report 0 for a buffer that still holds device
+# capacity. `Base.summarysize` only sees the small CPU-side handle.
+function AdaptiveArrayPools._inactive_storage_bytes(tp::CuTypedPool, first::Int, last::Int)
+    total = 0
+    for i in first:last
+        @inbounds total += getfield(tp.vectors[i], :maxsize)
+    end
+    return total
+end
+
+function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool; force_gc::Bool = false)
+    # Mirror the CPU type-stable accumulation, but reuse the GPU-specific
+    # `foreach_fixed_slot` (different fixed-slot field set: float16, no Bit).
+    # `Ref` accumulators keep the closure box-free, so the summary stays concrete.
+    slots = Ref(0)
+    wrappers = Ref(0)
+    bytes = Ref(0)
+    AdaptiveArrayPools.foreach_fixed_slot(pool) do tp
+        c = AdaptiveArrayPools._trim_counts!(tp)
+        slots[] += c[1]
+        wrappers[] += c[2]
+        bytes[] += c[3]
+    end
+    oc = AdaptiveArrayPools._trim_others_counts!(values(pool.others))
+    # force_gc on CUDA also returns CUDA.jl's pooled blocks to the driver
+    # (CUDA.reclaim()), making freed VRAM available for reallocation. As the main
+    # docstring notes, immediate OS/VRAM return is still not guaranteed — the
+    # driver may keep the blocks in its own pool.
+    force_gc && (GC.gc(); CUDA.reclaim())
+    return AdaptiveArrayPools._trim_summary(
+        (slots[] + oc[1], wrappers[] + oc[2], bytes[] + oc[3]), force_gc,
+    )
+end
+
+# Guard + trim for one element type on CUDA. Mirrors the CPU `_trim_one_counts!`
+# but uses the CUDA fixed-type set; shared by the single-type and varargs forms.
+# Never creates a pool for a never-used fallback type (get_typed_pool! would
+# register a new fallback pool on a miss).
+@inline function AdaptiveArrayPools._trim_one_counts!(pool::CuAdaptiveArrayPool, ::Type{T})::NTuple{3, Int} where {T}
+    (!(T <: _CUDA_FIXED_TYPES) && !haskey(pool.others, T)) && return (0, 0, 0)
+    return AdaptiveArrayPools._trim_counts!(AdaptiveArrayPools.get_typed_pool!(pool, T))
+end
+
+@inline function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T}
+    counts = AdaptiveArrayPools._trim_one_counts!(pool, T)
+    # force_gc on CUDA also returns CUDA.jl's pooled blocks to the driver
+    # (CUDA.reclaim()), making freed VRAM available for reallocation. As the main
+    # docstring notes, immediate OS/VRAM return is still not guaranteed — the
+    # driver may keep the blocks in its own pool.
+    force_gc && (GC.gc(); CUDA.reclaim())
+    return AdaptiveArrayPools._trim_summary(counts, force_gc)
+end
+
+# Varargs form — parity with `checkpoint!`/`rewind!`, which also expose a `Type...`
+# overload on CUDA. Unrolled at compile time; one GC.gc()/CUDA.reclaim() after all
+# listed types are detached. Stays type-stable via `_trim_one_counts!`'s annotation.
+@generated function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool, types::Type...; force_gc::Bool = false)
+    n = length(types)
+    syms = [Symbol(:c, i) for i in 1:n]
+    assigns = [:($(syms[i]) = AdaptiveArrayPools._trim_one_counts!(pool, types[$i])) for i in 1:n]
+    counts_expr = n == 0 ? :((0, 0, 0)) :
+        Expr(:tuple, (Expr(:call, :+, (:($(s)[$j]) for s in syms)...) for j in 1:3)...)
+    return quote
+        $(assigns...)
+        counts = $counts_expr
+        force_gc && (GC.gc(); CUDA.reclaim())
+        return AdaptiveArrayPools._trim_summary(counts, force_gc)
+    end
+end

ext/AdaptiveArrayPoolsMetalExt/acquire.jl

@@ -195,8 +195,10 @@ Get an N-dimensional `MtlArray` from the pool with `setfield!`-based wrapper reu
 3. If different (rare: only after grow beyond capacity), update `:data` via refcount management
 
 ## Cache Miss (first call per (slot, N))
-Creates MtlArray wrapper sharing backing vector's GPU memory via `copy(vec.data)`,
-stores in `arr_wrappers[N][slot]` via `_store_arr_wrapper!` (reuses base module helper).
+Creates an MtlArray wrapper sharing the backing vector's GPU memory — the `DataRef`
+is passed directly (the constructor retains it; an extra `copy()` would double-count
+the refcount and pin the buffer). Stores it in `arr_wrappers[N][slot]` via
+`_store_arr_wrapper!` (reuses base module helper).
 """
 @inline function AdaptiveArrayPools.get_array!(tp::MetalTypedPool{T, S}, dims::NTuple{N, Int}) where {T, S, N}
     total_len = safe_prod(dims)
@@ -218,9 +220,12 @@ stores in `arr_wrappers[N][slot]` via `_store_arr_wrapper!` (reuses base module
         end
     end
 
-    # Cache miss: create wrapper sharing vec's GPU memory
+    # Cache miss: create wrapper sharing vec's GPU memory.
+    # Pass the DataRef directly — the MtlArray constructor already copies/retains
+    # it. An explicit copy() here double-counts the buffer's refcount, so it never
+    # reaches 0 and trim!/empty! cannot free the GPU memory.
     mtl = MtlArray{T, N, S}(
-        copy(getfield(vec, :data)), dims;
+        getfield(vec, :data), dims;
         maxsize = getfield(vec, :maxsize),
         offset = getfield(vec, :offset),
     )
@@ -283,9 +288,11 @@ Zero-allocation reshape for MtlArray using `setfield!`-based wrapper reuse.
         end
     end
 
-    # Cache miss (first call per slot+N): create wrapper, cache forever
+    # Cache miss (first call per slot+N): create wrapper, cache forever.
+    # Pass the DataRef directly (no copy()) — the constructor copies/retains it;
+    # an explicit copy() double-counts the refcount and leaks the buffer.
     mtl = MtlArray{T, N, S}(
-        copy(getfield(A, :data)), dims;
+        getfield(A, :data), dims;
         maxsize = getfield(A, :maxsize),
         offset = getfield(A, :offset),
     )

ext/AdaptiveArrayPoolsMetalExt/dispatch.jl

@@ -27,6 +27,10 @@ end
 @inline AdaptiveArrayPools.get_typed_pool!(p::MetalAdaptiveArrayPool, ::Type{ComplexF32}) = p.complexf32
 @inline AdaptiveArrayPools.get_typed_pool!(p::MetalAdaptiveArrayPool, ::Type{Bool}) = p.bool
 
+# Fixed-slot element types (dedicated struct fields above). Used by trim!(pool, T)
+# to recognize always-present types without creating an `others` entry.
+const _METAL_FIXED_TYPES = Union{Float32, Float16, Int32, Int64, ComplexF32, Bool}
+
 # Slow path: rare types via IdDict (with checkpoint correction!)
 # Explicitly reject Float64 and ComplexF64 (unsupported by Metal hardware).
 @inline function AdaptiveArrayPools.get_typed_pool!(p::MetalAdaptiveArrayPool, ::Type{T}) where {T}

ext/AdaptiveArrayPoolsMetalExt/state.jl

@@ -334,3 +334,72 @@ function Base.empty!(pool::MetalAdaptiveArrayPool)
 
     return pool
 end
+
+# ==============================================================================
+# trim! for MetalAdaptiveArrayPool (parity with CPU)
+# ==============================================================================
+
+# Metal backing storage lives on the GPU. Use `maxsize` (allocated device-buffer
+# bytes), NOT `sizeof`: in runtime-check mode (R>=1), rewind shrinks a released
+# slot's logical dims to 0 via `_resize_to_fit!` while PRESERVING `maxsize`, so
+# `sizeof` (logical length) would report 0 for a buffer that still holds device
+# capacity. `Base.summarysize` only sees the small CPU-side handle.
+function AdaptiveArrayPools._inactive_storage_bytes(tp::MetalTypedPool, first::Int, last::Int)
+    total = 0
+    for i in first:last
+        @inbounds total += getfield(tp.vectors[i], :maxsize)
+    end
+    return total
+end
+
+function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool; force_gc::Bool = false)
+    # Mirror the CPU type-stable accumulation, but reuse the Metal-specific
+    # `foreach_fixed_slot` (different fixed-slot field set). `Ref` accumulators keep
+    # the closure box-free, so the returned summary stays a concrete NamedTuple.
+    slots = Ref(0)
+    wrappers = Ref(0)
+    bytes = Ref(0)
+    AdaptiveArrayPools.foreach_fixed_slot(pool) do tp
+        c = AdaptiveArrayPools._trim_counts!(tp)
+        slots[] += c[1]
+        wrappers[] += c[2]
+        bytes[] += c[3]
+    end
+    oc = AdaptiveArrayPools._trim_others_counts!(values(pool.others))
+    force_gc && GC.gc()
+    return AdaptiveArrayPools._trim_summary(
+        (slots[] + oc[1], wrappers[] + oc[2], bytes[] + oc[3]), force_gc,
+    )
+end
+
+# Guard + trim for one element type on Metal. Mirrors the CPU `_trim_one_counts!`
+# but uses the Metal fixed-type set; shared by the single-type and varargs forms.
+# Never creates a pool for a never-used fallback type (get_typed_pool! would
+# register a new fallback pool on a miss).
+@inline function AdaptiveArrayPools._trim_one_counts!(pool::MetalAdaptiveArrayPool, ::Type{T})::NTuple{3, Int} where {T}
+    (!(T <: _METAL_FIXED_TYPES) && !haskey(pool.others, T)) && return (0, 0, 0)
+    return AdaptiveArrayPools._trim_counts!(AdaptiveArrayPools.get_typed_pool!(pool, T))
+end
+
+@inline function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T}
+    counts = AdaptiveArrayPools._trim_one_counts!(pool, T)
+    force_gc && GC.gc()
+    return AdaptiveArrayPools._trim_summary(counts, force_gc)
+end
+
+# Varargs form — parity with `checkpoint!`/`rewind!`, which also expose a `Type...`
+# overload on Metal. Unrolled at compile time; one GC.gc() after all listed types
+# are detached. Stays type-stable via `_trim_one_counts!`'s return annotation.
+@generated function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool, types::Type...; force_gc::Bool = false)
+    n = length(types)
+    syms = [Symbol(:c, i) for i in 1:n]
+    assigns = [:($(syms[i]) = AdaptiveArrayPools._trim_one_counts!(pool, types[$i])) for i in 1:n]
+    counts_expr = n == 0 ? :((0, 0, 0)) :
+        Expr(:tuple, (Expr(:call, :+, (:($(s)[$j]) for s in syms)...) for j in 1:3)...)
+    return quote
+        $(assigns...)
+        counts = $counts_expr
+        force_gc && GC.gc()
+        return AdaptiveArrayPools._trim_summary(counts, force_gc)
+    end
+end

scripts/trim_smoke_cuda.jl

@@ -0,0 +1,66 @@
+# CUDA `trim!` memory smoke test (run on an NVIDIA GPU machine).
+#
+# Not a unit test — a manual script that prints memory before/after `trim!` to
+# show whether reclamation actually happens. Run in an env with CUDA + AdaptiveArrayPools:
+#
+#   julia --project=<env> scripts/trim_smoke_cuda.jl
+#
+# Probes:
+#   julia_live = Base.gc_live_bytes()                        Julia-tracked live heap
+#   rss        = process RSS (ps)                             OS resident memory
+#   vram_used  = CUDA.total_memory() - CUDA.available_memory()  device VRAM in use
+#
+# On CUDA, trim!(force_gc=true) additionally calls CUDA.reclaim() to return
+# pooled blocks to the driver. NOTE: a Metal-specific wrapper double-copy bug
+# (since fixed) had prevented Metal from reclaiming. CUDA's CuArray constructor
+# takes ownership of its DataRef (no internal copy), so CUDA does not have that
+# bug and is expected to reclaim here. This script confirms it on real hardware:
+# vram_used should fall at step 3.
+
+using CUDA
+using AdaptiveArrayPools
+
+human(b) = Base.format_bytes(b)
+jlive() = Base.gc_live_bytes()
+rss() =
+try
+    parse(Int, readchomp(`ps -o rss= -p $(getpid())`)) * 1024
+catch
+    -1
+end
+vram_used() = Int(CUDA.total_memory() - CUDA.available_memory())
+
+function report(label)
+    return println(
+        rpad(label, 30),
+        "julia_live=", rpad(human(jlive()), 11),
+        "  rss=", rpad(rss() < 0 ? "?" : human(rss()), 11),
+        "  vram_used=", human(vram_used()),
+    )
+end
+
+const GiB = 1024^3
+const N = GiB ÷ sizeof(Float32)   # 1 GiB of Float32 per slot
+const NSLOTS = 4                  # 4 GiB total
+
+pool = get_task_local_cuda_pool()
+GC.gc(true); CUDA.reclaim()
+report("0) start")
+
+checkpoint!(pool)
+for i in 1:NSLOTS
+    v = acquire!(pool, Float32, N)
+    fill!(v, Float32(i))
+end
+CUDA.synchronize()
+report("1) allocated $(NSLOTS) GiB")
+
+rewind!(pool)                     # slots now inactive (retained for reuse)
+report("2) after rewind! (retained)")
+
+summary = trim!(pool; force_gc = true)   # detach inactive buffers + GC.gc() + CUDA.reclaim()
+GC.gc(true)
+report("3) after trim!(force_gc=true)")
+println("    summary = ", summary)
+println()
+CUDA.memory_status()              # detailed pool/driver VRAM breakdown