diff --git a/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl b/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl index 7f83b12e..eaa305c6 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl @@ -33,6 +33,9 @@ using CUDA # State management (checkpoint!, rewind!, reset!, empty!) include("state.jl") + # Auto-compact: service the timer-flagged request at `@with_pool :cuda` entry + include("auto_compact.jl") + # Safety: poisoning, escape detection, borrow tracking include("debug.jl") diff --git a/ext/AdaptiveArrayPoolsCUDAExt/auto_compact.jl b/ext/AdaptiveArrayPoolsCUDAExt/auto_compact.jl new file mode 100644 index 00000000..7ddd0649 --- /dev/null +++ b/ext/AdaptiveArrayPoolsCUDAExt/auto_compact.jl @@ -0,0 +1,25 @@ +# ============================================================================== +# CUDA Auto-Compact: service a flagged CUDA pool at its `@with_pool :cuda` entry. +# ============================================================================== +# +# Parity with the CPU `_maybe_auto_compact!(::AdaptiveArrayPool)` in src/auto_compact.jl +# and the Metal mirror. The single global Timer (base module, CPU thread) sets this pool's +# `@atomic _compact_requested` flag via the shared WeakRef registry; the owner task runs +# the full `compact!` here, at the `_current_depth == 1` safepoint where nothing is borrowed. +# +# The scope-entry hook itself is already emitted by the SHARED `@with_pool` codegen (it +# expands `@with_pool :cuda` through the same `_generate_block_inner`), so this concrete +# method is all that's needed to make the hook act on CUDA pools — before it, the hook +# resolved to the `::Any` no-op fallback. Reuses the base `_run_auto_compact!` (CUDA +# `compact!` takes the same factor/shrink_to/min_bytes/active kwargs as CPU). +# +# `@inline` so the common case (flag clear) is a cheap inlined monotonic read + compare; +# the cold `_run_auto_compact!` stays a non-inlined call. `:monotonic` matches CPU: the +# flag is a one-way eventual-visibility signal with no cross-field ordering invariant. + +@inline function AdaptiveArrayPools._maybe_auto_compact!(pool::CuAdaptiveArrayPool) + if (@atomic :monotonic pool._compact_requested) && pool._current_depth == 1 + AdaptiveArrayPools._run_auto_compact!(pool) + end + return nothing +end diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index 1bf7ab67..7ff5714f 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -466,10 +466,15 @@ end # reference. All buffer ownership goes through `_update_cuda_wrapper_data!` # (unsafe_free! old + copy/retain new), so the old device buffer's refcount reaches # zero and CUDA frees it. -function AdaptiveArrayPools._compact_slot!(tp::CuTypedPool{T}, slot::Int, target::Int, used::Int) where {T} +# `copy_live` (parity with CPU `_compact_slot!`): copy the `used` live elements into the +# new device buffer. TRUE for ACTIVE slots — still held, must survive the shrink. FALSE +# for INACTIVE slots: their contents are dead (the next `acquire!` returns uninitialized +# device memory the caller fills), so the GPU→GPU `copyto!` is pure waste; skipping it +# makes inactive compaction allocate-only (no device copy / kernel launch). +function AdaptiveArrayPools._compact_slot!(tp::CuTypedPool{T}, slot::Int, target::Int, used::Int, copy_live::Bool) where {T} v = @inbounds tp.vectors[slot] nv = CuVector{T}(undef, target) - used > 0 && copyto!(nv, 1, v, 1, used) + copy_live && copyto!(nv, 1, v, 1, used) _update_cuda_wrapper_data!(v, nv) # backing takes nv's buffer (refcounted) setfield!(v, :dims, (target,)) for wrappers in tp.arr_wrappers @@ -487,12 +492,18 @@ end function AdaptiveArrayPools._maybe_compact_slot!(tp::CuTypedPool{T}, slot::Int, factor::Real, shrink_to::Real, min_bytes::Int) where {T} used = AdaptiveArrayPools._slot_used(tp, slot) used == 0 && return 0 + # RUNTIME_CHECK (S=1) invalidation poisons a released slot AND shrinks its logical length + # to 0 (device buffer retained) so an escaped view reads NaN/typemax. Compacting it would + # swap in fresh device storage and UNDO the poison — skip it (compacts once re-acquired). + # Parity with CPU `_maybe_compact_slot!`; `length == 0 && used > 0` ⟺ poisoned. + @inbounds(length(tp.vectors[slot])) == 0 && return 0 cap = AdaptiveArrayPools._slot_capacity(@inbounds tp.vectors[slot]) cap >= factor * used || return 0 target = max(used, ceil(Int, min(Float64(cap), shrink_to * used))) # clamp to [used, cap] reclaim = (cap - target) * _aligned_sizeof(T) reclaim >= min_bytes || return 0 - AdaptiveArrayPools._compact_slot!(tp, slot, target, used) + # Copy live data only for ACTIVE slots (1:n_active); inactive slots hold dead data. + AdaptiveArrayPools._compact_slot!(tp, slot, target, used, slot <= tp.n_active) return reclaim end diff --git a/ext/AdaptiveArrayPoolsCUDAExt/task_local_pool.jl b/ext/AdaptiveArrayPoolsCUDAExt/task_local_pool.jl index 2306652a..e81b79df 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/task_local_pool.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/task_local_pool.jl @@ -36,6 +36,10 @@ Uses `Dict{Int, CuAdaptiveArrayPool{RUNTIME_CHECK}}` in task-local storage, keye if pool === nothing pool = CuAdaptiveArrayPool() # Uses RUNTIME_CHECK for initial S pools[dev_id] = pool + # Register for auto-compaction (gated by the compile-time const; mirrors CPU + # get_task_local_pool). One entry per (task, device) pool. DCE'd when the + # Preference is off, so no registry traffic on the zero-overhead build. + AdaptiveArrayPools.AUTO_COMPACT && AdaptiveArrayPools.register_auto_compact!(pool) end return pool::CuAdaptiveArrayPool{RUNTIME_CHECK} diff --git a/ext/AdaptiveArrayPoolsCUDAExt/types.jl b/ext/AdaptiveArrayPoolsCUDAExt/types.jl index f3908540..e4998582 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/types.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/types.jl @@ -125,6 +125,11 @@ mutable struct CuAdaptiveArrayPool{S} <: AbstractArrayPool _pending_callsite::String _pending_return_site::String _borrow_log::Union{Nothing, IdDict{Any, String}} + + # Auto-compact request flag (parity with CPU AdaptiveArrayPool): set by the base + # module's background Timer sweep (a different thread), read + reset by the owner task + # at the `@with_pool :cuda` entry safepoint. Atomic for the cross-thread handoff. + @atomic _compact_requested::Bool end function CuAdaptiveArrayPool{S}() where {S} @@ -145,7 +150,8 @@ function CuAdaptiveArrayPool{S}() where {S} CUDA.deviceid(dev), "", # _pending_callsite "", # _pending_return_site - nothing # _borrow_log: lazily created when S >= 1 + nothing, # _borrow_log: lazily created when S >= 1 + false # _compact_requested: no pending auto-compact request ) end @@ -206,7 +212,8 @@ function _transfer_cuda_pool(::Val{V}, old::CuAdaptiveArrayPool) where {V} old.device_id, "", # _pending_callsite: reset "", # _pending_return_site: reset - V >= 1 ? IdDict{Any, String}() : nothing # _borrow_log + V >= 1 ? IdDict{Any, String}() : nothing, # _borrow_log + false # _compact_requested: reset on migration ) end diff --git a/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl b/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl index 4de72ca0..db1211a7 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl @@ -27,6 +27,7 @@ using Metal include("acquire.jl") include("task_local_pool.jl") include("state.jl") + include("auto_compact.jl") include("debug.jl") include("utils.jl") include("macros.jl") diff --git a/ext/AdaptiveArrayPoolsMetalExt/auto_compact.jl b/ext/AdaptiveArrayPoolsMetalExt/auto_compact.jl new file mode 100644 index 00000000..8c9b9d18 --- /dev/null +++ b/ext/AdaptiveArrayPoolsMetalExt/auto_compact.jl @@ -0,0 +1,25 @@ +# ============================================================================== +# Metal Auto-Compact: service a flagged Metal pool at its `@with_pool :metal` entry. +# ============================================================================== +# +# Parity with the CPU `_maybe_auto_compact!(::AdaptiveArrayPool)` in src/auto_compact.jl. +# The single global Timer (base module, CPU thread) sets this pool's `@atomic +# _compact_requested` flag via the shared WeakRef registry; the owner task runs the full +# `compact!` here, at the `_current_depth == 1` safepoint where nothing is borrowed. +# +# The scope-entry hook itself is already emitted by the SHARED `@with_pool` codegen +# (it expands `@with_pool :metal` through the same `_generate_block_inner`), so this +# concrete method is all that's needed to make the hook act on Metal pools — before it, +# the hook resolved to the `::Any` no-op fallback. Reuses the base `_run_auto_compact!` +# (Metal `compact!` takes the same factor/shrink_to/min_bytes/active kwargs as CPU). +# +# `@inline` so the common case (flag clear) is a cheap inlined monotonic read + compare; +# the cold `_run_auto_compact!` stays a non-inlined call. `:monotonic` matches CPU: the +# flag is a one-way eventual-visibility signal with no cross-field ordering invariant. + +@inline function AdaptiveArrayPools._maybe_auto_compact!(pool::MetalAdaptiveArrayPool) + if (@atomic :monotonic pool._compact_requested) && pool._current_depth == 1 + AdaptiveArrayPools._run_auto_compact!(pool) + end + return nothing +end diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 5d5e5083..add26a27 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -439,10 +439,15 @@ end # reference. All buffer ownership goes through `_update_metal_wrapper_data!` # (unsafe_free! old + copy/retain new), so the old device buffer's refcount reaches # zero and Metal frees it. -function AdaptiveArrayPools._compact_slot!(tp::MetalTypedPool{T, S}, slot::Int, target::Int, used::Int) where {T, S} +# `copy_live` (parity with CPU `_compact_slot!`): copy the `used` live elements into the +# new device buffer. TRUE for ACTIVE slots — still held, must survive the shrink. FALSE +# for INACTIVE slots: their contents are dead (the next `acquire!` returns uninitialized +# device memory the caller fills), so the GPU→GPU `copyto!` is pure waste; skipping it +# makes inactive compaction allocate-only (no device copy / kernel launch). +function AdaptiveArrayPools._compact_slot!(tp::MetalTypedPool{T, S}, slot::Int, target::Int, used::Int, copy_live::Bool) where {T, S} v = @inbounds tp.vectors[slot] nv = MtlArray{T, 1, S}(undef, target) - used > 0 && copyto!(nv, 1, v, 1, used) + copy_live && copyto!(nv, 1, v, 1, used) _update_metal_wrapper_data!(v, nv) # backing takes nv's buffer (refcounted) setfield!(v, :dims, (target,)) for wrappers in tp.arr_wrappers @@ -460,12 +465,18 @@ end function AdaptiveArrayPools._maybe_compact_slot!(tp::MetalTypedPool{T, S}, slot::Int, factor::Real, shrink_to::Real, min_bytes::Int) where {T, S} used = AdaptiveArrayPools._slot_used(tp, slot) used == 0 && return 0 + # RUNTIME_CHECK (R=1) invalidation poisons a released slot AND shrinks its logical length + # to 0 (device buffer retained) so an escaped view reads NaN/typemax. Compacting it would + # swap in fresh device storage and UNDO the poison — skip it (compacts once re-acquired). + # Parity with CPU `_maybe_compact_slot!`; `length == 0 && used > 0` ⟺ poisoned. + @inbounds(length(tp.vectors[slot])) == 0 && return 0 cap = AdaptiveArrayPools._slot_capacity(@inbounds tp.vectors[slot]) cap >= factor * used || return 0 target = max(used, ceil(Int, min(Float64(cap), shrink_to * used))) # clamp to [used, cap] reclaim = (cap - target) * sizeof(T) reclaim >= min_bytes || return 0 - AdaptiveArrayPools._compact_slot!(tp, slot, target, used) + # Copy live data only for ACTIVE slots (1:n_active); inactive slots hold dead data. + AdaptiveArrayPools._compact_slot!(tp, slot, target, used, slot <= tp.n_active) return reclaim end diff --git a/ext/AdaptiveArrayPoolsMetalExt/task_local_pool.jl b/ext/AdaptiveArrayPoolsMetalExt/task_local_pool.jl index 2e3e5796..f77984f0 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/task_local_pool.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/task_local_pool.jl @@ -37,6 +37,10 @@ Uses `Dict{UInt64, MetalAdaptiveArrayPool{RUNTIME_CHECK, METAL_STORAGE}}` in tas if pool === nothing pool = MetalAdaptiveArrayPool() # Uses RUNTIME_CHECK for initial R pools[dev_key] = pool + # Register for auto-compaction (gated by the compile-time const; mirrors CPU + # get_task_local_pool). One entry per (task, device) pool. DCE'd when the + # Preference is off, so no registry traffic on the zero-overhead build. + AdaptiveArrayPools.AUTO_COMPACT && AdaptiveArrayPools.register_auto_compact!(pool) end return pool::MetalAdaptiveArrayPool{RUNTIME_CHECK, METAL_STORAGE} diff --git a/ext/AdaptiveArrayPoolsMetalExt/types.jl b/ext/AdaptiveArrayPoolsMetalExt/types.jl index ab6156b6..bab1d661 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/types.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/types.jl @@ -121,6 +121,11 @@ mutable struct MetalAdaptiveArrayPool{R, S} <: AbstractArrayPool _pending_callsite::String _pending_return_site::String _borrow_log::Union{Nothing, IdDict{Any, String}} + + # Auto-compact request flag (parity with CPU AdaptiveArrayPool): set by the base + # module's background Timer sweep (a different thread), read + reset by the owner task + # at the `@with_pool :metal` entry safepoint. Atomic for the cross-thread handoff. + @atomic _compact_requested::Bool end function MetalAdaptiveArrayPool{R, S}() where {R, S} @@ -139,6 +144,7 @@ function MetalAdaptiveArrayPool{R, S}() where {R, S} "", # _pending_callsite "", # _pending_return_site nothing, # _borrow_log: lazily created when R >= 1 + false, # _compact_requested: no pending auto-compact request ) end diff --git a/src/AdaptiveArrayPools.jl b/src/AdaptiveArrayPools.jl index 53bd61c2..19349695 100644 --- a/src/AdaptiveArrayPools.jl +++ b/src/AdaptiveArrayPools.jl @@ -17,6 +17,7 @@ export @with_pool, @maybe_with_pool, @safe_with_pool, @safe_maybe_with_pool export STATIC_POOLING, MAYBE_POOLING, RUNTIME_CHECK export PoolEscapeError, EscapePoint, PoolMutationError, MutationPoint export checkpoint!, rewind!, reset!, trim!, compact! +export enable_auto_compact!, disable_auto_compact!, auto_compact_enabled, AUTO_COMPACT export get_task_local_cuda_pool, get_task_local_cuda_pools # CUDA (stubs, overridden by extension) export get_task_local_metal_pool, get_task_local_metal_pools # Metal (stubs, overridden by extension) @@ -33,6 +34,7 @@ export DisabledPool, DISABLED_CPU, pooling_enabled # Disabled pool support include("bitarray.jl") include("convenience.jl") include("state.jl") + include("auto_compact.jl") include("task_local_pool.jl") include("debug.jl") include("macros.jl") @@ -44,6 +46,7 @@ else include("legacy/bitarray.jl") include("convenience.jl") include("legacy/state.jl") + include("legacy/auto_compact.jl") include("task_local_pool.jl") include("debug.jl") include("macros.jl") diff --git a/src/auto_compact.jl b/src/auto_compact.jl new file mode 100644 index 00000000..25da8e8c --- /dev/null +++ b/src/auto_compact.jl @@ -0,0 +1,241 @@ +# ============================================================================== +# Auto-Compact: timer-driven background capacity compaction. +# Design: docs/plans/DESIGN_auto_compact.md +# +# A single global Timer periodically *requests* compaction by setting each live +# pool's `@atomic _compact_requested` flag (via a WeakRef registry). Each pool's +# owner task *executes* the request — the full `compact!` — at its next `@with_pool` +# ENTRY (the `_current_depth == 1` safepoint), where nothing is borrowed. +# +# Tier 1 (compile-time): `AUTO_COMPACT` gates whether the scope-entry flag check is +# emitted (macros.jl) and whether pools register (task_local_pool.jl). Default ON +# (~0.6 ns/entry); set the Preference false to compile it OUT (full DCE → zero cost). +# Tier 2 (runtime): `enable_auto_compact!` / `disable_auto_compact!`. +# ============================================================================== + +using Preferences: @load_preference + +""" + AUTO_COMPACT::Bool + +Compile-time master switch (Tier 1) for auto-compaction. Default `true`: the scope-entry +flag check is emitted (a ~0.6 ns inlined check per `@with_pool` entry), new task-local +pools auto-register, and `__init__` auto-starts the background Timer. Set `false` in +`LocalPreferences.toml` (then restart) to compile the feature OUT entirely — the hook is +dead-code-eliminated, restoring zero hot-path cost: + +```toml +[AdaptiveArrayPools] +auto_compact = false +``` + +`enable_auto_compact!` / `disable_auto_compact!` control the background Timer at runtime +(Tier 2). +""" +const AUTO_COMPACT = @load_preference("auto_compact", true)::Bool + +# Registry of pools eligible for auto-compaction. `WeakRef` so a pool dies with its +# task; dead refs are swap-removed on each sweep. All access under the lock below. +const _AUTO_COMPACT_REGISTRY = WeakRef[] +const _AUTO_COMPACT_LOCK = ReentrantLock() + +# Mutable global config, set by `enable_auto_compact!`, read by `_run_auto_compact!`. +# Mirrors the `compact!` knobs; `factor`/`shrink_to` are `Float64` for stable storage. +mutable struct _AutoCompactConfig + factor::Float64 + shrink_to::Float64 + min_bytes::Int + active::Bool +end +const _AUTO_COMPACT_CONFIG = _AutoCompactConfig(10.0, 1.5, 2^20, true) + +""" + register_auto_compact!(pool) + +Add `pool` to the auto-compact registry (under lock) so the background Timer's sweep +will flag it. Called from the task-local-pool slow path when `AUTO_COMPACT` is on +(one entry per task-local pool). The same registry serves every backend (CPU +`AdaptiveArrayPool`, GPU `MetalAdaptiveArrayPool`/`CuAdaptiveArrayPool`); each carries +its own `@atomic _compact_requested` flag and services it at its own scope entry. + +A pool *without* that flag (e.g. `DisabledPool`, or a foreign `AbstractArrayPool` +subtype) is silently skipped: registering one would make the sweep throw on its missing +`_compact_requested` and abort the cycle, starving every pool after it in the registry. +""" +function register_auto_compact!(pool::AbstractArrayPool) + # Only pools carrying the cross-thread flag can be swept (const-folds per concrete type). + hasfield(typeof(pool), :_compact_requested) || return nothing + lock(() -> push!(_AUTO_COMPACT_REGISTRY, WeakRef(pool)), _AUTO_COMPACT_LOCK) + return nothing +end + +# The Timer callback body: set every live pool's request flag and prune dead refs. +# Takes the timer (unused) so it can be handed straight to `Timer(cb, …)`. The caller +# wraps it in try/catch (Phase 2) — a throwing callback silently kills a Julia Timer. +function _auto_compact_sweep!(_timer) + lock(_AUTO_COMPACT_LOCK) do + i = 1 + while i <= length(_AUTO_COMPACT_REGISTRY) + pool = _AUTO_COMPACT_REGISTRY[i].value + if pool === nothing + _AUTO_COMPACT_REGISTRY[i] = _AUTO_COMPACT_REGISTRY[end] + pop!(_AUTO_COMPACT_REGISTRY) # dead → swap-remove + else + p = pool::AbstractArrayPool # CPU or any GPU backend pool + @atomic :monotonic p._compact_requested = true # request compaction + i += 1 + end + end + end + return nothing +end + +# Owner-side execution: reset the flag BEFORE the work (so a concurrent sweep's re-request +# is not lost), snapshot the config under the lock (it may be rewritten from another thread +# by `enable_auto_compact!`), then run the full `compact!`. Wrapped in try/catch so a +# background-maintenance failure (e.g. OOM allocating the new backing) is logged and the +# cycle skipped — it must never surface as a user-visible error at the `@with_pool` boundary. +# Returns nothing; the summary is discarded (no allocation). Shared across backends: +# `compact!` takes the same factor/shrink_to/min_bytes/active kwargs for every +# `AbstractArrayPool`, so one implementation drives CPU and GPU pools alike. +function _run_auto_compact!(pool::AbstractArrayPool) + @atomic :monotonic pool._compact_requested = false + factor, shrink_to, min_bytes, active = lock(_AUTO_COMPACT_LOCK) do + cfg = _AUTO_COMPACT_CONFIG + (cfg.factor, cfg.shrink_to, cfg.min_bytes, cfg.active) + end + try + compact!(pool; factor, shrink_to, min_bytes, active) + catch e + @warn "auto-compact failed; skipping this cycle" exception = (e, catch_backtrace()) maxlog = 3 + end + return nothing +end + +# Scope-ENTRY hook body, generated before the checkpoint in every `@with_pool` (gated by +# AUTO_COMPACT). At the `_current_depth == 1` safepoint (entering the outermost scope from +# global), run a flagged pool's compaction. Servicing the request at entry handles EVERY +# exit type of the previous scope — normal / early return / break / continue / exception — +# with a single hook point, and never runs compaction inside a `finally` during unwind. +# `@inline` so the common case (flag clear) is a cheap inlined monotonic read + compare; the +# cold `_run_auto_compact!` stays a non-inlined call. The `::Any` fallback keeps the SHARED +# `@with_pool` codegen safe for pools that don't opt in (e.g. `DisabledPool`, or a backend +# without its own method) — they no-op here. GPU backends define their own concrete method +# in their extension (e.g. `_maybe_auto_compact!(::MetalAdaptiveArrayPool)`), reusing the +# shared `_run_auto_compact!`. `:monotonic` suffices: the flag is a one-way eventual-visibility +# signal with no cross-field ordering invariant (cheaper than seq-cst on weak-ordered CPUs). +@inline _maybe_auto_compact!(::Any) = nothing +@inline function _maybe_auto_compact!(pool::AdaptiveArrayPool) + if (@atomic :monotonic pool._compact_requested) && pool._current_depth == 1 + _run_auto_compact!(pool) + end + return nothing +end + +# ── Timer lifecycle (Tier 2 runtime control) ───────────────────────────────── + +# Handle to the single global auto-compact Timer (`nothing` = stopped). +const _AUTO_COMPACT_TIMER = Ref{Union{Nothing, Timer}}(nothing) + +# Timer callback: the sweep wrapped so a throw can't silently kill the Timer (an +# uncaught error in a Julia Timer callback stops it permanently). +function _safe_auto_compact_sweep!(timer) + try + _auto_compact_sweep!(timer) + catch e + @warn "auto-compact sweep failed (timer continues)" exception = (e, catch_backtrace()) maxlog = 3 + end + return nothing +end + +# Lock-guarded timer lifecycle: the firing callback and concurrent enable!/disable! must not +# interleave a read-modify-write on `_AUTO_COMPACT_TIMER` (→ a leaked or double-closed Timer). +function _start_auto_compact_timer!(interval::Real) + lock(_AUTO_COMPACT_LOCK) do + _stop_auto_compact_timer_unlocked!() # replace any existing + _AUTO_COMPACT_TIMER[] = Timer(_safe_auto_compact_sweep!, interval; interval = interval) + end + return nothing +end + +_stop_auto_compact_timer!() = (lock(_stop_auto_compact_timer_unlocked!, _AUTO_COMPACT_LOCK); nothing) + +function _stop_auto_compact_timer_unlocked!() + t = _AUTO_COMPACT_TIMER[] + if t !== nothing + close(t) + _AUTO_COMPACT_TIMER[] = nothing + end + return nothing +end + +""" + enable_auto_compact!(; interval = 30.0, factor = 10, shrink_to = 1.5, + min_bytes = 2^20, active = true) + +Start (or restart) the background auto-compact `Timer` with the given config. Every +`interval` seconds it flags every registered pool; each pool then runs the full +`compact!` at its next `@with_pool` **entry** (the `_current_depth == 1` safepoint). + +Requires the `auto_compact` compile-time Preference (`AUTO_COMPACT`) for the scope-entry +hook and pool auto-registration. With the Preference **off**, the `@with_pool` hook is +compiled out, so even though this warns-then-starts the timer (and the timer keeps setting +flags), nothing ever *services* those flags — no automatic compaction happens, and +registering pools manually does not change that. Set the Preference and restart, or call +`compact!` yourself. + +!!! note "What gets compacted, when, and the zero-alloc trade-off" + Auto-compaction targets the **task-local pool** (`get_task_local_pool()`) — the pool + `@with_pool` uses and the only one that auto-registers. A hand-made `AdaptiveArrayPool()` + is **not** auto-registered (so not auto-compacted); `register_auto_compact!` it explicitly + to include it in the sweep, or just call `compact!` on it directly. The work + runs at a `@with_pool` **entry** (`_current_depth == 1`), so an idle task with no + `@with_pool` activity won't compact until it next enters a scope. Because a compaction + allocates the new right-sized backing(s), a `@with_pool` that services a pending request + is **not** strictly zero-allocation; set `auto_compact = false` (Preference) for a + guaranteed-zero-alloc hot path. + +See also [`disable_auto_compact!`](@ref), [`auto_compact_enabled`](@ref). +""" +function enable_auto_compact!(; + interval::Real = 30.0, factor::Real = 10, shrink_to::Real = 1.5, + min_bytes::Int = 2^20, active::Bool = true, + ) + AUTO_COMPACT || @warn "enable_auto_compact!: the `auto_compact` Preference is off — the `@with_pool` hook is compiled out, so the timer's flags are never serviced and no automatic compaction happens (registering pools does not change this). Set the Preference and restart to activate, or call `compact!` manually." maxlog = 1 + lock(_AUTO_COMPACT_LOCK) do + cfg = _AUTO_COMPACT_CONFIG + cfg.factor = Float64(factor) + cfg.shrink_to = Float64(shrink_to) + cfg.min_bytes = min_bytes + cfg.active = active + end + _start_auto_compact_timer!(interval) + return nothing +end + +""" + disable_auto_compact!() + +Stop the background auto-compact `Timer`. Idempotent. +""" +disable_auto_compact!() = _stop_auto_compact_timer!() + +""" + auto_compact_enabled() -> Bool + +Whether the background auto-compact `Timer` is currently running. +""" +auto_compact_enabled() = _AUTO_COMPACT_TIMER[] !== nothing + +# Auto-start the background timer when AUTO_COMPACT is on — but NOT while generating +# precompile output. A live Timer is an open libuv handle, so a precompile worker (e.g. +# precompiling a CUDA/Metal extension, which loads this package and runs `__init__`) +# would hang with "unfinished IO handles with timer events". `jl_generating_output()` is +# nonzero exactly during precompile/sysimage output. Also close the timer at process exit. +function __init__() + if AUTO_COMPACT && ccall(:jl_generating_output, Cint, ()) == 0 + enable_auto_compact!() + atexit(disable_auto_compact!) + end + return nothing +end diff --git a/src/legacy/auto_compact.jl b/src/legacy/auto_compact.jl new file mode 100644 index 00000000..ea8e02f0 --- /dev/null +++ b/src/legacy/auto_compact.jl @@ -0,0 +1,49 @@ +# ============================================================================== +# Auto-Compact — legacy (Julia < 1.12) no-op shims. +# +# The modern src/auto_compact.jl is NOT included on the legacy path, but the SHARED +# macros.jl and task_local_pool.jl reference AUTO_COMPACT / _maybe_auto_compact! / +# register_auto_compact!. Define them here as no-ops so the shared code compiles and the +# generated @with_pool scope-entry hook DCEs away (AUTO_COMPACT === false). The public +# enable/disable/enabled API stays callable across the whole supported Julia range. +# Upgrade to Julia 1.12+ for actual background auto-compaction. +# ============================================================================== + +# Compile-time OFF: constant-folds the shared @with_pool scope-entry hook to nothing. +const AUTO_COMPACT = false + +# Referenced (and DCE'd) by get_task_local_pool's `AUTO_COMPACT && register_auto_compact!`. +register_auto_compact!(pool) = nothing + +# Referenced (and DCE'd) by the @with_pool scope-entry hook. +@inline _maybe_auto_compact!(::Any) = nothing + +""" + enable_auto_compact!(; interval = 30.0, factor = 10, shrink_to = 1.5, + min_bytes = 2^20, active = true) + +No-op on Julia < 1.12 (the legacy pool architecture has no capacity compaction). Warns +once. Upgrade to Julia 1.12+ for background auto-compaction. +""" +function enable_auto_compact!(; + interval::Real = 30.0, factor::Real = 10, shrink_to::Real = 1.5, + min_bytes::Int = 2^20, active::Bool = true, + ) + @warn "enable_auto_compact! is a no-op on Julia < 1.12 (legacy pool architecture). " * + "Upgrade to Julia 1.12+ for background auto-compaction." maxlog = 1 + return nothing +end + +""" + disable_auto_compact!() + +No-op on Julia < 1.12. See [`enable_auto_compact!`](@ref). +""" +disable_auto_compact!() = nothing + +""" + auto_compact_enabled() -> Bool + +Always `false` on Julia < 1.12 (no auto-compaction). See [`enable_auto_compact!`](@ref). +""" +auto_compact_enabled() = false diff --git a/src/macros.jl b/src/macros.jl index 11eed70d..9a49e836 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -743,6 +743,7 @@ function _generate_block_inner(pool_name, expr, safe::Bool, source) if safe transformed_expr = _transform_return_stmts(transformed_expr, pool_name) return quote + $(_auto_compact_hook(pool_name)) $checkpoint_call try local _result = $(esc(transformed_expr)) @@ -768,6 +769,7 @@ function _generate_block_inner(pool_name, expr, safe::Bool, source) return quote local $(esc(entry_depth_var)) = $(esc(pool_name))._current_depth + $(_auto_compact_hook(pool_name)) $checkpoint_call local _result = $(esc(transformed_expr)) # Leaked scope cleanup BEFORE validation: if an inner @with_pool threw @@ -821,6 +823,7 @@ function _generate_function_inner(pool_name, expr, safe::Bool, source) if safe transformed_expr = _transform_return_stmts(transformed_expr, pool_name) return quote + $(_auto_compact_hook(pool_name)) $checkpoint_call try local _result = $(esc(transformed_expr)) @@ -846,6 +849,7 @@ function _generate_function_inner(pool_name, expr, safe::Bool, source) return quote local $(esc(entry_depth_var)) = $(esc(pool_name))._current_depth + $(_auto_compact_hook(pool_name)) $checkpoint_call local _result = $(esc(transformed_expr)) # Leaked scope cleanup BEFORE validation: if an inner @with_pool threw @@ -1587,6 +1591,16 @@ const _RUNTIME_CHECK_REF = GlobalRef(@__MODULE__, :_runtime_check) const _WARN_LEAKED_SCOPE_REF = GlobalRef(@__MODULE__, :_warn_leaked_scope) const _REWIND_REF = GlobalRef(@__MODULE__, :rewind!) const _LAZY_REWIND_REF = GlobalRef(@__MODULE__, :_lazy_rewind!) + +# Auto-compact scope-ENTRY hook, generated before the checkpoint in every `@with_pool`. +# Gated by the `AUTO_COMPACT` compile-time const (constant-folded → DCE'd to nothing when +# off → zero cost). Dispatches through `_maybe_auto_compact!` so non-CPU (GPU) pools safely +# no-op. Placed BEFORE the checkpoint, so `_current_depth == 1` ⇒ the OUTERMOST scope is +# being entered from global (nothing borrowed → safe). Servicing the flag at entry handles +# every exit type of the previous scope (return/break/throw), with no compaction in `finally`. +const _AUTO_COMPACT_REF = GlobalRef(@__MODULE__, :AUTO_COMPACT) +const _MAYBE_AUTO_COMPACT_REF = GlobalRef(@__MODULE__, :_maybe_auto_compact!) +_auto_compact_hook(pool_name) = :($_AUTO_COMPACT_REF && $_MAYBE_AUTO_COMPACT_REF($(esc(pool_name)))) const _TYPED_LAZY_REWIND_REF = GlobalRef(@__MODULE__, :_typed_lazy_rewind!) const _CAN_USE_TYPED_PATH_REF = GlobalRef(@__MODULE__, :_can_use_typed_path) const _TRACKED_MASK_REF = GlobalRef(@__MODULE__, :_tracked_mask_for_types) diff --git a/src/state.jl b/src/state.jl index 48a8891c..e9c0fbfe 100644 --- a/src/state.jl +++ b/src/state.jl @@ -879,10 +879,16 @@ end # Shrink one slot's backing to capacity `target` in place: allocate a smaller # `Memory`, copy the live `used` elements, swap the backing `Vector`'s `:ref`/`:size` # (identity preserved → views follow), then re-sync the slot's cached wrappers. -function _compact_slot!(tp::TypedPool{T}, slot::Int, target::Int, used::Int) where {T} +# `copy_live`: copy the `used` live elements into the new backing. TRUE for ACTIVE +# slots — their data is still held and must survive the shrink. FALSE for INACTIVE +# slots: their contents are dead (the next `acquire!` returns uninitialized memory the +# caller fills), so copying `used` elements is pure waste. Skipping it makes inactive +# compaction allocation-only (no `memcpy`) — the only zero-copy shrink available, since +# Julia's `Memory` cannot be shrunk in place (any real reclaim must allocate anew). +function _compact_slot!(tp::TypedPool{T}, slot::Int, target::Int, used::Int, copy_live::Bool) where {T} v = @inbounds tp.vectors[slot] nv = Vector{T}(undef, target) - copyto!(nv, 1, v, 1, used) + copy_live && copyto!(nv, 1, v, 1, used) setfield!(v, :size, (target,)) setfield!(v, :ref, getfield(nv, :ref)) for wrappers in tp.arr_wrappers @@ -902,6 +908,13 @@ _maybe_compact_slot!(::AbstractTypedPool, ::Int, ::Real, ::Real, ::Int) = 0 function _maybe_compact_slot!(tp::TypedPool{T}, slot::Int, factor::Real, shrink_to::Real, min_bytes::Int) where {T} used = _slot_used(tp, slot) used == 0 && return 0 + # RUNTIME_CHECK (S=1) invalidation poisons a released slot AND shrinks its backing to + # logical length 0 (keeping the retained Memory) so an escaped `acquire_view!` reads + # NaN/typemax. Such a slot still looks bloated (large capacity, nonzero recorded extent), + # but compacting it would swap in fresh storage and UNDO the poison. Skip it — it compacts + # normally once re-acquired. `length == 0 && used > 0` ⟺ poisoned (the two are set together; + # an un-checked slot keeps its logical length, so it is never length-0 with `used > 0`). + @inbounds(length(tp.vectors[slot])) == 0 && return 0 cap = _slot_capacity(@inbounds tp.vectors[slot]) cap >= factor * used || return 0 # Clamp the shrink goal to `[used, cap]`: never below the live size, and never @@ -910,7 +923,8 @@ function _maybe_compact_slot!(tp::TypedPool{T}, slot::Int, factor::Real, shrink_ target = max(used, ceil(Int, min(Float64(cap), shrink_to * used))) reclaim = (cap - target) * sizeof(T) reclaim >= min_bytes || return 0 - _compact_slot!(tp, slot, target, used) + # Copy live data only for ACTIVE slots (1:n_active); inactive slots hold dead data. + _compact_slot!(tp, slot, target, used, slot <= tp.n_active) return reclaim end diff --git a/src/task_local_pool.jl b/src/task_local_pool.jl index f1ec8f84..475f8e86 100644 --- a/src/task_local_pool.jl +++ b/src/task_local_pool.jl @@ -96,6 +96,8 @@ type-asserts directly to `AdaptiveArrayPool{RUNTIME_CHECK}`. # This branch is rarely taken (only once per Task) pool = AdaptiveArrayPool() task_local_storage(_POOL_KEY, pool) + # Register for background auto-compaction (DCE'd to nothing when AUTO_COMPACT is off). + AUTO_COMPACT && register_auto_compact!(pool) end return pool::AdaptiveArrayPool{RUNTIME_CHECK} diff --git a/src/types.jl b/src/types.jl index d4016c39..142a9e72 100644 --- a/src/types.jl +++ b/src/types.jl @@ -383,6 +383,11 @@ mutable struct AdaptiveArrayPool{S} <: AbstractArrayPool _pending_callsite::String # "" = no pending; set by macro before acquire _pending_return_site::String # "" = no pending; set by macro before validate _borrow_log::Union{Nothing, IdDict{Any, String}} # vector_obj => callsite string + + # Auto-compact request flag: set by the background Timer sweep (a different + # thread), read + reset by the owner task at the `_current_depth == 1` safepoint. + # Atomic for the cross-thread handoff; ignored when auto-compact is disabled. + @atomic _compact_requested::Bool end function AdaptiveArrayPool{S}() where {S} @@ -404,7 +409,8 @@ function AdaptiveArrayPool{S}() where {S} Int[0], # _others_ptr_bounds_checkpoints: sentinel "", # _pending_callsite: no pending "", # _pending_return_site: no pending - nothing # _borrow_log: lazily created at S=1 + nothing, # _borrow_log: lazily created at S=1 + false, # _compact_requested: no pending auto-compact request ) end @@ -456,7 +462,8 @@ _make_pool(runtime_check::Bool, old::AdaptiveArrayPool) = _make_pool(Int(runtime old._others_ptr_bounds_checkpoints, "", # _pending_callsite: reset "", # _pending_return_site: reset - S >= 1 ? IdDict{Any, String}() : nothing # _borrow_log + S >= 1 ? IdDict{Any, String}() : nothing, # _borrow_log + false, # _compact_requested: reset on migration ) level == 0 && return _new(Val(0)) return _new(Val(1)) diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl index 85f94e40..a7d997fd 100644 --- a/test/cuda/runtests.jl +++ b/test/cuda/runtests.jl @@ -51,5 +51,6 @@ else include("test_scope_depth_validation.jl") include("test_trim.jl") include("test_compact.jl") + include("test_auto_compact.jl") end end diff --git a/test/cuda/test_auto_compact.jl b/test/cuda/test_auto_compact.jl new file mode 100644 index 00000000..84d48dd7 --- /dev/null +++ b/test/cuda/test_auto_compact.jl @@ -0,0 +1,181 @@ +# CUDA Auto-Compact Tests +# ======================== +# Parity with the CPU suite (test/test_auto_compact.jl) and the Metal mirror, for CUDA. +# +# The single global Timer lives in the BASE module and runs on a CPU thread. It sweeps +# the shared `Vector{WeakRef}` registry and sets each live pool's `@atomic _compact_requested` +# flag — CPU and CUDA pools alike. Each pool's owner task runs the full `compact!` at its +# next scope entry (the `_current_depth == 1` safepoint): a CPU pool at `@with_pool`, a CUDA +# pool at `@with_pool :cuda`. The two are independent — one sweep flags both, and each +# services (and clears) only its own flag at its own boundary. +# +# AUTO_COMPACT defaults on, so `__init__` auto-starts the timer; we stop it up front so each +# testset is deterministic. Uses Float32 (GPU's primary type). + +import AdaptiveArrayPools as AAP + +AAP.disable_auto_compact!() # stop the __init__-started timer for deterministic tests + +@testset "CUDA auto_compact" begin + + # Helpers are scoped INSIDE this testset (a local scope) so they don't define top-level + # methods in `Main` that collide with the identically-named helpers in the CPU/Metal + # auto-compact test files when the full suite loads several backends' files together + # (otherwise: "WARNING: Method definition _clear_registry!() ... overwritten"). + _ccap(v) = Int(getfield(v, :maxsize) ÷ sizeof(eltype(v))) # CUDA device capacity (elements) + _cpucap(v) = length(getfield(v, :ref).mem) # CPU backing capacity (coexistence test) + _clear_registry!() = lock(() -> empty!(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + # An inactive, bloated Float32 CUDA slot (cap = `big` high-water, last extent = `small`). + function _cuda_inactive_bloated(big::Int, small::Int) + pool = CuAdaptiveArrayPool{0}() + checkpoint!(pool); a = acquire!(pool, Float32, big); fill!(a, 0.0f0); rewind!(pool) + checkpoint!(pool); b = acquire!(pool, Float32, small); fill!(b, 0.0f0); rewind!(pool) + return pool + end + + # ── Data layer: flag field + registry + run ────────────────────────────── + @testset "CUDA pool carries an atomic _compact_requested flag (default false)" begin + pool = CuAdaptiveArrayPool{0}() + @test (@atomic pool._compact_requested) == false + @atomic pool._compact_requested = true + @test (@atomic pool._compact_requested) == true + end + + @testset "registry sweep flags a registered CUDA pool" begin + _clear_registry!() + p = CuAdaptiveArrayPool{0}() + AAP.register_auto_compact!(p) + @test _registry_len() == 1 + AAP._auto_compact_sweep!(nothing) # the timer callback body + @test (@atomic p._compact_requested) == true + _clear_registry!() + end + + @testset "_run_auto_compact! resets the flag and compacts inactive CUDA bloat" begin + pool = _cuda_inactive_bloated(1_000_000, 100) + @atomic pool._compact_requested = true + @test _ccap(pool.float32.vectors[1]) >= 1_000_000 + + AAP._run_auto_compact!(pool) + + @test (@atomic pool._compact_requested) == false # reset (before the work) + @test _ccap(pool.float32.vectors[1]) < 1000 # device buffer shrunk in place + end + + # ── Scope-entry hook (dispatches on CuAdaptiveArrayPool) ────────────────── + @testset "_maybe_auto_compact!(::CuAdaptiveArrayPool) fires only at depth 1" begin + pool = _cuda_inactive_bloated(1_000_000, 100) + checkpoint!(pool) # depth → 2 + @atomic pool._compact_requested = true + AAP._maybe_auto_compact!(pool) # depth 2 ≠ 1 → must NOT act + @test (@atomic pool._compact_requested) == true # flag preserved for the outer exit + @test _ccap(pool.float32.vectors[1]) >= 1_000_000 # not compacted + rewind!(pool) # depth → 1 + AAP._maybe_auto_compact!(pool) # depth 1 → fires + @test (@atomic pool._compact_requested) == false # reset + @test _ccap(pool.float32.vectors[1]) < 1000 # compacted + end + + @testset "_maybe_auto_compact! is a no-op when the CUDA flag is clear" begin + pool = _cuda_inactive_bloated(1_000_000, 100) + @test (@atomic pool._compact_requested) == false + AAP._maybe_auto_compact!(pool) + @test _ccap(pool.float32.vectors[1]) >= 1_000_000 # untouched (no request) + end + + @testset "macro wires the gated scope-entry hook into @with_pool :cuda" begin + s = string( + @macroexpand @with_pool :cuda pool begin + x = acquire!(pool, Float32, 4) + sum(x) + end + ) + @test occursin("_maybe_auto_compact!", s) # hook emitted… + @test occursin("AUTO_COMPACT", s) # …gated by the compile-time const + end + + # ── Auto-registration via the CUDA task-local slow path ─────────────────── + @testset "get_task_local_cuda_pool auto-registers the pool (gate on)" begin + _clear_registry!() + fetch(Threads.@spawn get_task_local_cuda_pool()) # fresh task → creates + registers + @test _registry_len() >= 1 + _clear_registry!() + end + + @testset "end-to-end: bloat auto-compacts at a @with_pool :cuda boundary" begin + AAP.disable_auto_compact!() + _clear_registry!() + pool = get_task_local_cuda_pool() + empty!(pool) # clean slate + AAP.register_auto_compact!(pool) + @with_pool :cuda p begin # grow slot to 1M high-water + x = acquire!(p, Float32, 1_000_000); fill!(x, 0.0f0) + end + @with_pool :cuda p begin # reuse small → bloated, inactive + x = acquire!(p, Float32, 100); fill!(x, 0.0f0) + end + cap0 = _ccap(pool.float32.vectors[1]) + @test cap0 >= 1_000_000 + + # Set the request flag directly (deterministic). This exercises the unique end-to-end + # piece: the MACRO-emitted hook firing at a real @with_pool :cuda entry and dispatching + # to _maybe_auto_compact!(::CuAdaptiveArrayPool). + @atomic pool._compact_requested = true + @with_pool :cuda p begin # scope ENTRY at depth 1 → hook fires + acquire!(p, Float32, 4) + end + + @test (@atomic pool._compact_requested) == false # hook consumed the request + @test _ccap(pool.float32.vectors[1]) < cap0 # auto-compacted, no manual compact! + empty!(pool) + _clear_registry!() + end + + # ── Coexistence: a CPU pool and a CUDA pool, one Timer, independent compaction ── + # This pins the cross-backend behavior: the single CPU-side Timer flags BOTH a CPU and a + # CUDA pool in one sweep (structurally like two tasks each holding a pool — the registry + # tracks pools, not tasks). Each then compacts independently at its OWN scope boundary; + # servicing one never touches the other's flag or buffers. + @testset "CPU + CUDA coexist: one sweep flags both; each compacts independently" begin + AAP.disable_auto_compact!() + _clear_registry!() + + # CPU pool, inactive-bloated Float64. + cpu = AdaptiveArrayPool{0}() + checkpoint!(cpu); a = acquire!(cpu, Float64, 1_000_000); a .= 0.0; rewind!(cpu) + checkpoint!(cpu); b = acquire!(cpu, Float64, 100); b .= 0.0; rewind!(cpu) + # CUDA pool, inactive-bloated Float32. + cu = _cuda_inactive_bloated(1_000_000, 100) + + AAP.register_auto_compact!(cpu) + AAP.register_auto_compact!(cu) + @test _registry_len() == 2 + + cpu_cap0 = _cpucap(cpu.float64.vectors[1]) + cu_cap0 = _ccap(cu.float32.vectors[1]) + @test cpu_cap0 >= 1_000_000 + @test cu_cap0 >= 1_000_000 + + # ONE sweep (the single CPU-side Timer body) flags BOTH pools. + AAP._auto_compact_sweep!(nothing) + @test (@atomic cpu._compact_requested) == true + @test (@atomic cu._compact_requested) == true + + # Service the CPU pool at ITS boundary → only the CPU pool compacts; CUDA untouched. + AAP._maybe_auto_compact!(cpu) + @test (@atomic cpu._compact_requested) == false # CPU flag consumed + @test _cpucap(cpu.float64.vectors[1]) < cpu_cap0 # CPU compacted + @test (@atomic cu._compact_requested) == true # CUDA flag still pending + @test _ccap(cu.float32.vectors[1]) >= 1_000_000 # CUDA NOT compacted yet + + # Service the CUDA pool at ITS boundary → now the CUDA pool compacts, on its own. + AAP._maybe_auto_compact!(cu) + @test (@atomic cu._compact_requested) == false # CUDA flag consumed + @test _ccap(cu.float32.vectors[1]) < 1000 # CUDA compacted independently + + _clear_registry!() + end +end + +AAP.disable_auto_compact!() # leave the timer stopped for later CUDA test files diff --git a/test/cuda/test_compact.jl b/test/cuda/test_compact.jl index 8c845c11..d818b7e2 100644 --- a/test/cuda/test_compact.jl +++ b/test/cuda/test_compact.jl @@ -78,4 +78,63 @@ @test compact!(ext.DISABLED_CUDA, Float32, Int32) == ZERO end + @testset "active=true: per-slot copy decision (active preserved, inactive shrinks)" begin + # Parity with CPU `_compact_slot!`'s `copy_live`: the active-vs-inactive copy + # decision is PER SLOT (slot ≤ n_active ⇒ active ⇒ copy live data; else inactive ⇒ + # SKIP the GPU→GPU copy of dead data). Build two bloated slots, hold ONLY slot 1, + # compact both: both shrink, slot 1's held data survives, and slot 2 (inactive, + # copy skipped) still re-acquires correctly. (Inactive contents are dead → never + # asserted — the point is correctness despite the skipped copy.) + pool = CuAdaptiveArrayPool{0}() + checkpoint!(pool) # round 1: grow two slots to 1M + acquire!(pool, Float32, 1_000_000) + acquire!(pool, Float32, 1_000_000) + rewind!(pool) + checkpoint!(pool) # round 2: reuse both at 100 → bloated + acquire!(pool, Float32, 100) + acquire!(pool, Float32, 100) + rewind!(pool) + checkpoint!(pool) # round 3: hold ONLY slot 1 active + h1 = acquire!(pool, Float32, 100) + copyto!(h1, Float32.(1:100)) + tp = pool.float32 + @test tp.n_active == 1 # slot 1 active, slot 2 inactive + @test _ccap(tp.vectors[1]) >= 1_000_000 + @test _ccap(tp.vectors[2]) >= 1_000_000 + + s = compact!(pool; active = true) + + @test s.slots_compacted == 2 # BOTH slots shrink + @test _ccap(tp.vectors[1]) < 1000 # active slot 1 shrunk + @test _ccap(tp.vectors[2]) < 1000 # inactive slot 2 shrunk + @test Array(h1) == Float32.(1:100) # ACTIVE slot: live data preserved (copy_live) + rewind!(pool) + + checkpoint!(pool) # compacted INACTIVE slot re-acquires OK + acquire!(pool, Float32, 120) # slot 1 + r2 = acquire!(pool, Float32, 120) # slot 2 (was inactive, copy-skipped) + copyto!(r2, fill(5.0f0, 120)) + @test length(r2) == 120 + @test Array(r2) == fill(5.0f0, 120) # usable despite the skipped copy + rewind!(pool) + end + + @testset "RUNTIME_CHECK=1: compact! leaves poisoned inactive slots untouched" begin + # Parity with CPU: under S=1, releasing a slot poisons it AND shrinks its logical + # length to 0 (device buffer retained) so an escaped view reads NaN/typemax. compact! + # must NOT swap fresh device storage into such a slot — that would undo the poison. + pool = CuAdaptiveArrayPool{1}() + checkpoint!(pool); a = acquire!(pool, Float32, 1_000_000); fill!(a, 0.0f0); rewind!(pool) + checkpoint!(pool); b = acquire!(pool, Float32, 100); fill!(b, 0.0f0); rewind!(pool) + tp = pool.float32 + @test tp.n_active == 0 + @test length(tp.vectors[1]) == 0 # invalidated: logical length 0 + cap0 = _ccap(tp.vectors[1]) # retained device capacity (~1M) + @test cap0 >= 1_000_000 + + @test compact!(pool).slots_compacted == 0 # poisoned slot skipped + @test _ccap(tp.vectors[1]) == cap0 # device capacity (and poison) preserved + @test compact!(pool; active = false).slots_compacted == 0 + end + end diff --git a/test/metal/runtests.jl b/test/metal/runtests.jl index 63404b7d..4812ae95 100644 --- a/test/metal/runtests.jl +++ b/test/metal/runtests.jl @@ -53,5 +53,6 @@ else include("test_task_local_pool.jl") include("test_trim.jl") include("test_compact.jl") + include("test_auto_compact.jl") end end diff --git a/test/metal/test_auto_compact.jl b/test/metal/test_auto_compact.jl new file mode 100644 index 00000000..9602d67e --- /dev/null +++ b/test/metal/test_auto_compact.jl @@ -0,0 +1,181 @@ +# Metal Auto-Compact Tests +# ========================= +# Parity with the CPU suite (test/test_auto_compact.jl) for the Metal backend. +# +# The single global Timer lives in the BASE module and runs on a CPU thread. It sweeps +# the shared `Vector{WeakRef}` registry and sets each live pool's `@atomic _compact_requested` +# flag — CPU and Metal pools alike. Each pool's owner task runs the full `compact!` at its +# next scope entry (the `_current_depth == 1` safepoint): a CPU pool at `@with_pool`, a Metal +# pool at `@with_pool :metal`. The two are independent — one sweep flags both, and each +# services (and clears) only its own flag at its own boundary. +# +# AUTO_COMPACT defaults on, so `__init__` auto-starts the timer; we stop it up front so each +# testset is deterministic. Uses Float32 (Metal's primary type; no Float64 on Metal). + +import AdaptiveArrayPools as AAP + +AAP.disable_auto_compact!() # stop the __init__-started timer for deterministic tests + +@testset "Metal auto_compact" begin + + # Helpers are scoped INSIDE this testset (a local scope) so they don't define top-level + # methods in `Main` that collide with the identically-named helpers in the CPU/CUDA + # auto-compact test files when the full suite loads several backends' files together + # (otherwise: "WARNING: Method definition _clear_registry!() ... overwritten"). + _mcap(v) = Int(getfield(v, :maxsize) ÷ sizeof(eltype(v))) # Metal device capacity (elements) + _cpucap(v) = length(getfield(v, :ref).mem) # CPU backing capacity (coexistence test) + _clear_registry!() = lock(() -> empty!(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + # An inactive, bloated Float32 Metal slot (cap = `big` high-water, last extent = `small`). + function _metal_inactive_bloated(big::Int, small::Int) + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + checkpoint!(pool); a = acquire!(pool, Float32, big); fill!(a, 0.0f0); rewind!(pool) + checkpoint!(pool); b = acquire!(pool, Float32, small); fill!(b, 0.0f0); rewind!(pool) + return pool + end + + # ── Data layer: flag field + registry + run ────────────────────────────── + @testset "Metal pool carries an atomic _compact_requested flag (default false)" begin + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + @test (@atomic pool._compact_requested) == false + @atomic pool._compact_requested = true + @test (@atomic pool._compact_requested) == true + end + + @testset "registry sweep flags a registered Metal pool" begin + _clear_registry!() + p = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + AAP.register_auto_compact!(p) + @test _registry_len() == 1 + AAP._auto_compact_sweep!(nothing) # the timer callback body + @test (@atomic p._compact_requested) == true + _clear_registry!() + end + + @testset "_run_auto_compact! resets the flag and compacts inactive Metal bloat" begin + pool = _metal_inactive_bloated(1_000_000, 100) + @atomic pool._compact_requested = true + @test _mcap(pool.float32.vectors[1]) >= 1_000_000 + + AAP._run_auto_compact!(pool) + + @test (@atomic pool._compact_requested) == false # reset (before the work) + @test _mcap(pool.float32.vectors[1]) < 1000 # device buffer shrunk in place + end + + # ── Scope-entry hook (dispatches on MetalAdaptiveArrayPool) ─────────────── + @testset "_maybe_auto_compact!(::MetalAdaptiveArrayPool) fires only at depth 1" begin + pool = _metal_inactive_bloated(1_000_000, 100) + checkpoint!(pool) # depth → 2 + @atomic pool._compact_requested = true + AAP._maybe_auto_compact!(pool) # depth 2 ≠ 1 → must NOT act + @test (@atomic pool._compact_requested) == true # flag preserved for the outer exit + @test _mcap(pool.float32.vectors[1]) >= 1_000_000 # not compacted + rewind!(pool) # depth → 1 + AAP._maybe_auto_compact!(pool) # depth 1 → fires + @test (@atomic pool._compact_requested) == false # reset + @test _mcap(pool.float32.vectors[1]) < 1000 # compacted + end + + @testset "_maybe_auto_compact! is a no-op when the Metal flag is clear" begin + pool = _metal_inactive_bloated(1_000_000, 100) + @test (@atomic pool._compact_requested) == false + AAP._maybe_auto_compact!(pool) + @test _mcap(pool.float32.vectors[1]) >= 1_000_000 # untouched (no request) + end + + @testset "macro wires the gated scope-entry hook into @with_pool :metal" begin + s = string( + @macroexpand @with_pool :metal pool begin + x = acquire!(pool, Float32, 4) + sum(x) + end + ) + @test occursin("_maybe_auto_compact!", s) # hook emitted… + @test occursin("AUTO_COMPACT", s) # …gated by the compile-time const + end + + # ── Auto-registration via the Metal task-local slow path ────────────────── + @testset "get_task_local_metal_pool auto-registers the pool (gate on)" begin + _clear_registry!() + fetch(Threads.@spawn get_task_local_metal_pool()) # fresh task → creates + registers + @test _registry_len() >= 1 + _clear_registry!() + end + + @testset "end-to-end: bloat auto-compacts at a @with_pool :metal boundary" begin + AAP.disable_auto_compact!() + _clear_registry!() + pool = get_task_local_metal_pool() + empty!(pool) # clean slate + AAP.register_auto_compact!(pool) + @with_pool :metal p begin # grow slot to 1M high-water + x = acquire!(p, Float32, 1_000_000); fill!(x, 0.0f0) + end + @with_pool :metal p begin # reuse small → bloated, inactive + x = acquire!(p, Float32, 100); fill!(x, 0.0f0) + end + cap0 = _mcap(pool.float32.vectors[1]) + @test cap0 >= 1_000_000 + + # Set the request flag directly (deterministic). This exercises the unique end-to-end + # piece: the MACRO-emitted hook firing at a real @with_pool :metal entry and dispatching + # to _maybe_auto_compact!(::MetalAdaptiveArrayPool). + @atomic pool._compact_requested = true + @with_pool :metal p begin # scope ENTRY at depth 1 → hook fires + acquire!(p, Float32, 4) + end + + @test (@atomic pool._compact_requested) == false # hook consumed the request + @test _mcap(pool.float32.vectors[1]) < cap0 # auto-compacted, no manual compact! + empty!(pool) + _clear_registry!() + end + + # ── Coexistence: a CPU pool and a Metal pool, one Timer, independent compaction ── + # This pins the behavior asked about: the single CPU-side Timer flags BOTH a CPU and a + # Metal pool in one sweep (structurally like two tasks each holding a pool — the registry + # tracks pools, not tasks). Each then compacts independently at its OWN scope boundary; + # servicing one never touches the other's flag or buffers. + @testset "CPU + Metal coexist: one sweep flags both; each compacts independently" begin + AAP.disable_auto_compact!() + _clear_registry!() + + # CPU pool, inactive-bloated Float64. + cpu = AdaptiveArrayPool{0}() + checkpoint!(cpu); a = acquire!(cpu, Float64, 1_000_000); a .= 0.0; rewind!(cpu) + checkpoint!(cpu); b = acquire!(cpu, Float64, 100); b .= 0.0; rewind!(cpu) + # Metal pool, inactive-bloated Float32. + met = _metal_inactive_bloated(1_000_000, 100) + + AAP.register_auto_compact!(cpu) + AAP.register_auto_compact!(met) + @test _registry_len() == 2 + + cpu_cap0 = _cpucap(cpu.float64.vectors[1]) + met_cap0 = _mcap(met.float32.vectors[1]) + @test cpu_cap0 >= 1_000_000 + @test met_cap0 >= 1_000_000 + + # ONE sweep (the single CPU-side Timer body) flags BOTH pools. + AAP._auto_compact_sweep!(nothing) + @test (@atomic cpu._compact_requested) == true + @test (@atomic met._compact_requested) == true + + # Service the CPU pool at ITS boundary → only the CPU pool compacts; Metal untouched. + AAP._maybe_auto_compact!(cpu) + @test (@atomic cpu._compact_requested) == false # CPU flag consumed + @test _cpucap(cpu.float64.vectors[1]) < cpu_cap0 # CPU compacted + @test (@atomic met._compact_requested) == true # Metal flag still pending + @test _mcap(met.float32.vectors[1]) >= 1_000_000 # Metal NOT compacted yet + + # Service the Metal pool at ITS boundary → now the Metal pool compacts, on its own. + AAP._maybe_auto_compact!(met) + @test (@atomic met._compact_requested) == false # Metal flag consumed + @test _mcap(met.float32.vectors[1]) < 1000 # Metal compacted independently + + _clear_registry!() + end +end + +AAP.disable_auto_compact!() # leave the timer stopped for later Metal test files diff --git a/test/metal/test_compact.jl b/test/metal/test_compact.jl index 4a2fd693..8f925814 100644 --- a/test/metal/test_compact.jl +++ b/test/metal/test_compact.jl @@ -81,4 +81,63 @@ @test compact!(ext.DISABLED_METAL, Float32, Int32) == ZERO end + @testset "active=true: per-slot copy decision (active preserved, inactive shrinks)" begin + # Parity with CPU `_compact_slot!`'s `copy_live`: the active-vs-inactive copy + # decision is PER SLOT (slot ≤ n_active ⇒ active ⇒ copy live data; else inactive ⇒ + # SKIP the GPU→GPU copy of dead data). Build two bloated slots, hold ONLY slot 1, + # compact both: both shrink, slot 1's held data survives, and slot 2 (inactive, + # copy skipped) still re-acquires correctly. (Inactive contents are dead → never + # asserted — the point is correctness despite the skipped copy.) + pool = MetalAdaptiveArrayPool{0, METAL_STORAGE}() + checkpoint!(pool) # round 1: grow two slots to 1M + acquire!(pool, Float32, 1_000_000) + acquire!(pool, Float32, 1_000_000) + rewind!(pool) + checkpoint!(pool) # round 2: reuse both at 100 → bloated + acquire!(pool, Float32, 100) + acquire!(pool, Float32, 100) + rewind!(pool) + checkpoint!(pool) # round 3: hold ONLY slot 1 active + h1 = acquire!(pool, Float32, 100) + copyto!(h1, Float32.(1:100)) + tp = pool.float32 + @test tp.n_active == 1 # slot 1 active, slot 2 inactive + @test _mcap(tp.vectors[1]) >= 1_000_000 + @test _mcap(tp.vectors[2]) >= 1_000_000 + + s = compact!(pool; active = true) + + @test s.slots_compacted == 2 # BOTH slots shrink + @test _mcap(tp.vectors[1]) < 1000 # active slot 1 shrunk + @test _mcap(tp.vectors[2]) < 1000 # inactive slot 2 shrunk + @test Array(h1) == Float32.(1:100) # ACTIVE slot: live data preserved (copy_live) + rewind!(pool) + + checkpoint!(pool) # compacted INACTIVE slot re-acquires OK + acquire!(pool, Float32, 120) # slot 1 + r2 = acquire!(pool, Float32, 120) # slot 2 (was inactive, copy-skipped) + copyto!(r2, fill(5.0f0, 120)) + @test length(r2) == 120 + @test Array(r2) == fill(5.0f0, 120) # usable despite the skipped copy + rewind!(pool) + end + + @testset "RUNTIME_CHECK=1: compact! leaves poisoned inactive slots untouched" begin + # Parity with CPU: under R=1, releasing a slot poisons it AND shrinks its logical + # length to 0 (device buffer retained) so an escaped view reads NaN/typemax. compact! + # must NOT swap fresh device storage into such a slot — that would undo the poison. + pool = MetalAdaptiveArrayPool{1, METAL_STORAGE}() + checkpoint!(pool); a = acquire!(pool, Float32, 1_000_000); fill!(a, 0.0f0); rewind!(pool) + checkpoint!(pool); b = acquire!(pool, Float32, 100); fill!(b, 0.0f0); rewind!(pool) + tp = pool.float32 + @test tp.n_active == 0 + @test length(tp.vectors[1]) == 0 # invalidated: logical length 0 + cap0 = _mcap(tp.vectors[1]) # retained device capacity (~1M) + @test cap0 >= 1_000_000 + + @test compact!(pool).slots_compacted == 0 # poisoned slot skipped + @test _mcap(tp.vectors[1]) == cap0 # device capacity (and poison) preserved + @test compact!(pool; active = false).slots_compacted == 0 + end + end diff --git a/test/runtests.jl b/test/runtests.jl index a638f433..8ec5c0b1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,6 +3,10 @@ using AdaptiveArrayPools using AdaptiveArrayPools: get_typed_pool! import AdaptiveArrayPools: checkpoint!, rewind! +# AUTO_COMPACT defaults on, so __init__ auto-starts the background timer. Stop it so the +# suite is deterministic; testsets that exercise auto-compaction enable it explicitly. +AdaptiveArrayPools.disable_auto_compact!() + # Version-specific helpers (always defined, even for ARGS path) @static if VERSION >= v"1.12-" _test_nd_cache_preserved(tp) = !isempty(tp.arr_wrappers) @@ -24,6 +28,7 @@ else include("test_state.jl") include("test_trim.jl") include("test_compact.jl") + include("test_auto_compact.jl") include("test_multidimensional.jl") include("test_macros.jl") include("test_task_local_pool.jl") @@ -56,6 +61,7 @@ else include("test_state.jl") include("test_trim_legacy.jl") include("test_compact_legacy.jl") + include("test_auto_compact_legacy.jl") include("test_multidimensional.jl") include("test_macros.jl") include("test_task_local_pool.jl") diff --git a/test/test_auto_compact.jl b/test/test_auto_compact.jl new file mode 100644 index 00000000..6305b9ac --- /dev/null +++ b/test/test_auto_compact.jl @@ -0,0 +1,324 @@ +using Test +using Base.CoreLogging: with_logger, NullLogger # from Base — no `Logging` dep needed +using AdaptiveArrayPools +using AdaptiveArrayPools: checkpoint!, rewind! +import AdaptiveArrayPools as AAP + +# ============================================================================== +# Auto-compact — timer-driven background capacity compaction. +# Design: docs/plans/DESIGN_auto_compact.md +# +# ONE global Timer sweeps a Vector{WeakRef} registry, setting each live pool's +# `@atomic _compact_requested` flag. Each pool's owner task runs the full compact! +# at its next `@with_pool` entry (the `_current_depth == 1` safepoint; the macro hook). +# +# AUTO_COMPACT defaults to `true`, so __init__ auto-starts the timer and the +# @with_pool hook is live. We stop that timer up front so each testset controls it +# explicitly and the suite stays deterministic. +# ============================================================================== + +AAP.disable_auto_compact!() # stop the __init__-started timer for deterministic tests + +@testset "auto_compact" begin + + # Helpers are scoped INSIDE this testset (a local scope) so they don't define top-level + # methods in `Main` that collide with identically-named helpers in other test files when + # the full suite loads them together — e.g. `_cap` is also defined at the top of + # test_compact.jl (otherwise: "WARNING: Method definition _cap(Any) ... overwritten"). + _cap(v) = length(getfield(v, :ref).mem) + _clear_registry!() = lock(() -> empty!(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + # An inactive, bloated Float64 slot (cap = `big` high-water, last extent = `small`). + function _inactive_bloated(big::Int, small::Int) + pool = AdaptiveArrayPool{0}() + checkpoint!(pool); a = acquire!(pool, Float64, big); a .= 0.0; rewind!(pool) + checkpoint!(pool); b = acquire!(pool, Float64, small); b .= 0.0; rewind!(pool) + return pool + end + + # ── Data layer: flag field + registry + run ────────────────────────────── + @testset "pool carries an atomic _compact_requested flag (default false)" begin + pool = AdaptiveArrayPool{0}() + @test (@atomic pool._compact_requested) == false + @atomic pool._compact_requested = true + @test (@atomic pool._compact_requested) == true + end + + @testset "registry: sweep flags every live pool and prunes dead refs" begin + _clear_registry!() + p1 = AdaptiveArrayPool{0}() + p2 = AdaptiveArrayPool{0}() + AAP.register_auto_compact!(p1) + AAP.register_auto_compact!(p2) + lock(() -> push!(AAP._AUTO_COMPACT_REGISTRY, WeakRef(nothing)), AAP._AUTO_COMPACT_LOCK) + @test _registry_len() == 3 + + AAP._auto_compact_sweep!(nothing) # the timer callback body + + @test (@atomic p1._compact_requested) == true + @test (@atomic p2._compact_requested) == true + @test _registry_len() == 2 # dead WeakRef pruned + _clear_registry!() + end + + @testset "register_auto_compact! skips flagless pools (no sweep poison)" begin + # A pool without the `_compact_requested` flag (DisabledPool) must NOT be registered: + # otherwise the sweep would throw on it and starve every later pool in the registry. + _clear_registry!() + AAP.register_auto_compact!(DISABLED_CPU) # DisabledPool: no @atomic flag + @test _registry_len() == 0 # silently skipped, not registered + # And a real pool registered alongside is still swept normally (no poison). + p = AdaptiveArrayPool{0}() + AAP.register_auto_compact!(p) + AAP._auto_compact_sweep!(nothing) + @test (@atomic p._compact_requested) == true + _clear_registry!() + end + + @testset "_run_auto_compact! resets the flag and compacts inactive bloat" begin + pool = _inactive_bloated(1_000_000, 100) + @atomic pool._compact_requested = true + @test _cap(pool.float64.vectors[1]) >= 1_000_000 + + AAP._run_auto_compact!(pool) + + @test (@atomic pool._compact_requested) == false # reset (before the work) + @test _cap(pool.float64.vectors[1]) == 150 # shrunk to ceil(1.5 * 100) + end + + @testset "_run_auto_compact! uses the global config defaults" begin + @test AAP._AUTO_COMPACT_CONFIG.factor == 10 + @test AAP._AUTO_COMPACT_CONFIG.shrink_to == 1.5 + @test AAP._AUTO_COMPACT_CONFIG.min_bytes == 2^20 + @test AAP._AUTO_COMPACT_CONFIG.active == true + end + + # ── Timer lifecycle (enable/disable). AUTO_COMPACT is on → enable does NOT warn ─ + @testset "enable! sets config and starts a timer that flags registered pools" begin + _clear_registry!() + AAP.disable_auto_compact!() + pool = AdaptiveArrayPool{0}() + AAP.register_auto_compact!(pool) + + AAP.enable_auto_compact!( + interval = 0.1, factor = 20, shrink_to = 2.0, min_bytes = 4096, active = false, + ) + @test AAP.auto_compact_enabled() == true + @test AAP._AUTO_COMPACT_CONFIG.factor == 20.0 + @test AAP._AUTO_COMPACT_CONFIG.shrink_to == 2.0 + @test AAP._AUTO_COMPACT_CONFIG.min_bytes == 4096 + @test AAP._AUTO_COMPACT_CONFIG.active == false + + sleep(0.35) # ~3 ticks @ 0.1s + @test (@atomic pool._compact_requested) == true + + AAP.disable_auto_compact!() + @test AAP.auto_compact_enabled() == false + AAP._AUTO_COMPACT_CONFIG.factor = 10.0 # restore defaults + AAP._AUTO_COMPACT_CONFIG.shrink_to = 1.5 + AAP._AUTO_COMPACT_CONFIG.min_bytes = 2^20 + AAP._AUTO_COMPACT_CONFIG.active = true + _clear_registry!() + end + + @testset "disable! stops the timer and is idempotent" begin + _clear_registry!() + pool = AdaptiveArrayPool{0}() + AAP.register_auto_compact!(pool) + AAP.enable_auto_compact!(interval = 0.1) + AAP.disable_auto_compact!() + @test AAP.auto_compact_enabled() == false + @atomic pool._compact_requested = false + sleep(0.3) + @test (@atomic pool._compact_requested) == false # no ticks after disable + AAP.disable_auto_compact!() # idempotent: must not throw + @test AAP.auto_compact_enabled() == false + _clear_registry!() + end + + @testset "enable! replaces an existing timer (closes the old one)" begin + AAP.disable_auto_compact!() + AAP.enable_auto_compact!(interval = 0.1) + t1 = AAP._AUTO_COMPACT_TIMER[] + AAP.enable_auto_compact!(interval = 0.2) + t2 = AAP._AUTO_COMPACT_TIMER[] + @test t1 !== t2 + @test !isopen(t1) # old timer closed + @test isopen(t2) + AAP.disable_auto_compact!() + end + + @testset "a throwing sweep is caught — the timer survives" begin + _clear_registry!() + bogus = Ref(42) # heap object, NOT a pool + lock(() -> push!(AAP._AUTO_COMPACT_REGISTRY, WeakRef(bogus)), AAP._AUTO_COMPACT_LOCK) + # The _safe wrapper catches the sweep's throw (a non-pool in the registry) and @warns + # so the timer survives. Silence that EXPECTED warning (NullLogger) so it doesn't read + # as a real failure in the suite log; we still assert the wrapper does not rethrow. + with_logger(NullLogger()) do + @test (AAP._safe_auto_compact_sweep!(nothing); true) # _safe wrapper swallows the throw + end + @test bogus[] == 42 # keep `bogus` alive past the sweep + _clear_registry!() + end + + # ── Scope-entry hook ───────────────────────────────────────────────────── + @testset "_maybe_auto_compact! fires only at _current_depth == 1" begin + pool = _inactive_bloated(1_000_000, 100) + checkpoint!(pool) # depth → 2 + @atomic pool._compact_requested = true + AAP._maybe_auto_compact!(pool) # depth 2 ≠ 1 → must NOT act + @test (@atomic pool._compact_requested) == true # flag preserved for the outer exit + @test _cap(pool.float64.vectors[1]) >= 1_000_000 # not compacted + rewind!(pool) # depth → 1 + AAP._maybe_auto_compact!(pool) # depth 1 → fires + @test (@atomic pool._compact_requested) == false # reset + @test _cap(pool.float64.vectors[1]) == 150 # compacted + end + + @testset "_maybe_auto_compact! is a no-op when the flag is clear" begin + pool = _inactive_bloated(1_000_000, 100) + @test (@atomic pool._compact_requested) == false + AAP._maybe_auto_compact!(pool) + @test _cap(pool.float64.vectors[1]) >= 1_000_000 # untouched (no request) + end + + @testset "_maybe_auto_compact! no-ops on non-CPU pools (GPU fallback)" begin + @test AAP._maybe_auto_compact!(42) === nothing # ::Any fallback, no throw + @test AAP._maybe_auto_compact!("not a pool") === nothing + end + + @testset "macro wires the gated scope-entry hook into @with_pool" begin + s = string( + @macroexpand @with_pool pool begin + x = acquire!(pool, Float64, 4) + sum(x) + end + ) + @test occursin("_maybe_auto_compact!", s) # hook emitted… + @test occursin("AUTO_COMPACT", s) # …gated by the compile-time const + end + + # ── Default-on integration (AUTO_COMPACT === true) ─────────────────────── + @testset "AUTO_COMPACT is on by default" begin + @test AAP.AUTO_COMPACT == true + end + + @testset "__init__ auto-starts the timer when AUTO_COMPACT is on" begin + AAP.disable_auto_compact!() + @test AAP.auto_compact_enabled() == false + AAP.__init__() # re-run module init + @test AAP.auto_compact_enabled() == true # auto-started + AAP.disable_auto_compact!() + end + + @testset "get_task_local_pool auto-registers the pool (gate on)" begin + _clear_registry!() + fetch(Threads.@spawn get_task_local_pool()) # fresh task → creates + registers + @test _registry_len() >= 1 + _clear_registry!() + end + + @testset "end-to-end: bloat auto-compacts at a @with_pool boundary" begin + AAP.disable_auto_compact!() + _clear_registry!() + pool = get_task_local_pool() + empty!(pool) # clean slate + AAP.register_auto_compact!(pool) + @with_pool p begin # grow slot to 1M high-water + x = acquire!(p, Float64, 1_000_000); x .= 0.0 + end + @with_pool p begin # reuse small → bloated, inactive + x = acquire!(p, Float64, 100); x .= 0.0 + end + cap0 = _cap(pool.float64.vectors[1]) + @test cap0 >= 1_000_000 + + # Set the request flag directly (deterministic — the timer→flag link is covered by + # the "enable! flags registered pools" test). This exercises the unique end-to-end + # piece: the MACRO-emitted hook firing at a real @with_pool exit and running compact!. + @atomic pool._compact_requested = true + @with_pool p begin # scope ENTRY at depth 1 → hook fires + acquire!(p, Float64, 4) + end + + @test (@atomic pool._compact_requested) == false # hook consumed the request + @test _cap(pool.float64.vectors[1]) < cap0 # auto-compacted, no manual compact! + empty!(pool) + _clear_registry!() + end + + @testset "request is serviced for scopes that early return / break" begin + # The scope-ENTRY hook services a pending flag when the NEXT @with_pool is entered at + # depth 1, regardless of how the previous scope exited — so a scope that early-returns + # or breaks (which skip the normal exit path) still gets compacted at its entry. + pool = get_task_local_pool() + AAP.disable_auto_compact!() + + empty!(pool) # early RETURN + @with_pool p begin + x = acquire!(p, Float64, 1_000_000); x .= 0.0 + end + @with_pool p begin + x = acquire!(p, Float64, 100); x .= 0.0 + end + cap0 = _cap(pool.float64.vectors[1]) + @atomic pool._compact_requested = true + ret() = @with_pool p begin + acquire!(p, Float64, 4) + return 42 + end + @test ret() == 42 + @test (@atomic pool._compact_requested) == false # serviced at the next scope ENTRY + @test _cap(pool.float64.vectors[1]) < cap0 + + empty!(pool) # early BREAK + @with_pool p begin + x = acquire!(p, Float64, 1_000_000); x .= 0.0 + end + @with_pool p begin + x = acquire!(p, Float64, 100); x .= 0.0 + end + cap1 = _cap(pool.float64.vectors[1]) + @atomic pool._compact_requested = true + for _ in 1:3 + @with_pool p begin + acquire!(p, Float64, 4) + break + end + end + @test (@atomic pool._compact_requested) == false # serviced at the scope ENTRY + @test _cap(pool.float64.vectors[1]) < cap1 + empty!(pool) + end + + @testset "@safe_with_pool fires the hook during exception unwind (and preserves it)" begin + pool = get_task_local_pool() + AAP.disable_auto_compact!() + empty!(pool) + @with_pool p begin + x = acquire!(p, Float64, 1_000_000); x .= 0.0 + end + @with_pool p begin + x = acquire!(p, Float64, 100); x .= 0.0 + end + cap0 = _cap(pool.float64.vectors[1]) + @atomic pool._compact_requested = true + threw = false + try + @safe_with_pool p begin + acquire!(p, Float64, 4) + error("boom") + end + catch e + threw = true + @test e isa ErrorException # original exception NOT masked + end + @test threw + @test (@atomic pool._compact_requested) == false # hook fired at scope ENTRY (before the throw) + @test _cap(pool.float64.vectors[1]) < cap0 + empty!(pool) + end +end + +AAP.disable_auto_compact!() # leave the timer stopped for later test files diff --git a/test/test_auto_compact_legacy.jl b/test/test_auto_compact_legacy.jl new file mode 100644 index 00000000..7ae1f345 --- /dev/null +++ b/test/test_auto_compact_legacy.jl @@ -0,0 +1,32 @@ +using Test +using AdaptiveArrayPools +using AdaptiveArrayPools: get_task_local_pool + +# Legacy path (Julia < 1.12): the auto-compact API is defined but a NO-OP (the legacy +# pool architecture has no capacity compaction). `const AUTO_COMPACT = false` makes the +# shared @with_pool scope-entry hook DCE away, and the public enable/disable/enabled API +# stays callable across the supported Julia range. Mirrors test_compact_legacy.jl. + +@testset "auto-compact (legacy no-op, Julia < 1.12)" begin + @test AdaptiveArrayPools.AUTO_COMPACT == false + @test auto_compact_enabled() == false + + # enable! warns once and stays a no-op; disable!/enabled are safe and idempotent. + @test (@test_logs (:warn,) match_mode = :any enable_auto_compact!()) === nothing + @test auto_compact_enabled() == false + @test enable_auto_compact!(interval = 5.0, factor = 20, active = false) === nothing + @test disable_auto_compact!() === nothing + @test auto_compact_enabled() == false + + # The internal hooks the shared code references are no-ops too. + @test AdaptiveArrayPools.register_auto_compact!(get_task_local_pool()) === nothing + @test AdaptiveArrayPools._maybe_auto_compact!(get_task_local_pool()) === nothing + + # @with_pool still works on legacy (the hook DCEs to nothing; no auto-compaction). + r = @with_pool pool begin + x = acquire!(pool, Float64, 4) + x .= 1.0 + sum(x) + end + @test r == 4.0 +end diff --git a/test/test_compact.jl b/test/test_compact.jl index 3f6ec136..7f99c361 100644 --- a/test/test_compact.jl +++ b/test/test_compact.jl @@ -304,6 +304,47 @@ const _ZERO_COMPACT = (; slots_compacted = 0, bytes_reclaimed = 0, gc_triggered @test _cap(pool.float64.vectors[1]) == 150 end + @testset "active=true: per-slot copy decision (active data preserved, inactive shrinks)" begin + # The active-vs-inactive copy decision is PER SLOT (slot ≤ n_active ⇒ active). + # Build two Float64 slots BOTH bloated, then hold ONLY slot 1 active while slot + # 2 stays inactive. active=true must shrink BOTH, copy-preserve the held slot's + # live data, and the inactive slot must still re-acquire correctly afterward. + # (Inactive contents are dead — never asserted — so compaction may skip copying + # them; this pins that the live/dead choice keys off n_active per slot.) + pool = AdaptiveArrayPool{0}() + checkpoint!(pool) # round 1: grow two slots to 1M + acquire!(pool, Float64, 1_000_000) + acquire!(pool, Float64, 1_000_000) + rewind!(pool) + checkpoint!(pool) # round 2: reuse both at 100 → bloated + acquire!(pool, Float64, 100) + acquire!(pool, Float64, 100) + rewind!(pool) + checkpoint!(pool) # round 3: hold ONLY slot 1 active + h1 = acquire!(pool, Float64, 100) + h1 .= 1.0:100.0 + tp = pool.float64 + @test tp.n_active == 1 # slot 1 active, slot 2 inactive + @test _cap(tp.vectors[1]) >= 1_000_000 + @test _cap(tp.vectors[2]) >= 1_000_000 + + s = compact!(pool; active = true) + + @test s.slots_compacted == 2 # BOTH slots shrink + @test _cap(tp.vectors[1]) == 150 # active slot 1 shrunk + @test _cap(tp.vectors[2]) == 150 # inactive slot 2 shrunk + @test h1 == collect(1.0:100.0) # ACTIVE slot: live data preserved + rewind!(pool) + + checkpoint!(pool) # compacted INACTIVE slot re-acquires OK + acquire!(pool, Float64, 120) # slot 1 + r2 = acquire!(pool, Float64, 120) # slot 2 (was inactive, compacted) + r2 .= 5.0 + @test length(r2) == 120 + @test all(r2 .== 5.0) + rewind!(pool) + end + @testset "varargs compact!(pool, T...) compacts each listed type" begin pool = _inactive_bloated_pool(1_000_000, 100) # Float64 inactive bloat checkpoint!(pool); ai = acquire!(pool, Int64, 1_000_000); ai .= 0; rewind!(pool) @@ -401,4 +442,31 @@ const _ZERO_COMPACT = (; slots_compacted = 0, bytes_reclaimed = 0, gc_triggered @test compact!(pool; shrink_to = 1.0e6).slots_compacted == 0 end + @testset "RUNTIME_CHECK=1: compact! leaves poisoned inactive slots untouched" begin + # Under S=1, releasing a slot poisons it (NaN/typemax) AND shrinks its backing to + # logical length 0, so an escaped acquire_view! view reads poison instead of live + # data. compact! must NOT swap fresh storage into such a slot — that would undo the + # poison and let the escaped view read plausible memory. The slot stays untouched + # (retained capacity + poison preserved) until it is re-acquired. + pool = AdaptiveArrayPool{1}() + checkpoint!(pool); a = acquire!(pool, Float64, 1_000_000); a .= 0.0; rewind!(pool) + checkpoint!(pool); b = acquire!(pool, Float64, 100); b .= 0.0; rewind!(pool) + tp = pool.float64 + @test tp.n_active == 0 + @test length(tp.vectors[1]) == 0 # invalidated: logical length 0 + cap0 = _cap(tp.vectors[1]) # retained Memory capacity (~1M) + @test cap0 >= 1_000_000 + + s = compact!(pool) # default active=true; all slots inactive here + @test s.slots_compacted == 0 # poisoned slot skipped (no escape-detection undo) + @test _cap(tp.vectors[1]) == cap0 # capacity (and poison) preserved + + @test compact!(pool; active = false).slots_compacted == 0 # explicit inactive scan: also skipped + + # Re-acquiring clears the length-0 marker; the slot then compacts normally. + checkpoint!(pool); c = acquire!(pool, Float64, 100); c .= 0.0; rewind!(pool) + @test length(tp.vectors[1]) == 0 # re-released → invalidated again + @test compact!(pool).slots_compacted == 0 # still skipped while poisoned + end + end