From 20540407d79410eeedfcb48b5c8d5b9220ecef0f Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 14:40:39 -0700 Subject: [PATCH 1/7] =?UTF-8?q?(feat):=20`trim!`=20=E2=80=94=20manual=20re?= =?UTF-8?q?clamation=20of=20inactive=20pool=20slots=20(CPU=20+=20Metal)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `trim!(pool; force_gc=false)` to release backing buffers and cached wrappers for INACTIVE retained slots (n_active+1:end) while keeping ACTIVE slots intact. Fills the gap between `reset!` (keeps all memory) and `empty!` (drops everything). - Generic core on AbstractTypedPool: truncates `vectors` AND `arr_wrappers` to n_active. Cached wrappers pin backing buffers via :ref, so both must be dropped for the memory to become GC-eligible. - Capacity-aware byte estimate: `summarysize` (CPU), `sizeof` (Metal GPU). - Variants: per-type `trim!(pool, T)`, zero-arg `trim!()` (task-local pool), `DisabledPool` no-op. Returns `(; slots_released, wrappers_released, estimated_bytes_released, gc_triggered)`. - Metal parity (CuTypedPool deferred). Legacy (Julia <1.12) defines `trim!` as an exported no-op (zero summary + one-time @warn) so dependents stay portable across the full [compat] julia range. - Tests pin: active slots untouched (data + identity), inactive dropped, wrapper-cache truncation, byte accounting (S=0/S=1), self-heal, Bit/fallback. Verified on Julia 1.12 (CPU + Metal) and 1.10 LTS (legacy no-op). Docs: document `trim!` and correct a prior inaccuracy (`reset!` preserves buffers; it does not release memory). --- docs/design/selective_memory_reclamation.md | 414 ++++++++++++++++++++ docs/src/basics/api-essentials.md | 49 ++- docs/src/reference/api.md | 4 +- ext/AdaptiveArrayPoolsMetalExt/state.jl | 50 +++ src/AdaptiveArrayPools.jl | 2 +- src/legacy/state.jl | 33 ++ src/state.jl | 145 +++++++ test/metal/runtests.jl | 1 + test/metal/test_trim.jl | 109 ++++++ test/runtests.jl | 2 + test/test_trim.jl | 227 +++++++++++ test/test_trim_legacy.jl | 25 ++ 12 files changed, 1051 insertions(+), 10 deletions(-) create mode 100644 docs/design/selective_memory_reclamation.md create mode 100644 test/metal/test_trim.jl create mode 100644 test/test_trim.jl create mode 100644 test/test_trim_legacy.jl diff --git a/docs/design/selective_memory_reclamation.md b/docs/design/selective_memory_reclamation.md new file mode 100644 index 00000000..96c9cecb --- /dev/null +++ b/docs/design/selective_memory_reclamation.md @@ -0,0 +1,414 @@ +# Inactive Slot Trimming for AdaptiveArrayPools + +- **Date:** 2026-06-28 +- **Status:** Revised design — pending user review (not yet implemented) +- **Scope:** New manual `trim!` API. CPU + Metal in first cut; CUDA deferred. +- **Author:** design session + +--- + +## 1. Motivation + +AdaptiveArrayPools keeps backing buffers after scopes exit so later `acquire!` +calls can reuse memory and stay allocation-free on the hot path. This is the +right default for normal workloads. + +The downside is memory retention after an unusual workload. A rare large input +can grow one or more slots far beyond the usual working set. After the scope +rewinds, those slots are no longer active, but the pool still keeps strong +references to their backing buffers and cached wrappers. Julia's GC cannot +collect those buffers while the pool still references them. + +We want a manual operation that says: + +> Keep the arrays currently in use. Drop the inactive pooled buffers that are +> only being retained for possible future reuse. + +This is not an immediate OS/VRAM release guarantee. The goal is to release the +pool-held references so the old buffers become GC-eligible once no user code +still references them. + +## 2. Core Concept: Active vs Inactive Slots + +Each typed pool has: + +- **active slots:** `1:tp.n_active` + - Currently claimed by the live call stack. + - Must remain valid. +- **inactive retained slots:** `(tp.n_active + 1):length(tp.vectors)` + - Already released by `rewind!` or previous scope exit. + - Kept only as a reuse cache. + - Safe for a manual cleanup operation to detach from the pool. + +`trim!` operates only on inactive retained slots. It never touches active slots. + +This makes the API easier to reason about than "deep" or "level" terminology: +the persistent thing is the typed pool's slot array, not a durable level object. + +## 3. Proposed API + +```julia +trim!(pool::AdaptiveArrayPool; force_gc::Bool = false) +trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) +trim!(; force_gc::Bool = false) +``` + +Metal extension: + +```julia +trim!(pool::MetalAdaptiveArrayPool; force_gc::Bool = false) +trim!(pool::MetalAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) +``` + +Disabled pools: + +```julia +trim!(::DisabledPool; force_gc::Bool = false) +trim!(::DisabledPool, ::Type{T}; force_gc::Bool = false) +``` + +Return a small summary: + +```julia +(; + slots_released::Int, + wrappers_released::Int, + estimated_bytes_released::Int, + gc_triggered::Bool, +) +``` + +`force_gc=false` is the default. When `true`, `trim!` calls `GC.gc()` after +detaching inactive references. This still does not promise immediate OS or GPU +memory return; it only asks Julia to run collection after the pool stops holding +the buffers. + +## 4. Semantics + +For each affected typed pool: + +1. Let `keep = tp.n_active`. +2. Estimate storage held by inactive slots for the return summary. +3. Drop backing-buffer references by truncating `tp.vectors` to `keep`. +4. Drop cached wrapper references for inactive slots by truncating each + wrapper-cache vector to `keep`. +5. Leave checkpoint stacks, depth state, typed-pool objects, and live slots + unchanged. + +Illustrative core: + +```julia +function _trim_inactive_typed_pool!(tp) + keep = tp.n_active + old_len = length(tp.vectors) + + released = old_len - keep + bytes = _inactive_storage_bytes(tp, keep + 1, old_len) + + resize!(tp.vectors, keep) + _trim_inactive_wrappers!(tp, keep) + + return (; slots_released = released, + wrappers_released = ..., + estimated_bytes_released = bytes) +end +``` + +No replacement vectors are allocated. If a later workload needs more slots, the +existing `acquire!` cold path will allocate the needed backing buffers at the +actual requested size. + +This avoids speculative "right-sizing" or reserve heuristics. A cleanup API +should not guess future sizes or allocate warm reserves on the user's behalf. + +## 5. Relationship to Existing Operations + +| Function | `n_active` | Backing buffers | Wrapper caches | Depth/checkpoint state | +|---|---:|---|---|---| +| `rewind!` | restored by one scope | keep all | keep all | pop one scope | +| `reset!` | set to 0 | keep all | keep all | reset to sentinel | +| `empty!` | set to 0 | drop all | drop all | reset to sentinel | +| `trim!` | unchanged | drop inactive only | drop inactive only | unchanged | + +Important distinction: + +- `trim!` is not `reset!`. If slots are still active, `trim!` preserves them. +- If user code has allowed `n_active` to grow at top level, call `reset!(pool)` + first if the intent is to mark those slots inactive, then call `trim!(pool)`. +- At normal top level after `@with_pool` scopes have rewound, `n_active` is + usually `0`, so `trim!(pool)` releases all retained backing buffers while + keeping the pool object and lightweight structural state. + +## 6. Why Not Allocate Smaller Replacement Buffers? + +An earlier idea was to replace a huge inactive buffer with a smaller new buffer +such as `1.5x` or `2x` the expected normal size. + +That is not part of this design. + +Reasons: + +- It performs a new allocation during a memory-reclamation call. +- The "normal future size" is only a guess without hot-path usage tracking. +- If the user actually needs that slot again, the existing `acquire!` path will + allocate the right size when it has the real request. +- The simplest and most honest behavior is: detach inactive references now, + allocate later only if needed. + +## 7. Implementation Notes + +### 7.1 Backing Buffers + +CPU typed pools can detach inactive backing buffers with: + +```julia +resize!(tp.vectors, tp.n_active) +``` + +This removes references to the inactive `Vector`s. The old vectors become +GC-eligible once no other references exist. + +For `BitTypedPool`, the same outer-vector truncation works; the dropped +`BitVector`s own their chunk arrays. + +For Metal, the same structural operation removes Julia references to inactive +`MtlArray`s. Actual VRAM return still depends on Julia GC and Metal.jl's memory +pool behavior, matching the caveat already documented for `empty!`. + +### 7.2 Wrapper Caches + +Dropping `tp.vectors` is not enough when cached wrappers still reference the old +backing storage. + +For each cached wrapper vector in `tp.arr_wrappers`, truncate it to `tp.n_active` +or clear inactive entries before truncation: + +```julia +for wrappers in tp.arr_wrappers + wrappers === nothing && continue + old_len = length(wrappers) + released += max(0, old_len - tp.n_active) + resize!(wrappers, min(length(wrappers), tp.n_active)) +end +``` + +The wrapper cache remains structurally reusable. Future `acquire!` calls past +the current active region will take the existing cold path and create fresh +wrappers if needed. + +### 7.3 Byte Accounting + +`estimated_bytes_released` is an estimate of storage that the pool stopped +retaining, not proof that RSS/VRAM immediately dropped. + +Suggested helpers: + +- CPU `Vector`: `Base.summarysize(v)` because `resize!(v, 0)` can leave capacity + retained even when `length(v) == 0`. +- `BitVector`: `Base.summarysize(v)` for the same capacity reason. +- Metal `MtlArray`: use GPU allocation size (`sizeof(v)` or `maxsize`-based + helper consistent with existing Metal stats), plus CPU wrapper size only if + useful for diagnostics. + +Avoid `length(v) * sizeof(T)` for released slots in safety mode: those vectors +may have logical length `0` while still retaining large capacity. + +## 8. Safety Contract + +`trim!` preserves active slots and their wrappers. Arrays currently in use by +the live call stack remain valid. + +Inactive retained slots are different: they have already been released by pool +scope management. User code should not rely on arrays/views from those slots. +After `trim!`, any pool-owned references to those inactive buffers are gone. + +In runtime-check mode (`S >= 1`), released slots have already been invalidated +by `rewind!`/`reset!` where applicable. `trim!` does not need to add new hot-path +safety metadata. + +`trim!` should be idempotent: + +- Calling it twice in a row releases nothing on the second call. +- Calling it inside a nested computation keeps currently active slots. +- Calling it at top level after normal rewind releases retained buffers. + +## 9. Backend Structure + +Generic core in `src/state.jl`: + +- `_trim_inactive_typed_pool!(tp::AbstractTypedPool)` +- `_trim_inactive_wrappers!(tp, keep::Int)` +- `_inactive_storage_bytes(tp, first::Int, last::Int)` + +Per-pool summaries are accumulated inline in the `trim!(pool)` loops over the +fixed slots and `others` (no separate merge helper is needed). + +CPU pool entry points: + +- `trim!(pool::AdaptiveArrayPool; force_gc=false)` +- `trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc=false)` +- `trim!(; force_gc=false)` for `get_task_local_pool()` + +Metal extension entry points: + +- Thin methods in `ext/AdaptiveArrayPoolsMetalExt/state.jl` +- Reuse generic typed-pool helpers where field layout matches. +- Add Metal-specific storage-byte helper if needed. + +CUDA: + +- Deferred until CPU + Metal behavior is proven. +- Expected to mirror Metal structurally. + +Legacy (Julia < 1.12): + +- The legacy pool (`src/legacy/`) uses an N-way set-associative cache with a + different field layout (`nd_arrays`/`nd_dims`/… instead of `arr_wrappers`), so + real reclamation is not implemented there. +- `trim!` is still **defined and exported** on the legacy path as a NO-OP that + returns a zero summary and warns once (`@warn ... maxlog=1`). This keeps the + public API callable across the full `[compat] julia = "1.10"` range, so + dependent packages can `import`/call `trim!` on any supported Julia version + without a load-time or runtime `UndefVarError`. Real reclamation requires + Julia 1.12+. (Adding a real legacy implementation later would be non-breaking.) + +## 10. Phase Plan (TDD) + +Execute phases one at a time. + +### Phase 1 — CPU `trim!` + +**Goal:** Implement `trim!` for CPU pools and typed pools. + +**Test strategy:** + +- Active slots are preserved. +- Inactive `tp.vectors` entries are removed. +- Inactive wrapper-cache entries are removed. +- `reset!(pool); trim!(pool)` drops all retained CPU buffers. +- `trim!(pool, T)` affects only that type. +- `DisabledPool` returns a zero summary. +- `force_gc=true` sets `gc_triggered=true`. + +**Tasks:** + +- [ ] RED: Add CPU tests in `test/test_state.jl` or a new targeted test file. +- [ ] GREEN: Add generic typed-pool trimming helpers and CPU API methods. +- [ ] REFACTOR: Share summary aggregation helpers and docstrings. + +**Quality gate:** + +- [ ] Targeted CPU state tests pass. +- [ ] Existing allocation/zero-allocation tests for normal acquire paths still pass. +- [ ] `trim!` does not add code to hot acquire or automatic rewind paths. + +**Rollback:** Revert `src/state.jl`, export changes, and new tests. + +### Phase 2 — Safety Mode + Wrapper Reference Verification + +**Goal:** Prove trimming actually detaches pool-held references, including +cached wrappers. + +**Test strategy:** + +- Create cached `Array`/`BitArray` wrappers for inactive slots, call `trim!`, + and assert wrapper cache length is truncated. +- In `S=1`, released slots that were invalidated to length zero are still + counted using capacity-aware byte accounting. +- Subsequent `acquire!` after trim self-heals by allocating fresh backing storage + and wrapper entries. + +**Tasks:** + +- [ ] RED: Add tests for wrapper-cache truncation and S=1 byte estimates. +- [ ] GREEN: Fix byte accounting and wrapper-cache cleanup details. +- [ ] REFACTOR: Keep helper names and docstrings aligned with "inactive slots". + +**Quality gate:** + +- [ ] Targeted CPU safety tests pass. +- [ ] No stale wrapper retains old backing storage through the pool cache. + +**Rollback:** Revert helper changes and tests from this phase. + +### Phase 3 — Metal Parity + +**Goal:** Add the same inactive-slot trimming API to Metal pools. + +**Test strategy:** + +- Active Metal slots are preserved. +- Inactive Metal backing `MtlArray` references are removed. +- Inactive Metal wrapper-cache entries are removed. +- `force_gc=true` calls Julia GC after references are detached. +- Return summary uses Metal-appropriate storage estimates. + +**Tasks:** + +- [ ] RED: Add Metal trim tests under `test/metal/`. +- [ ] GREEN: Add thin Metal pool entry points and byte helper overrides. +- [ ] REFACTOR: Keep CPU/Metal implementation split consistent with `reset!` and `empty!`. + +**Quality gate:** + +- [ ] Targeted Metal tests pass on the M1 machine. +- [ ] CPU tests from Phases 1-2 remain green. + +**Rollback:** Revert Metal extension changes and Metal tests. + +### Phase 4 — Public Docs + +**Goal:** Document `trim!` as the manual inactive-slot cleanup operation. + +**Test strategy:** + +- Doctest or docs build if available. +- API reference includes `trim!`. +- User docs clearly distinguish `trim!`, `reset!`, and `empty!`. + +**Tasks:** + +- [ ] RED: Add docs references that fail until API exists. +- [ ] GREEN: Update API docs, examples, and exported function list. +- [ ] REFACTOR: Tighten wording around GC eligibility vs immediate memory return. + +**Quality gate:** + +- [ ] Docs build succeeds. +- [ ] No public docs imply immediate OS/VRAM release. + +**Rollback:** Revert docs and export/reference updates. + +## 11. Deferred Ideas + +These are intentionally out of scope for the first implementation: + +- Outlier-only trimming (`policy=:largest`, `min_bytes`, `ratio`). +- Dry-run reporting. +- Per-type or per-backend default thresholds. +- Hysteresis or usage tracking. +- Automatic periodic trimming. +- Replacement/reserve buffers for "right-sized" warm reuse. + +All of these add policy. The first version should be a simple primitive: + +> Drop inactive retained references; preserve active working set. + +## 12. Risks + +| Risk | Impact | Likelihood | Mitigation | Detection | +|---|---:|---:|---|---| +| Cached wrappers keep old buffers alive | High | Medium | Truncate wrapper caches with `tp.vectors` | Wrapper-cache reference tests | +| Byte summary reports zero for resized inactive buffers | Medium | High in `S=1` | Use capacity-aware estimates (`summarysize`, Metal allocation size) | S=1 byte-accounting tests | +| Users expect immediate RSS/VRAM drop | Medium | Medium | Document GC eligibility and `force_gc` limits | Docs review | +| `trim!` confused with `reset!` | Medium | Medium | Keep `n_active` unchanged and document `reset!; trim!` pattern | API docs/tests | +| Metal memory pool behavior differs from CPU | Medium | Medium | Match `empty!` caveats; test reference removal, not immediate VRAM | Metal tests | + +## 13. Open Items + +- Final summary field names: `slots_released` vs `inactive_slots_released`, + `estimated_bytes_released` vs `bytes_detached`. +- Whether zero-argument `trim!()` should stay CPU-only or eventually accept a + backend selector. +- Whether `force_gc=true` should call `GC.gc()` once or use a stronger collection + pattern. Default remains `false`. diff --git a/docs/src/basics/api-essentials.md b/docs/src/basics/api-essentials.md index 513f1cd9..c52015a1 100644 --- a/docs/src/basics/api-essentials.md +++ b/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 **inactive** retained buffers (those already +freed by `rewind!`/scope exit), while preserving the **active** arrays still in +use. This is the operation for reclaiming memory after a rare large workload +without discarding the warm, normal-sized buffers that keep the hot path +allocation-free. + +```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 diff --git a/docs/src/reference/api.md b/docs/src/reference/api.md index 7829832f..8080c00f 100644 --- a/docs/src/reference/api.md +++ b/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)`. 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 diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 0d8dffb2..371df3d1 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -334,3 +334,53 @@ function Base.empty!(pool::MetalAdaptiveArrayPool) return pool end + +# ============================================================================== +# trim! for MetalAdaptiveArrayPool (parity with CPU) +# ============================================================================== + +# Metal backing storage lives on the GPU. `sizeof(MtlArray)` = length * sizeof(T) +# = the real GPU buffer size, whereas `Base.summarysize` (the generic CPU helper) +# only sees the small CPU-side handle. Override so the byte estimate is correct. +# (Metal inactive slots are not shrunk to length 0, so length == capacity here.) +function AdaptiveArrayPools._inactive_storage_bytes(tp::MetalTypedPool, first::Int, last::Int) + total = 0 + for i in first:last + @inbounds total += sizeof(tp.vectors[i]) + end + return total +end + +function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool; force_gc::Bool = false) + slots = 0 + wrappers = 0 + bytes = 0 + AdaptiveArrayPools.foreach_fixed_slot(pool) do tp + s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) + slots += s.slots_released + wrappers += s.wrappers_released + bytes += s.estimated_bytes_released + end + for tp in values(pool.others) + s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) + slots += s.slots_released + wrappers += s.wrappers_released + bytes += s.estimated_bytes_released + end + force_gc && GC.gc() + return (; + slots_released = slots, wrappers_released = wrappers, + estimated_bytes_released = bytes, gc_triggered = force_gc, + ) +end + +@inline function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} + s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) + force_gc && GC.gc() + return (; + slots_released = s.slots_released, + wrappers_released = s.wrappers_released, + estimated_bytes_released = s.estimated_bytes_released, + gc_triggered = force_gc, + ) +end diff --git a/src/AdaptiveArrayPools.jl b/src/AdaptiveArrayPools.jl index aa89cd86..959d241e 100644 --- a/src/AdaptiveArrayPools.jl +++ b/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! +export checkpoint!, rewind!, reset!, trim! 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) diff --git a/src/legacy/state.jl b/src/legacy/state.jl index 9a92054f..bad2e26e 100644 --- a/src/legacy/state.jl +++ b/src/legacy/state.jl @@ -753,3 +753,36 @@ reset!(::DisabledPool, ::Type) = nothing reset!(::DisabledPool, types::Type...) = nothing Base.empty!(::DisabledPool) = nothing + +# ============================================================================== +# trim! — no-op on the legacy path (Julia < 1.12) +# ============================================================================== +# +# Modern `trim!` reclaims inactive slots via the `arr_wrappers` setfield! cache, +# which only exists on Julia 1.12+. The legacy pool uses an N-way set-associative +# cache (`nd_arrays`/`nd_dims`/…) with a different layout, so `trim!` is provided +# here as a defined-and-exported NO-OP (returns a zero summary, warns once). +# +# This keeps the public API callable across the full `[compat] julia` range, so +# dependent packages can `import`/call `trim!` on any supported Julia version +# without load-time or runtime `UndefVarError`. Upgrade to Julia 1.12+ for actual +# inactive-slot reclamation. + +const _ZERO_TRIM_SUMMARY = (; + slots_released = 0, wrappers_released = 0, + estimated_bytes_released = 0, gc_triggered = false, +) + +function trim!(pool::AdaptiveArrayPool; force_gc::Bool = false) + @warn "trim! is a no-op on Julia < 1.12 (legacy pool architecture). " * + "Upgrade to Julia 1.12+ for inactive-slot reclamation." maxlog = 1 + return _ZERO_TRIM_SUMMARY +end + +trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} = + trim!(pool; force_gc = force_gc) + +trim!(; force_gc::Bool = false) = trim!(get_task_local_pool(); force_gc = force_gc) + +trim!(::DisabledPool; force_gc::Bool = false) = _ZERO_TRIM_SUMMARY +trim!(::DisabledPool, ::Type{T}; force_gc::Bool = false) where {T} = _ZERO_TRIM_SUMMARY diff --git a/src/state.jl b/src/state.jl index 8d85fbbe..6571375b 100644 --- a/src/state.jl +++ b/src/state.jl @@ -728,6 +728,148 @@ in `touched_mask` is also set in `tracked_mask`. return (touched_mask & ~tracked_mask) == UInt16(0) && !has_others end +# ============================================================================== +# State Management - trim! (manual inactive-slot trimming) +# ============================================================================== +# +# trim! drops the pool's strong references to INACTIVE retained slots +# (indices n_active+1 : end of each typed pool) while preserving ACTIVE slots +# (1 : n_active). This releases backing buffers (and cached wrappers) so they +# become GC-eligible once no user code still references them. It does NOT promise +# immediate OS/VRAM return. +# +# trim! is a manual operation. It adds no code to the hot acquire path or the +# automatic rewind path. + +# Zero summary: returned by no-op trims (DisabledPool, nothing to release). +const _ZERO_TRIM_SUMMARY = (; + slots_released = 0, wrappers_released = 0, + estimated_bytes_released = 0, gc_triggered = false, +) + +# Drop cached N-D wrappers for inactive slots. A cached Array/BitArray holds a +# `:ref` into its slot's backing buffer, so leaving inactive wrappers in place +# would pin the very buffers we just detached. Each per-slot wrapper vector +# (indexed by dimensionality N) is truncated to `keep`. Returns the count of +# non-nothing wrapper objects dropped. +function _trim_inactive_wrappers!(tp::AbstractTypedPool, keep::Int) + released = 0 + for wrappers in tp.arr_wrappers + wrappers === nothing && continue + old_len = length(wrappers) + old_len <= keep && continue + for i in (keep + 1):old_len + @inbounds wrappers[i] === nothing || (released += 1) + end + resize!(wrappers, keep) + end + return released +end + +# Estimate the storage held by inactive backing vectors in slots `first:last`. +# Uses `Base.summarysize`, which is capacity-aware: it sizes the backing `Memory`, +# so it still reports the retained capacity for vectors that `rewind!`/`reset!` +# shrank to logical length 0 in runtime-check mode (S >= 1). This is an estimate, +# not a guarantee that RSS dropped. Dispatch point: extensions (e.g. Metal) may +# override for backend-appropriate (GPU) storage sizes. +function _inactive_storage_bytes(tp::AbstractTypedPool, first::Int, last::Int) + total = 0 + for i in first:last + @inbounds total += Base.summarysize(tp.vectors[i]) + end + return total +end + +# Trim one typed pool: drop inactive backing-vector and wrapper references by +# truncating `tp.vectors` and the wrapper caches to the active count. Returns a +# per-pool summary NamedTuple. Bytes are measured before truncation. +function _trim_inactive_typed_pool!(tp::AbstractTypedPool) + keep = tp.n_active + old_len = length(tp.vectors) + released = old_len - keep + wrappers_released = _trim_inactive_wrappers!(tp, keep) + bytes = 0 + if released > 0 + bytes = _inactive_storage_bytes(tp, keep + 1, old_len) + resize!(tp.vectors, keep) + end + return (; + slots_released = max(0, released), + wrappers_released = wrappers_released, + estimated_bytes_released = bytes, + ) +end + +""" + trim!(pool::AdaptiveArrayPool; force_gc::Bool = false) + +Release the pool's references to **inactive retained slots** (those already +released by `rewind!`/scope exit), keeping all **active** slots intact. + +Unlike `reset!` (keeps all buffers) and `empty!` (drops all buffers and resets +state), `trim!` drops only the inactive backing buffers and cached wrappers and +leaves `n_active`, checkpoint stacks, and depth state unchanged. + +Returns `(; slots_released, wrappers_released, estimated_bytes_released, gc_triggered)`. + +`force_gc=true` calls `GC.gc()` after detaching references. This still does not +guarantee immediate OS/VRAM return; it only asks Julia to collect now that the +pool no longer holds the buffers. + +!!! note "Julia version" + Actual reclamation requires Julia 1.12+. On older Julia (the legacy pool + architecture) `trim!` is a defined no-op: it returns a zero summary and warns + once. It remains callable so dependent packages compile across the full + supported Julia range. + +See also: [`reset!`](@ref), [`empty!`](@ref). +""" +function trim!(pool::AdaptiveArrayPool; force_gc::Bool = false) + slots = 0 + wrappers = 0 + bytes = 0 + foreach_fixed_slot(pool) do tp + s = _trim_inactive_typed_pool!(tp) + slots += s.slots_released + wrappers += s.wrappers_released + bytes += s.estimated_bytes_released + end + for tp in pool._others_values + s = _trim_inactive_typed_pool!(tp) + slots += s.slots_released + wrappers += s.wrappers_released + bytes += s.estimated_bytes_released + end + force_gc && GC.gc() + return (; + slots_released = slots, wrappers_released = wrappers, + estimated_bytes_released = bytes, gc_triggered = force_gc, + ) +end + +""" + trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) + +Trim inactive slots for a single element type `T` only. See [`trim!`](@ref). +""" +function trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} + s = _trim_inactive_typed_pool!(get_typed_pool!(pool, T)) + force_gc && GC.gc() + return (; + slots_released = s.slots_released, + wrappers_released = s.wrappers_released, + estimated_bytes_released = s.estimated_bytes_released, + gc_triggered = force_gc, + ) +end + +""" + trim!(; force_gc::Bool = false) + +Trim inactive slots of the current task-local pool (`get_task_local_pool()`). +""" +trim!(; force_gc::Bool = false) = trim!(get_task_local_pool(); force_gc = force_gc) + # ============================================================================== # DisabledPool State Management (no-ops) # ============================================================================== @@ -746,3 +888,6 @@ reset!(::DisabledPool, ::Type) = nothing reset!(::DisabledPool, types::Type...) = nothing Base.empty!(::DisabledPool) = nothing + +trim!(::DisabledPool; force_gc::Bool = false) = _ZERO_TRIM_SUMMARY +trim!(::DisabledPool, ::Type{T}; force_gc::Bool = false) where {T} = _ZERO_TRIM_SUMMARY diff --git a/test/metal/runtests.jl b/test/metal/runtests.jl index 006c523e..70c8ba1a 100644 --- a/test/metal/runtests.jl +++ b/test/metal/runtests.jl @@ -51,5 +51,6 @@ else include("test_scope_depth_validation.jl") include("test_reshape.jl") include("test_task_local_pool.jl") + include("test_trim.jl") end end diff --git a/test/metal/test_trim.jl b/test/metal/test_trim.jl new file mode 100644 index 00000000..17bc855a --- /dev/null +++ b/test/metal/test_trim.jl @@ -0,0 +1,109 @@ +# Metal trim! Tests +# Parity with CPU test_trim.jl: drop inactive retained slots, keep active. +# Uses Float32 (Metal's primary type; no Float64 on Metal hardware). + +@testset "Metal trim!" begin + + @testset "keeps active slots, drops deeper inactive (Metal, S=0)" begin + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + + checkpoint!(pool) # depth 2 + acquire!(pool, Float32, 100) # slot 1 (stays active) + acquire!(pool, Float32, 50) # slot 2 (stays active) + + checkpoint!(pool) # depth 3 + acquire!(pool, Float32, 25) # slot 3 + acquire!(pool, Float32, 10) # slot 4 + rewind!(pool) # back to depth 2: n_active → 2 + + @test pool.float32.n_active == 2 + @test length(pool.float32.vectors) == 4 + + s = trim!(pool) + + @test pool.float32.n_active == 2 + @test length(pool.float32.vectors) == 2 + @test s.slots_released == 2 + @test s.gc_triggered == false + end + + @testset "active slots untouched (Metal, identity + data)" begin + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + + checkpoint!(pool) # depth 2 — active + a = acquire!(pool, Float32, 8); fill!(a, 7.0f0) + b = acquire!(pool, Float32, 4); fill!(b, 100.0f0) + + checkpoint!(pool) # depth 3 — become inactive on rewind + acquire!(pool, Float32, 1000) + acquire!(pool, Float32, 2000) + rewind!(pool) + + v1 = pool.float32.vectors[1] + v2 = pool.float32.vectors[2] + + s = trim!(pool) + + @test s.slots_released == 2 + @test pool.float32.n_active == 2 + @test length(pool.float32.vectors) == 2 + @test pool.float32.vectors[1] === v1 # same GPU buffers, untouched + @test pool.float32.vectors[2] === v2 + @test all(==(7.0f0), Array(a)) # data intact (device → host round-trip) + @test all(==(100.0f0), Array(b)) + end + + @testset "byte estimate uses GPU size (sizeof), not summarysize" begin + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + + checkpoint!(pool) + acquire!(pool, Float32, 10_000) # 40_000 GPU bytes + rewind!(pool) + + s = trim!(pool) + @test s.slots_released == 1 + # summarysize(MtlArray) reports only the ~80-byte CPU handle; sizeof gives + # the real 40_000 GPU bytes. This asserts the Metal byte override is used. + @test s.estimated_bytes_released >= 40_000 + end + + @testset "wrapper truncation + per-type + force_gc (Metal)" begin + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + + checkpoint!(pool) + acquire!(pool, Float32, 4, 5) # 2-D → arr_wrappers[2][1] + acquire!(pool, Int32, 100) + rewind!(pool) + + @test pool.float32.arr_wrappers[2] !== nothing + @test length(pool.float32.arr_wrappers[2]) == 1 + + s = trim!(pool, Float32) # per-type: only Float32 + @test length(pool.float32.vectors) == 0 + @test length(pool.int32.vectors) == 1 # untouched + @test s.slots_released == 1 + for w in pool.float32.arr_wrappers + w === nothing && continue + @test length(w) <= pool.float32.n_active + end + + @test trim!(pool; force_gc = true).gc_triggered == true + end + + @testset "self-heal: acquire! after trim! (Metal)" begin + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + + checkpoint!(pool) + acquire!(pool, Float32, 1000) + rewind!(pool) + trim!(pool) + @test length(pool.float32.vectors) == 0 + + checkpoint!(pool) + v = acquire!(pool, Float32, 7) + @test length(v) == 7 + @test length(pool.float32.vectors) == 1 + rewind!(pool) + end + +end diff --git a/test/runtests.jl b/test/runtests.jl index 8c418c96..0e3bb653 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,6 +22,7 @@ else include("test_aqua.jl") include("test_basic.jl") include("test_state.jl") + include("test_trim.jl") include("test_multidimensional.jl") include("test_macros.jl") include("test_task_local_pool.jl") @@ -52,6 +53,7 @@ else include("test_aqua.jl") include("test_basic.jl") include("test_state.jl") + include("test_trim_legacy.jl") include("test_multidimensional.jl") include("test_macros.jl") include("test_task_local_pool.jl") diff --git a/test/test_trim.jl b/test/test_trim.jl new file mode 100644 index 00000000..d8919e37 --- /dev/null +++ b/test/test_trim.jl @@ -0,0 +1,227 @@ +using Test +using AdaptiveArrayPools +using AdaptiveArrayPools: get_typed_pool!, trim!, checkpoint!, rewind! + +# ============================================================================== +# trim! — manual inactive-slot trimming +# +# Concept: trim! drops the pool's strong references to INACTIVE retained slots +# (indices n_active+1 : end) while preserving ACTIVE slots (1 : n_active). +# ============================================================================== + +@testset "trim!" begin + + @testset "keeps active slots, drops deeper inactive (CPU, S=0)" begin + pool = AdaptiveArrayPool{0}() # S=0: retained vectors keep their length + + checkpoint!(pool) # depth 2 + acquire!(pool, Float64, 100) # slot 1 (stays active) + acquire!(pool, Float64, 50) # slot 2 (stays active) + + checkpoint!(pool) # depth 3 + acquire!(pool, Float64, 25) # slot 3 + acquire!(pool, Float64, 10) # slot 4 + rewind!(pool) # back to depth 2: n_active → 2, slots 3,4 retained + + @test pool.float64.n_active == 2 + @test length(pool.float64.vectors) == 4 # inactive slots still retained + + summary = trim!(pool) + + # Active slots preserved; inactive backing vectors dropped + @test pool.float64.n_active == 2 + @test length(pool.float64.vectors) == 2 + @test summary.slots_released == 2 + @test summary.gc_triggered == false + end + + @testset "active slots are byte-for-byte untouched (data + identity)" begin + pool = AdaptiveArrayPool{0}() + + checkpoint!(pool) # depth 2 — these stay ACTIVE + a = acquire!(pool, Float64, 8) + a .= 1.0:8.0 # known data in active slot 1 + b = acquire!(pool, Float64, 4) + b .= 100.0 # known data in active slot 2 + + checkpoint!(pool) # depth 3 — become INACTIVE on rewind + acquire!(pool, Float64, 1000) + acquire!(pool, Float64, 2000) + rewind!(pool) # n_active → 2; slots 3,4 retained (inactive) + + # Identity of the ACTIVE backing vectors before trimming. + v1 = pool.float64.vectors[1] + v2 = pool.float64.vectors[2] + + summary = trim!(pool) + + # Inactive dropped; active count intact. + @test summary.slots_released == 2 + @test pool.float64.n_active == 2 + @test length(pool.float64.vectors) == 2 + + # ACTIVE backing vectors are the SAME objects — not replaced or resized. + @test pool.float64.vectors[1] === v1 + @test pool.float64.vectors[2] === v2 + + # ACTIVE arrays' data is byte-for-byte intact and they stay usable. + @test a == collect(1.0:8.0) + @test all(==(100.0), b) + a[1] = 42.0 # still a live, writable view into the pool + @test pool.float64.vectors[1][1] == 42.0 + end + + @testset "drops inactive wrapper-cache entries (memory-pin fix)" begin + pool = AdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Float64, 4, 5) # 2-D → arr_wrappers[2][1] + acquire!(pool, Float64, 3, 3) # 2-D → arr_wrappers[2][2] + rewind!(pool) # n_active → 0; slots + wrappers retained + + tp = pool.float64 + @test tp.n_active == 0 + @test length(tp.arr_wrappers) >= 2 + @test tp.arr_wrappers[2] !== nothing + @test length(tp.arr_wrappers[2]) == 2 # two cached wrappers retained + + summary = trim!(pool) + + # Every wrapper-cache vector truncated to the active count (0 here): + # a cached Array holds a :ref into the old buffer, pinning it in memory. + for w in tp.arr_wrappers + w === nothing && continue + @test length(w) <= tp.n_active + end + @test summary.wrappers_released >= 2 + end + + @testset "per-type trim!(pool, T) affects only that type" begin + pool = AdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Float64, 100) + acquire!(pool, Int64, 200) + rewind!(pool) # both retained, n_active == 0 each + + @test length(pool.float64.vectors) == 1 + @test length(pool.int64.vectors) == 1 + + summary = trim!(pool, Float64) + + @test length(pool.float64.vectors) == 0 # trimmed + @test length(pool.int64.vectors) == 1 # untouched + @test summary.slots_released == 1 + end + + @testset "DisabledPool returns a zero summary" begin + s = trim!(DISABLED_CPU) + @test s.slots_released == 0 + @test s.wrappers_released == 0 + @test s.estimated_bytes_released == 0 + @test s.gc_triggered == false + + s2 = trim!(DISABLED_CPU, Float64) + @test s2.slots_released == 0 + end + + @testset "zero-arg trim!() operates on the task-local pool" begin + pool = get_task_local_pool() + empty!(pool) # clean slate (task-local is shared) + + checkpoint!(pool) + acquire!(pool, Float64, 100) + rewind!(pool) + @test length(pool.float64.vectors) == 1 + + summary = trim!() # no-arg → task-local pool + @test length(pool.float64.vectors) == 0 + @test summary.slots_released >= 1 + + empty!(pool) # cleanup to avoid cross-test contamination + end + + @testset "force_gc=true sets gc_triggered (regression)" begin + pool = AdaptiveArrayPool{0}() + @test trim!(pool; force_gc = true).gc_triggered == true + @test trim!(pool; force_gc = false).gc_triggered == false + end + + @testset "reset!(pool); trim!(pool) releases all retained buffers (regression)" begin + pool = AdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Float64, 100) + acquire!(pool, Float64, 50) + # (no rewind: simulate top-level n_active growth) + @test pool.float64.n_active == 2 + + reset!(pool) # mark all slots inactive (n_active → 0) + @test pool.float64.n_active == 0 + @test length(pool.float64.vectors) == 2 # reset! preserves buffers + + trim!(pool) # now release them + @test length(pool.float64.vectors) == 0 + end + + @testset "byte accounting" begin + @testset "S=0: estimate reflects retained buffer sizes" begin + pool = AdaptiveArrayPool{0}() + checkpoint!(pool) + acquire!(pool, Float64, 10_000) # 80_000 bytes of backing storage + rewind!(pool) + s = trim!(pool) + @test s.slots_released == 1 + @test s.estimated_bytes_released >= 80_000 + end + + @testset "S=1: capacity-aware estimate counts buffers shrunk to length 0" begin + pool = AdaptiveArrayPool{1}() + checkpoint!(pool) + acquire!(pool, Float64, 10_000) + rewind!(pool) # S=1: inactive slot resize!'d to length 0 + @test length(pool.float64.vectors[1]) == 0 # logical length 0... + s = trim!(pool) + @test s.slots_released == 1 + @test s.estimated_bytes_released >= 80_000 # ...but retained capacity counted + end + end + + @testset "self-heal: acquire! after trim! allocates fresh and works" begin + pool = AdaptiveArrayPool{0}() + checkpoint!(pool) + acquire!(pool, Float64, 1000) + rewind!(pool) + trim!(pool) + @test length(pool.float64.vectors) == 0 + + checkpoint!(pool) + v = acquire!(pool, Float64, 7) + v .= 3.0 + @test length(v) == 7 + @test all(v .== 3.0) + @test length(pool.float64.vectors) == 1 + rewind!(pool) + end + + @testset "Bit (BitVector) and fallback (others) types trim too" begin + pool = AdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Bit, 10_000) # BitTypedPool path + acquire!(pool, UInt8, 5_000) # fallback (others) path + rewind!(pool) + + @test length(pool.bits.vectors) == 1 + @test pool.others[UInt8].n_active == 0 + @test length(pool.others[UInt8].vectors) == 1 + + s = trim!(pool) + + @test length(pool.bits.vectors) == 0 + @test length(pool.others[UInt8].vectors) == 0 + @test s.slots_released == 2 + @test s.estimated_bytes_released > 0 + end + +end diff --git a/test/test_trim_legacy.jl b/test/test_trim_legacy.jl new file mode 100644 index 00000000..d30c2038 --- /dev/null +++ b/test/test_trim_legacy.jl @@ -0,0 +1,25 @@ +using Test +using AdaptiveArrayPools +using AdaptiveArrayPools: trim!, get_task_local_pool + +# Legacy path (Julia < 1.12): trim! is a defined, exported NO-OP that returns a +# zero summary and warns once. This keeps the public API portable across the full +# supported Julia range (dependents can import/call trim! on any version). + +@testset "trim! (legacy no-op, Julia < 1.12)" begin + pool = get_task_local_pool() + + # First call warns once and returns a zero summary. + s = @test_logs (:warn,) match_mode = :any trim!(pool) + @test s.slots_released == 0 + @test s.wrappers_released == 0 + @test s.estimated_bytes_released == 0 + @test s.gc_triggered == false + + # Same zero-summary shape for the other entry points. + @test trim!(pool, Float64).slots_released == 0 + @test trim!().slots_released == 0 + + # Legacy no-op ignores force_gc (no reclamation to do). + @test trim!(pool; force_gc = true).gc_triggered == false +end From fe36f8da632b1077e00c2055d50a5036fc5d45d2 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 14:55:14 -0700 Subject: [PATCH 2/7] (feat): CUDA parity for `trim!` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror the Metal `trim!` into the CUDA extension: - `trim!(pool::CuAdaptiveArrayPool; force_gc=false)` + per-type variant, reusing the generic `_trim_inactive_typed_pool!` core (CuTypedPool shares the CPU/Metal field layout, so the helpers apply unchanged). - `_inactive_storage_bytes(::CuTypedPool, …)` override using `sizeof` for true GPU buffer bytes (summarysize only sees the CPU-side handle). - `test/cuda/test_trim.jl` mirrors the Metal suite (active-untouched data+identity, byte accounting, wrapper-cache truncation, per-type, force_gc, self-heal); runs on CI with a functional CUDA device, skipped otherwise. Validated locally by extension precompilation (dispatch resolves); GPU execution deferred to CI. Spec/API docs updated to reflect CPU + Metal + CUDA support. --- docs/design/selective_memory_reclamation.md | 11 +- docs/src/reference/api.md | 2 +- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 50 +++++++++ test/cuda/runtests.jl | 1 + test/cuda/test_trim.jl | 109 ++++++++++++++++++++ 5 files changed, 168 insertions(+), 5 deletions(-) create mode 100644 test/cuda/test_trim.jl diff --git a/docs/design/selective_memory_reclamation.md b/docs/design/selective_memory_reclamation.md index 96c9cecb..4a4eba59 100644 --- a/docs/design/selective_memory_reclamation.md +++ b/docs/design/selective_memory_reclamation.md @@ -1,8 +1,8 @@ # Inactive Slot Trimming for AdaptiveArrayPools - **Date:** 2026-06-28 -- **Status:** Revised design — pending user review (not yet implemented) -- **Scope:** New manual `trim!` API. CPU + Metal in first cut; CUDA deferred. +- **Status:** Implemented (CPU + Metal validated on hardware; CUDA precompile-validated, GPU execution via CI). +- **Scope:** New manual `trim!` API. CPU + Metal + CUDA (CUDA mirrors Metal). - **Author:** design session --- @@ -257,8 +257,11 @@ Metal extension entry points: CUDA: -- Deferred until CPU + Metal behavior is proven. -- Expected to mirror Metal structurally. +- Implemented as a structural mirror of Metal (`CuTypedPool` has the same field + layout; `_inactive_storage_bytes` overridden to use `sizeof` for GPU bytes). +- Validated locally by extension precompilation (method dispatch resolves); + GPU execution covered by `test/cuda/test_trim.jl`, which runs on CI with a + functional CUDA device (skipped on non-CUDA machines). Legacy (Julia < 1.12): diff --git a/docs/src/reference/api.md b/docs/src/reference/api.md index 8080c00f..d9b83c49 100644 --- a/docs/src/reference/api.md +++ b/docs/src/reference/api.md @@ -24,7 +24,7 @@ | `pool_stats(pool)` | Prints detailed statistics about pool usage. | | `get_task_local_pool()` | Returns the task-local pool instance. | | `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)`. Reclaims on Julia 1.12+; defined no-op (zero summary + one-time warning) on older Julia. | +| `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 diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index bba450d0..547c0c0a 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -352,3 +352,53 @@ function Base.empty!(pool::CuAdaptiveArrayPool) return pool end + +# ============================================================================== +# trim! for CuAdaptiveArrayPool (parity with CPU/Metal) +# ============================================================================== + +# CUDA backing storage lives on the GPU. `sizeof(CuVector)` = length * sizeof(T) +# = the real device buffer size, whereas `Base.summarysize` (the generic CPU +# helper) only sees the small CPU-side handle. Override so the byte estimate is +# correct. (CUDA inactive slots are not shrunk to length 0, so length == capacity.) +function AdaptiveArrayPools._inactive_storage_bytes(tp::CuTypedPool, first::Int, last::Int) + total = 0 + for i in first:last + @inbounds total += sizeof(tp.vectors[i]) + end + return total +end + +function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool; force_gc::Bool = false) + slots = 0 + wrappers = 0 + bytes = 0 + AdaptiveArrayPools.foreach_fixed_slot(pool) do tp + s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) + slots += s.slots_released + wrappers += s.wrappers_released + bytes += s.estimated_bytes_released + end + for tp in values(pool.others) + s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) + slots += s.slots_released + wrappers += s.wrappers_released + bytes += s.estimated_bytes_released + end + force_gc && GC.gc() + return (; + slots_released = slots, wrappers_released = wrappers, + estimated_bytes_released = bytes, gc_triggered = force_gc, + ) +end + +@inline function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} + s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) + force_gc && GC.gc() + return (; + slots_released = s.slots_released, + wrappers_released = s.wrappers_released, + estimated_bytes_released = s.estimated_bytes_released, + gc_triggered = force_gc, + ) +end diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl index c5974fbb..df466279 100644 --- a/test/cuda/runtests.jl +++ b/test/cuda/runtests.jl @@ -49,5 +49,6 @@ else include("test_cuda_safety.jl") include("test_runtime_mutation.jl") include("test_scope_depth_validation.jl") + include("test_trim.jl") end end diff --git a/test/cuda/test_trim.jl b/test/cuda/test_trim.jl new file mode 100644 index 00000000..2b44a788 --- /dev/null +++ b/test/cuda/test_trim.jl @@ -0,0 +1,109 @@ +# CUDA trim! Tests +# Parity with CPU/Metal test_trim.jl: drop inactive retained slots, keep active. +# Uses Float32 (universal GPU type). + +@testset "CUDA trim!" begin + + @testset "keeps active slots, drops deeper inactive (CUDA, S=0)" begin + pool = CuAdaptiveArrayPool{0}() + + checkpoint!(pool) # depth 2 + acquire!(pool, Float32, 100) # slot 1 (stays active) + acquire!(pool, Float32, 50) # slot 2 (stays active) + + checkpoint!(pool) # depth 3 + acquire!(pool, Float32, 25) # slot 3 + acquire!(pool, Float32, 10) # slot 4 + rewind!(pool) # back to depth 2: n_active → 2 + + @test pool.float32.n_active == 2 + @test length(pool.float32.vectors) == 4 + + s = trim!(pool) + + @test pool.float32.n_active == 2 + @test length(pool.float32.vectors) == 2 + @test s.slots_released == 2 + @test s.gc_triggered == false + end + + @testset "active slots untouched (CUDA, identity + data)" begin + pool = CuAdaptiveArrayPool{0}() + + checkpoint!(pool) # depth 2 — active + a = acquire!(pool, Float32, 8); fill!(a, 7.0f0) + b = acquire!(pool, Float32, 4); fill!(b, 100.0f0) + + checkpoint!(pool) # depth 3 — become inactive on rewind + acquire!(pool, Float32, 1000) + acquire!(pool, Float32, 2000) + rewind!(pool) + + v1 = pool.float32.vectors[1] + v2 = pool.float32.vectors[2] + + s = trim!(pool) + + @test s.slots_released == 2 + @test pool.float32.n_active == 2 + @test length(pool.float32.vectors) == 2 + @test pool.float32.vectors[1] === v1 # same GPU buffers, untouched + @test pool.float32.vectors[2] === v2 + @test all(==(7.0f0), Array(a)) # data intact (device → host round-trip) + @test all(==(100.0f0), Array(b)) + end + + @testset "byte estimate uses GPU size (sizeof), not summarysize" begin + pool = CuAdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Float32, 10_000) # 40_000 GPU bytes + rewind!(pool) + + s = trim!(pool) + @test s.slots_released == 1 + # summarysize(CuArray) reports only the CPU-side handle; sizeof gives the + # real GPU bytes. This asserts the CUDA byte override is used. + @test s.estimated_bytes_released >= 40_000 + end + + @testset "wrapper truncation + per-type + force_gc (CUDA)" begin + pool = CuAdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Float32, 4, 5) # 2-D → arr_wrappers[2][1] + acquire!(pool, Int32, 100) + rewind!(pool) + + @test pool.float32.arr_wrappers[2] !== nothing + @test length(pool.float32.arr_wrappers[2]) == 1 + + s = trim!(pool, Float32) # per-type: only Float32 + @test length(pool.float32.vectors) == 0 + @test length(pool.int32.vectors) == 1 # untouched + @test s.slots_released == 1 + for w in pool.float32.arr_wrappers + w === nothing && continue + @test length(w) <= pool.float32.n_active + end + + @test trim!(pool; force_gc = true).gc_triggered == true + end + + @testset "self-heal: acquire! after trim! (CUDA)" begin + pool = CuAdaptiveArrayPool{0}() + + checkpoint!(pool) + acquire!(pool, Float32, 1000) + rewind!(pool) + trim!(pool) + @test length(pool.float32.vectors) == 0 + + checkpoint!(pool) + v = acquire!(pool, Float32, 7) + @test length(v) == 7 + @test length(pool.float32.vectors) == 1 + rewind!(pool) + end + +end From bfd7b1b606f141e4ab928ce397ac714be2524917 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 15:43:28 -0700 Subject: [PATCH 3/7] (fix): correct `trim!` GPU byte accounting + doc wording Address code-review findings: - [P2] GPU `_inactive_storage_bytes` used `sizeof`, which follows logical length. In runtime-check mode (R/S>=1), rewind shrinks released slots to dims 0 via `_resize_to_fit!` while preserving `maxsize`, so the estimate read 0 for buffers still holding device capacity. Use `getfield(v, :maxsize)` (allocated device bytes) on Metal + CUDA, and add R=1/S=1 safety-mode byte-accounting tests. - [P3] api-essentials wording implied `trim!` preserves "warm, normal-sized" buffers, but it drops ALL inactive retained slots (at top level, every reusable buffer). Reword to the actual guarantee: only active arrays are preserved. Verified: Metal suite green on Apple Silicon (incl. the new R=1 test); CUDA ext precompiles and `trim!` dispatch confirmed (GPU execution on CI). CPU unaffected (`summarysize` is already capacity-aware). --- docs/src/basics/api-essentials.md | 10 +++++----- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 11 ++++++----- ext/AdaptiveArrayPoolsMetalExt/state.jl | 11 ++++++----- test/cuda/test_trim.jl | 15 ++++++++++++++- test/metal/test_trim.jl | 15 ++++++++++++++- 5 files changed, 45 insertions(+), 17 deletions(-) diff --git a/docs/src/basics/api-essentials.md b/docs/src/basics/api-essentials.md index c52015a1..19d03109 100644 --- a/docs/src/basics/api-essentials.md +++ b/docs/src/basics/api-essentials.md @@ -105,11 +105,11 @@ reset!(get_task_local_pool()) ### `trim!(pool; force_gc=false)` -Releases the pool's references to **inactive** retained buffers (those already -freed by `rewind!`/scope exit), while preserving the **active** arrays still in -use. This is the operation for reclaiming memory after a rare large workload -without discarding the warm, normal-sized buffers that keep the hot path -allocation-free. +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 diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index 547c0c0a..5c736901 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -357,14 +357,15 @@ end # trim! for CuAdaptiveArrayPool (parity with CPU/Metal) # ============================================================================== -# CUDA backing storage lives on the GPU. `sizeof(CuVector)` = length * sizeof(T) -# = the real device buffer size, whereas `Base.summarysize` (the generic CPU -# helper) only sees the small CPU-side handle. Override so the byte estimate is -# correct. (CUDA inactive slots are not shrunk to length 0, so length == capacity.) +# 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 += sizeof(tp.vectors[i]) + @inbounds total += getfield(tp.vectors[i], :maxsize) end return total end diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 371df3d1..107f99c1 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -339,14 +339,15 @@ end # trim! for MetalAdaptiveArrayPool (parity with CPU) # ============================================================================== -# Metal backing storage lives on the GPU. `sizeof(MtlArray)` = length * sizeof(T) -# = the real GPU buffer size, whereas `Base.summarysize` (the generic CPU helper) -# only sees the small CPU-side handle. Override so the byte estimate is correct. -# (Metal inactive slots are not shrunk to length 0, so length == capacity here.) +# 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 += sizeof(tp.vectors[i]) + @inbounds total += getfield(tp.vectors[i], :maxsize) end return total end diff --git a/test/cuda/test_trim.jl b/test/cuda/test_trim.jl index 2b44a788..ad683328 100644 --- a/test/cuda/test_trim.jl +++ b/test/cuda/test_trim.jl @@ -62,11 +62,24 @@ s = trim!(pool) @test s.slots_released == 1 - # summarysize(CuArray) reports only the CPU-side handle; sizeof gives the + # summarysize(CuArray) reports only the CPU-side handle; maxsize gives the # real GPU bytes. This asserts the CUDA byte override is used. @test s.estimated_bytes_released >= 40_000 end + @testset "byte estimate is capacity-aware in safety mode (CUDA, S=1)" begin + pool = CuAdaptiveArrayPool{1}() # S=1: safety invalidation on + + checkpoint!(pool) + acquire!(pool, Float32, 10_000) # 40_000 device bytes + rewind!(pool) # S=1: released slot dims → 0, maxsize preserved + + @test length(pool.float32.vectors[1]) == 0 # logical length 0 (sizeof would be 0)... + s = trim!(pool) + @test s.slots_released == 1 + @test s.estimated_bytes_released >= 40_000 # ...but maxsize (device capacity) is counted + end + @testset "wrapper truncation + per-type + force_gc (CUDA)" begin pool = CuAdaptiveArrayPool{0}() diff --git a/test/metal/test_trim.jl b/test/metal/test_trim.jl index 17bc855a..08dc406e 100644 --- a/test/metal/test_trim.jl +++ b/test/metal/test_trim.jl @@ -62,11 +62,24 @@ s = trim!(pool) @test s.slots_released == 1 - # summarysize(MtlArray) reports only the ~80-byte CPU handle; sizeof gives + # summarysize(MtlArray) reports only the ~80-byte CPU handle; maxsize gives # the real 40_000 GPU bytes. This asserts the Metal byte override is used. @test s.estimated_bytes_released >= 40_000 end + @testset "byte estimate is capacity-aware in safety mode (Metal, R=1)" begin + pool = MetalAdaptiveArrayPool{1, METAL_STORAGE}() # R=1: safety invalidation on + + checkpoint!(pool) + acquire!(pool, Float32, 10_000) # 40_000 device bytes + rewind!(pool) # R=1: released slot dims → 0, maxsize preserved + + @test length(pool.float32.vectors[1]) == 0 # logical length 0 (sizeof would be 0)... + s = trim!(pool) + @test s.slots_released == 1 + @test s.estimated_bytes_released >= 40_000 # ...but maxsize (device capacity) is counted + end + @testset "wrapper truncation + per-type + force_gc (Metal)" begin pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() From 1691908b5dc13cfa39897c9fd036916898d8248a Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 17:32:14 -0700 Subject: [PATCH 4/7] =?UTF-8?q?(fix):=20Metal=20=E2=80=94=20free=20GPU=20m?= =?UTF-8?q?emory=20on=20trim!/empty!=20(drop=20redundant=20DataRef=20copy)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MtlArray wrapper-creation sites passed `copy(getfield(vec, :data))` to the MtlArray constructor, but that constructor already copies/retains its DataRef argument. The extra copy created a second retained DataRef that was never released, pinning each buffer's refcount above 0 — so the GPU buffer was never freed by trim!, empty!, GC, or even explicit unsafe_free!. Pass the DataRef directly (the constructor copies it) at both wrapper sites (get_array! cache-miss and reshape cache-miss). Buffer sharing is unchanged (write-through verified); the refcount now reaches 0 and Metal device memory is reclaimed. Verified on Apple Silicon: 4 GiB -> 384 KiB after trim!(force_gc=true); full Metal suite green incl. a new regression test asserting device memory drops. CUDA is unaffected: its CuArray constructor takes ownership of the DataRef (no internal copy), so CUDA's explicit copy() is required and correct. --- ext/AdaptiveArrayPoolsMetalExt/acquire.jl | 13 +++++++++---- test/metal/test_trim.jl | 23 +++++++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/ext/AdaptiveArrayPoolsMetalExt/acquire.jl b/ext/AdaptiveArrayPoolsMetalExt/acquire.jl index a1f1909a..726e8821 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/acquire.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/acquire.jl @@ -218,9 +218,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 +286,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), ) diff --git a/test/metal/test_trim.jl b/test/metal/test_trim.jl index 08dc406e..593db9f9 100644 --- a/test/metal/test_trim.jl +++ b/test/metal/test_trim.jl @@ -119,4 +119,27 @@ rewind!(pool) end + @testset "trim!(force_gc=true) actually frees Metal device memory" begin + # Regression guard for the wrapper double-copy bug: an extra copy() of the + # backing DataRef pinned the buffer refcount so it never reached 0, and + # device memory was never reclaimed (by trim! OR empty!). This asserts the + # device allocation actually drops. + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + dev = Metal.device() + GC.gc(true) + base = Int(dev.currentAllocatedSize) + + checkpoint!(pool) + v = acquire!(pool, Float32, 16_777_216) # 64 MiB + fill!(v, 1.0f0) + Metal.synchronize() + grown = Int(dev.currentAllocatedSize) + @test grown - base >= 64_000_000 # device memory grew by the allocation + + rewind!(pool) + trim!(pool; force_gc = true) + GC.gc(true) + @test Int(dev.currentAllocatedSize) <= grown - 48_000_000 # and was reclaimed + end + end From 25e6803e73615a2827c6d22c3912a0f480c36722 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 17:32:33 -0700 Subject: [PATCH 5/7] (feat): aggressive trim!(force_gc=true) reclaim + GPU memory smoke scripts - CUDA: force_gc=true now also calls CUDA.reclaim() to return pooled blocks to the driver (Metal/CPU have no further reclaim step). Docstring updated to describe the backend-aware behavior. - Add scripts/trim_smoke_{metal,cuda}.jl: manual scripts that print the Julia heap (gc_live_bytes), process RSS, and GPU device memory before/after trim! to confirm real reclamation. The Metal script runs on Apple Silicon; the CUDA script on an NVIDIA box. --- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 8 +++- scripts/trim_smoke_cuda.jl | 66 ++++++++++++++++++++++++++ scripts/trim_smoke_metal.jl | 62 ++++++++++++++++++++++++ src/state.jl | 8 ++-- 4 files changed, 139 insertions(+), 5 deletions(-) create mode 100644 scripts/trim_smoke_cuda.jl create mode 100644 scripts/trim_smoke_metal.jl diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index 5c736901..5270877f 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -386,7 +386,9 @@ function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool; force_gc::Bool = fa wrappers += s.wrappers_released bytes += s.estimated_bytes_released end - force_gc && GC.gc() + # force_gc on CUDA also returns CUDA.jl's pooled blocks to the driver + # (CUDA.reclaim()) so freed VRAM is actually released, not just GC-eligible. + force_gc && (GC.gc(); CUDA.reclaim()) return (; slots_released = slots, wrappers_released = wrappers, estimated_bytes_released = bytes, gc_triggered = force_gc, @@ -395,7 +397,9 @@ end @inline function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) - force_gc && GC.gc() + # force_gc on CUDA also returns CUDA.jl's pooled blocks to the driver + # (CUDA.reclaim()) so freed VRAM is actually released, not just GC-eligible. + force_gc && (GC.gc(); CUDA.reclaim()) return (; slots_released = s.slots_released, wrappers_released = s.wrappers_released, diff --git a/scripts/trim_smoke_cuda.jl b/scripts/trim_smoke_cuda.jl new file mode 100644 index 00000000..e9dbcdbb --- /dev/null +++ b/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= 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 diff --git a/scripts/trim_smoke_metal.jl b/scripts/trim_smoke_metal.jl new file mode 100644 index 00000000..1005fdfd --- /dev/null +++ b/scripts/trim_smoke_metal.jl @@ -0,0 +1,62 @@ +# Metal `trim!` memory smoke test (Apple Silicon). +# +# Not a unit test — a manual script that prints memory before/after `trim!` to +# show that reclamation actually happens. Run in an env with Metal + AdaptiveArrayPools: +# +# julia --project= scripts/trim_smoke_metal.jl +# +# Three probes: +# julia_live = Base.gc_live_bytes() Julia-tracked live heap +# rss = process RSS (ps) OS resident memory +# metal_alloc = Metal.device().currentAllocatedSize Metal device allocation +# (Apple Silicon is unified-memory, so this is shared RAM) +# +# Metal has no reclaim() (no caching allocator exposed); force_gc=true runs +# GC.gc(), whose finalizers free the MtlArray buffers. + +using Metal +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 +metal_alloc() = Int(Metal.device().currentAllocatedSize) + +function report(label) + return println( + rpad(label, 30), + "julia_live=", rpad(human(jlive()), 11), + " rss=", rpad(rss() < 0 ? "?" : human(rss()), 11), + " metal_alloc=", human(metal_alloc()), + ) +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_metal_pool() +GC.gc(true) +report("0) start") + +# Allocate NSLOTS x 1 GiB into the pool (fill! to commit), inside a scope. +checkpoint!(pool) +for i in 1:NSLOTS + v = acquire!(pool, Float32, N) + fill!(v, Float32(i)) +end +Metal.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.gc(true) +report("3) after trim!(force_gc=true)") +println(" summary = ", summary) diff --git a/src/state.jl b/src/state.jl index 6571375b..28d260ad 100644 --- a/src/state.jl +++ b/src/state.jl @@ -812,9 +812,11 @@ leaves `n_active`, checkpoint stacks, and depth state unchanged. Returns `(; slots_released, wrappers_released, estimated_bytes_released, gc_triggered)`. -`force_gc=true` calls `GC.gc()` after detaching references. This still does not -guarantee immediate OS/VRAM return; it only asks Julia to collect now that the -pool no longer holds the buffers. +`force_gc=true` runs `GC.gc()` after detaching references, and on backends with a +caching device allocator it also returns the pooled memory to the driver (CUDA: +`CUDA.reclaim()`; Metal/CPU have no further reclaim step). Even so, immediate +OS/VRAM return is not guaranteed — it only asks the runtime to collect now that +the pool no longer holds the buffers. !!! note "Julia version" Actual reclamation requires Julia 1.12+. On older Julia (the legacy pool From fd6efd95be80f6476df69a2072affcfeeee9cf03 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 19:11:03 -0700 Subject: [PATCH 6/7] (fix): trim!(pool, T) no longer creates a pool for unused types (Copilot review, PR #44) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-type `trim!` called `get_typed_pool!(pool, T)`, which registers a new fallback pool in `pool.others` for a never-used type — a surprising mutation/allocation for a reclamation API. It now no-ops (zero summary, still honoring `force_gc`) unless `T` is a fixed-slot type or already present in `others`. Applied to CPU + Metal + CUDA via a per-backend fixed-slot type Union; fixed-slot access is unaffected. Also fills the trim! coverage gaps codecov flagged on PR #44: - modern: per-type `force_gc` path and the new no-create path (test/test_trim.jl) - legacy (Julia <1.12): `DisabledPool` trim! no-ops (test/test_trim_legacy.jl) Verified: full CPU+Metal suite green (1.12); legacy test green (1.10); the trim! region of src/state.jl is fully covered (guard branch + per-type force_gc hit). --- ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl | 4 +++ ext/AdaptiveArrayPoolsCUDAExt/state.jl | 10 ++++++ ext/AdaptiveArrayPoolsMetalExt/dispatch.jl | 4 +++ ext/AdaptiveArrayPoolsMetalExt/state.jl | 10 ++++++ src/state.jl | 11 ++++++ src/types.jl | 5 +++ test/test_trim.jl | 39 ++++++++++++++++++++++ test/test_trim_legacy.jl | 8 +++++ 8 files changed, 91 insertions(+) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl b/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl index 5dac0a31..968c84c4 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl +++ b/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 diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index 5270877f..c8c98ab3 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -396,6 +396,16 @@ function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool; force_gc::Bool = fa end @inline function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} + # Never create a pool for a never-used type (see Copilot review, PR #44): + # get_typed_pool! would register a new fallback pool on a miss. Only proceed + # for fixed-slot types or fallback types already in the pool. + if !(T <: _CUDA_FIXED_TYPES) && !haskey(pool.others, T) + force_gc && (GC.gc(); CUDA.reclaim()) + return (; + slots_released = 0, wrappers_released = 0, + estimated_bytes_released = 0, gc_triggered = force_gc, + ) + end s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) # force_gc on CUDA also returns CUDA.jl's pooled blocks to the driver # (CUDA.reclaim()) so freed VRAM is actually released, not just GC-eligible. diff --git a/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl b/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl index 0c938f6a..879e76b2 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl +++ b/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} diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 107f99c1..b6668c0e 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -376,6 +376,16 @@ function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool; force_gc::Bool = end @inline function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} + # Never create a pool for a never-used type (see Copilot review, PR #44): + # get_typed_pool! would register a new fallback pool on a miss. Only proceed + # for fixed-slot types or fallback types already in the pool. + if !(T <: _METAL_FIXED_TYPES) && !haskey(pool.others, T) + force_gc && GC.gc() + return (; + slots_released = 0, wrappers_released = 0, + estimated_bytes_released = 0, gc_triggered = force_gc, + ) + end s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) force_gc && GC.gc() return (; diff --git a/src/state.jl b/src/state.jl index 28d260ad..f36de577 100644 --- a/src/state.jl +++ b/src/state.jl @@ -855,6 +855,17 @@ end Trim inactive slots for a single element type `T` only. See [`trim!`](@ref). """ function trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} + # Never create a pool for a type that was never used: get_typed_pool! would + # register a new fallback pool in `pool.others` on a miss — a surprising + # mutation/allocation for a reclamation call. Only proceed for fixed-slot + # types (always present) or fallback types already in the pool. + if !(T <: _FIXED_SLOT_TYPES) && !haskey(pool.others, T) + force_gc && GC.gc() + return (; + slots_released = 0, wrappers_released = 0, + estimated_bytes_released = 0, gc_triggered = force_gc, + ) + end s = _trim_inactive_typed_pool!(get_typed_pool!(pool, T)) force_gc && GC.gc() return (; diff --git a/src/types.jl b/src/types.jl index f38c6deb..5f030c5c 100644 --- a/src/types.jl +++ b/src/types.jl @@ -466,6 +466,11 @@ end @inline get_typed_pool!(p::AdaptiveArrayPool, ::Type{Bool}) = p.bool @inline get_typed_pool!(p::AdaptiveArrayPool, ::Type{Bit}) = p.bits +# Fixed-slot element types (those with a dedicated struct field above). Used by +# `trim!(pool, T)` to recognize types that always have a pool, without consulting +# or creating the `others` fallback. Keep in sync with the methods above. +const _FIXED_SLOT_TYPES = Union{Float64, Float32, Int64, Int32, ComplexF64, ComplexF32, Bool, Bit} + # Slow Path: rare types via IdDict @inline function get_typed_pool!(p::AdaptiveArrayPool, ::Type{T}) where {T} tp = get(p.others, T, nothing) diff --git a/test/test_trim.jl b/test/test_trim.jl index d8919e37..09bad4d9 100644 --- a/test/test_trim.jl +++ b/test/test_trim.jl @@ -145,6 +145,14 @@ using AdaptiveArrayPools: get_typed_pool!, trim!, checkpoint!, rewind! pool = AdaptiveArrayPool{0}() @test trim!(pool; force_gc = true).gc_triggered == true @test trim!(pool; force_gc = false).gc_triggered == false + + # per-type form: force_gc through the real-trim path + checkpoint!(pool) + acquire!(pool, Float64, 100) + rewind!(pool) + s = trim!(pool, Float64; force_gc = true) + @test s.slots_released == 1 + @test s.gc_triggered == true end @testset "reset!(pool); trim!(pool) releases all retained buffers (regression)" begin @@ -224,4 +232,35 @@ using AdaptiveArrayPools: get_typed_pool!, trim!, checkpoint!, rewind! @test s.estimated_bytes_released > 0 end + @testset "per-type trim! never creates a pool for an unused type" begin + pool = AdaptiveArrayPool{0}() + + # Unused fallback type (UInt8 is not a fixed slot, never acquired): + # trim! must NOT register a pool in `others` while trying to reclaim. + @test !haskey(pool.others, UInt8) + n_before = length(pool.others) + s = trim!(pool, UInt8) + @test s == (; + slots_released = 0, wrappers_released = 0, + estimated_bytes_released = 0, gc_triggered = false, + ) + @test length(pool.others) == n_before + @test !haskey(pool.others, UInt8) # still not created + + # force_gc is still honored on the no-op path. + @test trim!(pool, UInt8; force_gc = true).gc_triggered == true + @test !haskey(pool.others, UInt8) + + # Fixed-slot type (always present) is a safe no-op too — no `others` growth. + @test trim!(pool, Float64).slots_released == 0 + @test length(pool.others) == n_before + + # An existing fallback type is still trimmed normally. + checkpoint!(pool) + acquire!(pool, UInt8, 100) + rewind!(pool) + @test haskey(pool.others, UInt8) + @test trim!(pool, UInt8).slots_released == 1 + end + end diff --git a/test/test_trim_legacy.jl b/test/test_trim_legacy.jl index d30c2038..0fd9de42 100644 --- a/test/test_trim_legacy.jl +++ b/test/test_trim_legacy.jl @@ -22,4 +22,12 @@ using AdaptiveArrayPools: trim!, get_task_local_pool # Legacy no-op ignores force_gc (no reclamation to do). @test trim!(pool; force_gc = true).gc_triggered == false + + # DisabledPool: zero-summary no-op on the legacy path too. + zd = trim!(DISABLED_CPU) + @test zd.slots_released == 0 + @test zd.wrappers_released == 0 + @test zd.estimated_bytes_released == 0 + @test zd.gc_triggered == false + @test trim!(DISABLED_CPU, Float64).slots_released == 0 end From b483e92def5fa67d55f89c3f512dff2cadc5d14e Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sun, 28 Jun 2026 20:49:00 -0700 Subject: [PATCH 7/7] (refactor): type-stable trim! summary + Type... varargs parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type stability: trim!(pool) inferred its summary as NamedTuple{...,<:Tuple{Any,Any,Any,Bool}} — the foreach_fixed_slot closure boxed the accumulators (Core.Box) and the others-loop over Vector{Any} widened the count fields to Any. Replaced with an unrolled @generated fold (_trim_fixed_counts!) plus a barrier loop (_trim_others_counts!), and funneled every overload through one _trim_summary helper. The return is now the concrete @NamedTuple{...::Int, ..., gc_triggered::Bool} (also drops ~288B of boxing). Varargs parity: added trim!(pool, types::Type...) on CPU, CUDA, Metal, legacy, and DisabledPool, mirroring checkpoint!/rewind!/reset!. A shared _trim_one_counts! guard (skips never-used fallback types) backs both the single-type and varargs forms; GPU reuses the type-stable Ref-accumulator path. Tests: type-stability (@inferred + return_types for all positional forms), varargs behavior + per-type guard, N-D wrapper-cache rebuild after trim, and mid-scope trim leaving the checkpoint stack untouched before a later rewind. GPU varargs tests added to the CUDA/Metal suites. Comments: fixed the stale Metal get_array! docstring (DataRef passed directly, not copy()), softened the CUDA force_gc reclaim wording to match the docstring's "not guaranteed" hedge, and dropped PR-ref provenance from CUDA/Metal source. --- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 82 +++++++----- ext/AdaptiveArrayPoolsMetalExt/acquire.jl | 6 +- ext/AdaptiveArrayPoolsMetalExt/state.jl | 72 +++++----- src/legacy/state.jl | 4 + src/state.jl | 137 ++++++++++++------- test/cuda/test_trim.jl | 31 +++++ test/metal/test_trim.jl | 31 +++++ test/test_trim.jl | 153 ++++++++++++++++++++++ 8 files changed, 399 insertions(+), 117 deletions(-) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index c8c98ab3..f2a7542a 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -371,49 +371,61 @@ function AdaptiveArrayPools._inactive_storage_bytes(tp::CuTypedPool, first::Int, end function AdaptiveArrayPools.trim!(pool::CuAdaptiveArrayPool; force_gc::Bool = false) - slots = 0 - wrappers = 0 - bytes = 0 + # 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 - s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) - slots += s.slots_released - wrappers += s.wrappers_released - bytes += s.estimated_bytes_released - end - for tp in values(pool.others) - s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) - slots += s.slots_released - wrappers += s.wrappers_released - bytes += s.estimated_bytes_released + 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()) so freed VRAM is actually released, not just GC-eligible. + # (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 (; - slots_released = slots, wrappers_released = wrappers, - estimated_bytes_released = bytes, gc_triggered = force_gc, + 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} - # Never create a pool for a never-used type (see Copilot review, PR #44): - # get_typed_pool! would register a new fallback pool on a miss. Only proceed - # for fixed-slot types or fallback types already in the pool. - if !(T <: _CUDA_FIXED_TYPES) && !haskey(pool.others, T) - force_gc && (GC.gc(); CUDA.reclaim()) - return (; - slots_released = 0, wrappers_released = 0, - estimated_bytes_released = 0, gc_triggered = force_gc, - ) - end - s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) + counts = AdaptiveArrayPools._trim_one_counts!(pool, T) # force_gc on CUDA also returns CUDA.jl's pooled blocks to the driver - # (CUDA.reclaim()) so freed VRAM is actually released, not just GC-eligible. + # (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 (; - slots_released = s.slots_released, - wrappers_released = s.wrappers_released, - estimated_bytes_released = s.estimated_bytes_released, - gc_triggered = force_gc, - ) + 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 diff --git a/ext/AdaptiveArrayPoolsMetalExt/acquire.jl b/ext/AdaptiveArrayPoolsMetalExt/acquire.jl index 726e8821..93a9a6cc 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/acquire.jl +++ b/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) diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index b6668c0e..de2ceb81 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -353,45 +353,53 @@ function AdaptiveArrayPools._inactive_storage_bytes(tp::MetalTypedPool, first::I end function AdaptiveArrayPools.trim!(pool::MetalAdaptiveArrayPool; force_gc::Bool = false) - slots = 0 - wrappers = 0 - bytes = 0 + # 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 - s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) - slots += s.slots_released - wrappers += s.wrappers_released - bytes += s.estimated_bytes_released - end - for tp in values(pool.others) - s = AdaptiveArrayPools._trim_inactive_typed_pool!(tp) - slots += s.slots_released - wrappers += s.wrappers_released - bytes += s.estimated_bytes_released + 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 (; - slots_released = slots, wrappers_released = wrappers, - estimated_bytes_released = bytes, gc_triggered = force_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} - # Never create a pool for a never-used type (see Copilot review, PR #44): - # get_typed_pool! would register a new fallback pool on a miss. Only proceed - # for fixed-slot types or fallback types already in the pool. - if !(T <: _METAL_FIXED_TYPES) && !haskey(pool.others, 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 (; - slots_released = 0, wrappers_released = 0, - estimated_bytes_released = 0, gc_triggered = force_gc, - ) + return AdaptiveArrayPools._trim_summary(counts, force_gc) end - s = AdaptiveArrayPools._trim_inactive_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T)) - force_gc && GC.gc() - return (; - slots_released = s.slots_released, - wrappers_released = s.wrappers_released, - estimated_bytes_released = s.estimated_bytes_released, - gc_triggered = force_gc, - ) end diff --git a/src/legacy/state.jl b/src/legacy/state.jl index bad2e26e..c8253e98 100644 --- a/src/legacy/state.jl +++ b/src/legacy/state.jl @@ -782,7 +782,11 @@ end trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} = trim!(pool; force_gc = force_gc) +trim!(pool::AdaptiveArrayPool, types::Type...; force_gc::Bool = false) = + trim!(pool; force_gc = force_gc) + trim!(; force_gc::Bool = false) = trim!(get_task_local_pool(); force_gc = force_gc) trim!(::DisabledPool; force_gc::Bool = false) = _ZERO_TRIM_SUMMARY trim!(::DisabledPool, ::Type{T}; force_gc::Bool = false) where {T} = _ZERO_TRIM_SUMMARY +trim!(::DisabledPool, types::Type...; force_gc::Bool = false) = _ZERO_TRIM_SUMMARY diff --git a/src/state.jl b/src/state.jl index f36de577..34fad009 100644 --- a/src/state.jl +++ b/src/state.jl @@ -741,12 +741,18 @@ end # trim! is a manual operation. It adds no code to the hot acquire path or the # automatic rewind path. -# Zero summary: returned by no-op trims (DisabledPool, nothing to release). -const _ZERO_TRIM_SUMMARY = (; - slots_released = 0, wrappers_released = 0, - estimated_bytes_released = 0, gc_triggered = false, +# Build the public summary NamedTuple from raw `(slots, wrappers, bytes)` counts +# and the gc flag. Single definition point for the `trim!` return shape — every +# overload (CPU/GPU/per-type/zero) funnels through here, so they all return the +# identical concrete `@NamedTuple{...::Int, ..., gc_triggered::Bool}`. +@inline _trim_summary(c::NTuple{3, Int}, gc::Bool) = (; + slots_released = c[1], wrappers_released = c[2], + estimated_bytes_released = c[3], gc_triggered = gc, ) +# Zero summary: returned by no-op trims (DisabledPool, nothing to release). +const _ZERO_TRIM_SUMMARY = _trim_summary((0, 0, 0), false) + # Drop cached N-D wrappers for inactive slots. A cached Array/BitArray holds a # `:ref` into its slot's backing buffer, so leaving inactive wrappers in place # would pin the very buffers we just detached. Each per-slot wrapper vector @@ -781,9 +787,11 @@ function _inactive_storage_bytes(tp::AbstractTypedPool, first::Int, last::Int) end # Trim one typed pool: drop inactive backing-vector and wrapper references by -# truncating `tp.vectors` and the wrapper caches to the active count. Returns a -# per-pool summary NamedTuple. Bytes are measured before truncation. -function _trim_inactive_typed_pool!(tp::AbstractTypedPool) +# truncating `tp.vectors` and the wrapper caches to the active count. Returns the +# raw `(slots, wrappers, bytes)` counts (bytes measured before truncation). The +# `::NTuple{3, Int}` return type keeps callers concrete even when this is reached +# through dynamic dispatch over the `others` `Vector{Any}`. +function _trim_counts!(tp::AbstractTypedPool)::NTuple{3, Int} keep = tp.n_active old_len = length(tp.vectors) released = old_len - keep @@ -793,11 +801,42 @@ function _trim_inactive_typed_pool!(tp::AbstractTypedPool) bytes = _inactive_storage_bytes(tp, keep + 1, old_len) resize!(tp.vectors, keep) end - return (; - slots_released = max(0, released), - wrappers_released = wrappers_released, - estimated_bytes_released = bytes, - ) + return (max(0, released), wrappers_released, bytes) +end + +# Sum `_trim_counts!` over every fixed slot, unrolled at compile time. A plain +# `foreach_fixed_slot(pool) do tp ... end` closure would box the accumulators +# (Core.Box) and widen the summary to `Any`; this unrolled fold stays concrete +# (`NTuple{3, Int}`) and allocation-free. +@generated function _trim_fixed_counts!(pool::AdaptiveArrayPool) + syms = [Symbol(:c, i) for i in 1:length(FIXED_SLOT_FIELDS)] + assigns = [ + :($(syms[i]) = _trim_counts!(getfield(pool, $(QuoteNode(f))))) + for (i, f) in enumerate(FIXED_SLOT_FIELDS) + ] + sums = [Expr(:call, :+, [:($(s)[$j]) for s in syms]...) for j in 1:3] + return quote + Base.@_inline_meta + $(assigns...) + ($(sums...),) + end +end + +# Sum `_trim_counts!` over the `others` (non-fixed-slot) typed pools. A plain loop +# (no closure) avoids boxing; the `::NTuple{3, Int}` assertion keeps the running +# total concrete despite dynamic dispatch over the `Any`-typed iterable. Shared by +# the CPU (`_others_values`) and GPU (`values(pool.others)`) backends. +function _trim_others_counts!(others)::NTuple{3, Int} + slots = 0 + wrappers = 0 + bytes = 0 + for tp in others + c = _trim_counts!(tp)::NTuple{3, Int} + slots += c[1] + wrappers += c[2] + bytes += c[3] + end + return (slots, wrappers, bytes) end """ @@ -827,26 +866,21 @@ the pool no longer holds the buffers. See also: [`reset!`](@ref), [`empty!`](@ref). """ function trim!(pool::AdaptiveArrayPool; force_gc::Bool = false) - slots = 0 - wrappers = 0 - bytes = 0 - foreach_fixed_slot(pool) do tp - s = _trim_inactive_typed_pool!(tp) - slots += s.slots_released - wrappers += s.wrappers_released - bytes += s.estimated_bytes_released - end - for tp in pool._others_values - s = _trim_inactive_typed_pool!(tp) - slots += s.slots_released - wrappers += s.wrappers_released - bytes += s.estimated_bytes_released - end + counts = _trim_fixed_counts!(pool) .+ _trim_others_counts!(pool._others_values) force_gc && GC.gc() - return (; - slots_released = slots, wrappers_released = wrappers, - estimated_bytes_released = bytes, gc_triggered = force_gc, - ) + return _trim_summary(counts, force_gc) +end + +# Guard + trim for one element type, returning raw `(slots, wrappers, bytes)`. +# Shared by the single-type and varargs `trim!` forms. Never *creates* a pool: a +# never-used fallback type is skipped, because `get_typed_pool!` would register a +# new fallback pool in `pool.others` on a miss — a surprising mutation/allocation +# for a reclamation call. Fixed-slot types are always present; fallback types are +# trimmed only if already in the pool. The `::NTuple{3, Int}` return keeps the +# varargs fold type-stable even when the type is not a compile-time constant. +function _trim_one_counts!(pool::AdaptiveArrayPool, ::Type{T})::NTuple{3, Int} where {T} + (!(T <: _FIXED_SLOT_TYPES) && !haskey(pool.others, T)) && return (0, 0, 0) + return _trim_counts!(get_typed_pool!(pool, T)) end """ @@ -855,25 +889,31 @@ end Trim inactive slots for a single element type `T` only. See [`trim!`](@ref). """ function trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) where {T} - # Never create a pool for a type that was never used: get_typed_pool! would - # register a new fallback pool in `pool.others` on a miss — a surprising - # mutation/allocation for a reclamation call. Only proceed for fixed-slot - # types (always present) or fallback types already in the pool. - if !(T <: _FIXED_SLOT_TYPES) && !haskey(pool.others, T) + counts = _trim_one_counts!(pool, T) + force_gc && GC.gc() + return _trim_summary(counts, force_gc) +end + +""" + trim!(pool::AdaptiveArrayPool, types::Type...; force_gc::Bool = false) + +Trim inactive slots for several element types at once, running `GC.gc()` (when +`force_gc=true`) a single time after all listed types are detached. Returns one +combined summary. Uses `@generated` for zero-overhead compile-time unrolling, +mirroring [`reset!`](@ref). See [`trim!`](@ref). +""" +@generated function trim!(pool::AdaptiveArrayPool, types::Type...; force_gc::Bool = false) + n = length(types) + syms = [Symbol(:c, i) for i in 1:n] + assigns = [:($(syms[i]) = _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 (; - slots_released = 0, wrappers_released = 0, - estimated_bytes_released = 0, gc_triggered = force_gc, - ) + return _trim_summary(counts, force_gc) end - s = _trim_inactive_typed_pool!(get_typed_pool!(pool, T)) - force_gc && GC.gc() - return (; - slots_released = s.slots_released, - wrappers_released = s.wrappers_released, - estimated_bytes_released = s.estimated_bytes_released, - gc_triggered = force_gc, - ) end """ @@ -904,3 +944,4 @@ Base.empty!(::DisabledPool) = nothing trim!(::DisabledPool; force_gc::Bool = false) = _ZERO_TRIM_SUMMARY trim!(::DisabledPool, ::Type{T}; force_gc::Bool = false) where {T} = _ZERO_TRIM_SUMMARY +trim!(::DisabledPool, types::Type...; force_gc::Bool = false) = _ZERO_TRIM_SUMMARY diff --git a/test/cuda/test_trim.jl b/test/cuda/test_trim.jl index ad683328..998733ad 100644 --- a/test/cuda/test_trim.jl +++ b/test/cuda/test_trim.jl @@ -119,4 +119,35 @@ rewind!(pool) end + @testset "varargs trim!(pool, T1, T2, ...) trims only listed types (CUDA)" begin + # Parity with checkpoint!/rewind!, which also expose a Type... form on CUDA. + pool = CuAdaptiveArrayPool{0}() + checkpoint!(pool) + acquire!(pool, Float32, 100) + acquire!(pool, Int32, 50) + acquire!(pool, Float64, 25) # Float64 is a CUDA fixed slot too + rewind!(pool) + + s = trim!(pool, Float32, Int32) # trim two; leave Float64 + + @test length(pool.float32.vectors) == 0 + @test length(pool.int32.vectors) == 0 + @test length(pool.float64.vectors) == 1 # untouched + @test s.slots_released == 2 + @test s.gc_triggered == false + + # force_gc honored once across all listed types (+ CUDA.reclaim) + pool2 = CuAdaptiveArrayPool{0}() + checkpoint!(pool2) + acquire!(pool2, Float32, 10) + acquire!(pool2, Int32, 10) + rewind!(pool2) + @test trim!(pool2, Float32, Int32; force_gc = true).gc_triggered == true + + # A never-used fallback type in the list is skipped, never created. + pool3 = CuAdaptiveArrayPool{0}() + @test trim!(pool3, Float32, UInt16).slots_released == 0 + @test !haskey(pool3.others, UInt16) + end + end diff --git a/test/metal/test_trim.jl b/test/metal/test_trim.jl index 593db9f9..1cccdd6b 100644 --- a/test/metal/test_trim.jl +++ b/test/metal/test_trim.jl @@ -142,4 +142,35 @@ @test Int(dev.currentAllocatedSize) <= grown - 48_000_000 # and was reclaimed end + @testset "varargs trim!(pool, T1, T2, ...) trims only listed types (Metal)" begin + # Parity with checkpoint!/rewind!, which also expose a Type... form on Metal. + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + checkpoint!(pool) + acquire!(pool, Float32, 100) + acquire!(pool, Int32, 50) + acquire!(pool, Int64, 25) # Int64 is a Metal fixed slot too + rewind!(pool) + + s = trim!(pool, Float32, Int32) # trim two; leave Int64 + + @test length(pool.float32.vectors) == 0 + @test length(pool.int32.vectors) == 0 + @test length(pool.int64.vectors) == 1 # untouched + @test s.slots_released == 2 + @test s.gc_triggered == false + + # force_gc honored once across all listed types + pool2 = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + checkpoint!(pool2) + acquire!(pool2, Float32, 10) + acquire!(pool2, Int32, 10) + rewind!(pool2) + @test trim!(pool2, Float32, Int32; force_gc = true).gc_triggered == true + + # A never-used fallback type in the list is skipped, never created. + pool3 = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + @test trim!(pool3, Float32, UInt16).slots_released == 0 + @test !haskey(pool3.others, UInt16) + end + end diff --git a/test/test_trim.jl b/test/test_trim.jl index 09bad4d9..ba7065c3 100644 --- a/test/test_trim.jl +++ b/test/test_trim.jl @@ -263,4 +263,157 @@ using AdaptiveArrayPools: get_typed_pool!, trim!, checkpoint!, rewind! @test trim!(pool, UInt8).slots_released == 1 end + @testset "trim! return value is type-stable (no Any-typed fields)" begin + # The summary is a public return value: if its count fields infer as `Any`, + # every caller that reads `summary.slots_released` inherits a dynamic + # dispatch. The no-arg full-pool form must stay fully concrete even though + # it folds over fixed slots (closure) and the `others` Vector{Any} (dynamic + # dispatch) — exactly the two inference hazards. + pool = AdaptiveArrayPool{0}() + checkpoint!(pool) + acquire!(pool, Float64, 100) + acquire!(pool, Int64, 50) + acquire!(pool, UInt8, 10) # exercises the `others` (Vector{Any}) path + rewind!(pool) + + Summary = @NamedTuple{ + slots_released::Int, wrappers_released::Int, + estimated_bytes_released::Int, gc_triggered::Bool, + } + + rt = only(Base.return_types(trim!, (typeof(pool),))) + @test isconcretetype(rt) + @test rt == Summary + @test only(Base.return_types(trim!, (typeof(pool), Type{Float64}))) == Summary + @test only(Base.return_types(trim!, (typeof(pool), Type{UInt8}))) == Summary + + # Runtime inference guards: throw if a callsite isn't inferrable-concrete. + # Cover all three positional forms — single-type and varargs route through + # `_trim_one_counts!`/`get_typed_pool!`, the forms most prone to widening. + @test (@inferred trim!(pool)) isa Summary + @test (@inferred trim!(pool, Float64)) isa Summary + @test (@inferred trim!(pool, Float64, Int64)) isa Summary + end + + @testset "varargs trim!(pool, T1, T2, ...) trims only the listed types" begin + # Mirrors reset!/checkpoint!/rewind!, which all expose a Type... form. + pool = AdaptiveArrayPool{0}() + checkpoint!(pool) + acquire!(pool, Float64, 100) + acquire!(pool, Int64, 50) + acquire!(pool, Float32, 25) + rewind!(pool) + + s = trim!(pool, Float64, Int64) # trim two; leave Float32 + + @test length(pool.float64.vectors) == 0 + @test length(pool.int64.vectors) == 0 + @test length(pool.float32.vectors) == 1 # untouched + @test s.slots_released == 2 # 1 + 1 + @test s.wrappers_released == 2 # 1-D acquires cache a wrapper too (N=1) + @test s.estimated_bytes_released >= (100 + 50) * sizeof(Float64) # full summary shape + @test s.gc_triggered == false + + # force_gc honored once across all listed types + pool2 = AdaptiveArrayPool{0}() + checkpoint!(pool2) + acquire!(pool2, Float64, 10) + acquire!(pool2, Int64, 10) + rewind!(pool2) + s2 = trim!(pool2, Float64, Int64; force_gc = true) + @test s2.slots_released == 2 + @test s2.gc_triggered == true + + # Return stays type-stable (concrete summary), like every other form. + Summary = @NamedTuple{ + slots_released::Int, wrappers_released::Int, + estimated_bytes_released::Int, gc_triggered::Bool, + } + @test only(Base.return_types(trim!, (typeof(pool), Type{Float64}, Type{Int64}))) == Summary + + # DisabledPool varargs is a zero-summary no-op. + @test trim!(DISABLED_CPU, Float64, Int64).slots_released == 0 + + # A never-used fallback type in the list is skipped, never created. + pool3 = AdaptiveArrayPool{0}() + @test !haskey(pool3.others, UInt16) + s3 = trim!(pool3, Float64, UInt16) + @test s3.slots_released == 0 + @test !haskey(pool3.others, UInt16) + end + + @testset "N-D wrapper cache rebuilds correctly after trim!" begin + # trim! drops cached N-D wrappers (each pins its backing buffer via :ref). + # The teardown is tested elsewhere; this pins the REBUILD side: a re-acquire + # of an N-D array after trim must allocate fresh and return a correct, + # writable array (the wrapper cache was truncated and must regrow). + pool = AdaptiveArrayPool{0}() + checkpoint!(pool) + a = acquire!(pool, Float64, 4, 5) # 2-D wrapper cached at arr_wrappers[2][1] + a .= 7.0 + rewind!(pool) + + trim!(pool) + @test length(pool.float64.vectors) == 0 + for w in pool.float64.arr_wrappers + w === nothing && continue + @test length(w) <= pool.float64.n_active # caches truncated + end + + checkpoint!(pool) + b = acquire!(pool, Float64, 3, 6) # re-acquire a DIFFERENT 2-D shape + @test size(b) == (3, 6) + @test ndims(b) == 2 + b .= 2.5 + @test all(b .== 2.5) + b[1, 1] = 9.0 # writable view into a fresh buffer + @test b[1, 1] == 9.0 + @test pool.float64.vectors[1][1] == 9.0 + rewind!(pool) + end + + @testset "trim! mid-scope leaves the active scope and a later rewind! consistent" begin + # trim! only drops inactive slots and must leave n_active + checkpoint + # stacks untouched, so trimming WHILE a scope is open does not corrupt the + # later rewind! of that scope. + pool = AdaptiveArrayPool{0}() + + checkpoint!(pool) # outer scope + a = acquire!(pool, Float64, 10) # active in outer + a .= 1.0 + + checkpoint!(pool) # inner scope + acquire!(pool, Float64, 5) + acquire!(pool, Float64, 7) + rewind!(pool) # slots 2,3 become inactive, retained + @test pool.float64.n_active == 1 + @test length(pool.float64.vectors) == 3 + + # Capture the checkpoint/depth state: trim! must touch ONLY tp.vectors and + # tp.arr_wrappers — never n_active, the checkpoint stacks, or pool depth. + depth_before = pool._current_depth + cpdepths_before = copy(pool.float64._checkpoint_depths) + cpn_active_before = copy(pool.float64._checkpoint_n_active) + + s = trim!(pool) # trim while the outer scope is still open + @test s.slots_released == 2 + @test pool.float64.n_active == 1 # active count untouched + @test length(pool.float64.vectors) == 1 # inactive dropped + @test pool._current_depth == depth_before # depth intact + @test pool.float64._checkpoint_depths == cpdepths_before # checkpoint stack intact + @test pool.float64._checkpoint_n_active == cpn_active_before + @test a == fill(1.0, 10) # active data intact... + a[2] = 4.0 # ...and still writable + @test pool.float64.vectors[1][2] == 4.0 + + rewind!(pool) # exit outer scope cleanly + @test pool.float64.n_active == 0 + + # Re-acquire after the trimmed-then-rewound cycle still works. + checkpoint!(pool) + c = acquire!(pool, Float64, 3) + @test length(c) == 3 + rewind!(pool) + end + end