From bc9a74530036bd5a3fb1a73c2e1aefb656031bf4 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 14:13:14 -0700 Subject: [PATCH 01/14] perf(compact!): skip dead-data copy when shrinking inactive slots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_compact_slot!` copied the `used` live elements into the new backing unconditionally. For INACTIVE slots (slot > n_active) that data is dead — the next `acquire!` returns uninitialized memory the caller fills — so the copy is pure waste. Gate it on a per-slot active flag: active slots (1:n_active) still copy to preserve held data; inactive slots allocate the right-sized buffer with no copy. Eliminates compact!'s only non-sub-µs case: reclaiming a large-last-extent inactive slot drops from ~138µs to ~166ns (alloc-only). Active-slot data preservation is unchanged — guarded by existing tests plus a new per-slot characterization test (active slot 1 data preserved, inactive slot 2 shrinks + re-acquires correctly under active=true). Julia's `Memory` cannot be shrunk in place, so alloc+free stays the only real reclaim path; this just drops the needless memcpy. CPU only; the Metal/CUDA `_compact_slot!` methods carry the same opportunity (follow-up). --- src/state.jl | 13 ++++++++++--- test/test_compact.jl | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/src/state.jl b/src/state.jl index 48a8891c..a630a69b 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 @@ -910,7 +916,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/test/test_compact.jl b/test/test_compact.jl index 3f6ec136..00cdd17f 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) From 605d78e7ab9bd8f8b4cc1ff91f1d20dad76397cb Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 14:49:08 -0700 Subject: [PATCH 02/14] feat(auto-compact): pool flag field + registry + run (phase 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Foundation for timer-driven background compaction — the data layer only (no Timer, no macro hook yet; phases 2-3 follow). Every piece is exercised directly by test_auto_compact.jl. - AdaptiveArrayPool gains an `@atomic _compact_requested::Bool` field: the cross-thread handoff cell (Timer writes, owner reads+resets at its safepoint). 1 byte; both positional constructors updated; reset on pool migration. - New src/auto_compact.jl: * AUTO_COMPACT — compile-time Preference (Tier 1), default off → zero cost. * _AUTO_COMPACT_REGISTRY (Vector{WeakRef}) + _AUTO_COMPACT_LOCK. * _AutoCompactConfig + defaults mirroring compact! (factor=10, shrink_to=1.5, min_bytes=2^20, active=true). * register_auto_compact!, _auto_compact_sweep! (flag live pools, prune dead), _run_auto_compact! (reset-then-compact; depth guard is the caller's). CPU/modern path only. --- src/AdaptiveArrayPools.jl | 1 + src/auto_compact.jl | 95 +++++++++++++++++++++++++++++++++++++++ src/types.jl | 11 ++++- test/runtests.jl | 1 + test/test_auto_compact.jl | 76 +++++++++++++++++++++++++++++++ 5 files changed, 182 insertions(+), 2 deletions(-) create mode 100644 src/auto_compact.jl create mode 100644 test/test_auto_compact.jl diff --git a/src/AdaptiveArrayPools.jl b/src/AdaptiveArrayPools.jl index 53bd61c2..79fc160e 100644 --- a/src/AdaptiveArrayPools.jl +++ b/src/AdaptiveArrayPools.jl @@ -33,6 +33,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") diff --git a/src/auto_compact.jl b/src/auto_compact.jl new file mode 100644 index 00000000..1d751e49 --- /dev/null +++ b/src/auto_compact.jl @@ -0,0 +1,95 @@ +# ============================================================================== +# 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 +# `_current_depth == 1` scope-exit safepoint, where nothing is borrowed. +# +# Tier 1 (compile-time): `AUTO_COMPACT` gates whether the scope-exit flag check is +# emitted (macros.jl) and whether pools register (task_local_pool.jl). Default off +# → zero hot-path 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. When `false` (default), the +scope-exit flag check is not emitted and pools are not registered — zero hot-path +cost. Set in `LocalPreferences.toml` (then restart) to compile it in: + +```toml +[AdaptiveArrayPools] +auto_compact = true +``` + +When on, `enable_auto_compact!` / `disable_auto_compact!` control the background +Timer at runtime (Tier 2). +""" +const AUTO_COMPACT = @load_preference("auto_compact", false)::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). +""" +function register_auto_compact!(pool) + 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::AdaptiveArrayPool + @atomic 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), then run the full configured `compact!`. Called at the +# `_current_depth == 1` safepoint by the macro hook (Phase 3) — the depth guard is the +# caller's. Returns nothing; the summary is discarded (no allocation). +function _run_auto_compact!(pool::AdaptiveArrayPool) + @atomic pool._compact_requested = false + cfg = _AUTO_COMPACT_CONFIG + compact!( + pool; + factor = cfg.factor, shrink_to = cfg.shrink_to, + min_bytes = cfg.min_bytes, active = cfg.active, + ) + return nothing +end 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/runtests.jl b/test/runtests.jl index a638f433..58ba1e97 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,6 +24,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") diff --git a/test/test_auto_compact.jl b/test/test_auto_compact.jl new file mode 100644 index 00000000..03febcfe --- /dev/null +++ b/test/test_auto_compact.jl @@ -0,0 +1,76 @@ +using Test +using AdaptiveArrayPools +using AdaptiveArrayPools: checkpoint!, rewind! +import AdaptiveArrayPools as AAP + +# ============================================================================== +# Auto-compact — timer-driven background capacity compaction. +# Design: docs/plans/DESIGN_auto_compact.md +# +# Architecture: 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 `_current_depth == 1` scope-exit safepoint. +# +# Phase 1 covers the data layer (flag field + registry + config + run), with no +# timer and no macro hook yet — every piece is exercised directly. +# ============================================================================== + +# Allocated backing capacity (Memory length), mirroring test_compact.jl. +_cap(v) = length(getfield(v, :ref).mem) + +# 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 + +_clear_registry!() = lock(() -> empty!(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) +_registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) + +@testset "auto_compact — Phase 1: flag field + registry + run" begin + + @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) + # inject a guaranteed-dead entry to exercise pruning + 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 "_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 +end From acda0fced8346dabcedef88d869030027109b29b Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 14:55:17 -0700 Subject: [PATCH 03/14] feat(auto-compact): timer lifecycle + enable/disable API (phase 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runtime control (Tier 2) for the background compaction timer: - enable_auto_compact!(; interval=30, factor, shrink_to, min_bytes, active): set config + (re)start the single global Timer. Warns when the `auto_compact` Preference is off (hot-path compiled out) but still starts the timer, so manual registration / testing works. - disable_auto_compact!() — stop the timer (idempotent). - auto_compact_enabled() -> Bool. - _safe_auto_compact_sweep!: try/catch around the sweep so a throwing callback cannot silently kill the Timer. - __init__ auto-starts with defaults when AUTO_COMPACT is on (dead-code-eliminated when off → no timer on a normal load). No macro/hot-path integration yet (phase 3). Tests cover start→flag→stop lifecycle, config-setting, timer replacement, disable idempotency, and the error-swallowing wrapper. --- src/auto_compact.jl | 80 +++++++++++++++++++++++++++++++++++++++ test/test_auto_compact.jl | 73 +++++++++++++++++++++++++++++++++++ 2 files changed, 153 insertions(+) diff --git a/src/auto_compact.jl b/src/auto_compact.jl index 1d751e49..79c1d890 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -93,3 +93,83 @@ function _run_auto_compact!(pool::AdaptiveArrayPool) ) 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 + +function _start_auto_compact_timer!(interval::Real) + _stop_auto_compact_timer!() # replace any existing + _AUTO_COMPACT_TIMER[] = Timer(_safe_auto_compact_sweep!, interval; interval = interval) + return nothing +end + +function _stop_auto_compact_timer!() + 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 `_current_depth == 1` scope-exit safepoint. + +Requires the `auto_compact` compile-time Preference (`AUTO_COMPACT`) for the scope-exit +hook and pool auto-registration. With the Preference off this warns and the timer still +runs but stays inert unless you `register_auto_compact!` pools manually. + +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 scope-exit hook is compiled out and pools won't auto-register. The timer runs but stays inert unless you register pools manually. Set the Preference and restart to activate." + cfg = _AUTO_COMPACT_CONFIG + cfg.factor = Float64(factor) + cfg.shrink_to = Float64(shrink_to) + cfg.min_bytes = min_bytes + cfg.active = active + _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 when the Preference opts in. Dead-code-eliminated when `AUTO_COMPACT` is +# off (the default), so package load starts no timer unless the user set the Preference. +function __init__() + AUTO_COMPACT && enable_auto_compact!() + return nothing +end diff --git a/test/test_auto_compact.jl b/test/test_auto_compact.jl index 03febcfe..3a7713cc 100644 --- a/test/test_auto_compact.jl +++ b/test/test_auto_compact.jl @@ -74,3 +74,76 @@ _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPA @test AAP._AUTO_COMPACT_CONFIG.active == true end end + +@testset "auto_compact — Phase 2: timer lifecycle + enable/disable" begin + # In the default test build the `auto_compact` Preference is off, so + # `enable_auto_compact!` emits a guidance @warn — asserted via @test_logs (which + # also captures it, keeping output clean). The timer still starts, so the + # sweep→flag mechanism is exercisable against a manually-registered pool. + + @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) + + @test_logs (:warn,) match_mode = :any 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 + # restore config defaults for the other testsets + AAP._AUTO_COMPACT_CONFIG.factor = 10.0 + 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) + @test_logs (:warn,) match_mode = :any 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!() + @test_logs (:warn,) match_mode = :any AAP.enable_auto_compact!(interval = 0.1) + t1 = AAP._AUTO_COMPACT_TIMER[] + @test_logs (:warn,) match_mode = :any 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 `::AdaptiveArrayPool` assertion in the sweep throws on the bogus entry; + # the _safe wrapper must swallow it so the Timer would keep running. + @test (AAP._safe_auto_compact_sweep!(nothing); true) + @test bogus[] == 42 # keep `bogus` alive past the sweep + _clear_registry!() + end +end From 1146f39383f0e0a179a7713bad1839530b9605d5 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 15:08:16 -0700 Subject: [PATCH 04/14] feat(auto-compact): scope-exit hook + pool registration (phase 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the timer's flag into the hot path, fully gated by AUTO_COMPACT: - _maybe_auto_compact!(pool): @inline scope-exit hook body — at the _current_depth == 1 safepoint, if the flag is set, run the compaction (reset-then-compact). An ::Any fallback no-ops for non-CPU pools, keeping the SHARED @with_pool codegen safe (GPU auto-compact is a future phase). - macros.jl: every @with_pool exit (block + function, safe + non-safe) emits `AUTO_COMPACT && _maybe_auto_compact!(pool)` after the normal rewind. Injected once per generator by wrapping rewind_call, so early return/break sites are untouched (best-effort: catches the next normal exit at depth 1). - get_task_local_pool registers new task-local pools (AUTO_COMPACT-gated). Zero overhead when off: AUTO_COMPACT is a const false, so the hook is constant-folded out of the typed IR and LLVM entirely (verified: no trace in either, and @allocated == 0 across a @with_pool loop). ~0.6 ns inlined check when on. CPU/modern path only. --- src/auto_compact.jl | 14 ++++++++++++ src/macros.jl | 13 +++++++++++ src/task_local_pool.jl | 2 ++ test/test_auto_compact.jl | 47 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 76 insertions(+) diff --git a/src/auto_compact.jl b/src/auto_compact.jl index 79c1d890..2684f37c 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -94,6 +94,20 @@ function _run_auto_compact!(pool::AdaptiveArrayPool) return nothing end +# Scope-exit hook body, generated into every `@with_pool` exit (gated by AUTO_COMPACT). +# At the `_current_depth == 1` safepoint, run a flagged pool's compaction. `@inline` so +# the common case (flag clear) is a cheap inlined atomic-read + compare; the cold +# `_run_auto_compact!` stays a non-inlined call. The `::Any` fallback keeps the SHARED +# `@with_pool` codegen safe for non-CPU (GPU) pools — they no-op here, and a future phase +# adds their own methods. +@inline _maybe_auto_compact!(::Any) = nothing +@inline function _maybe_auto_compact!(pool::AdaptiveArrayPool) + if (@atomic 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). diff --git a/src/macros.jl b/src/macros.jl index 11eed70d..7a163424 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -752,6 +752,7 @@ function _generate_block_inner(pool_name, expr, safe::Bool, source) _result finally $rewind_call + $(_auto_compact_hook(pool_name)) end end else @@ -783,6 +784,7 @@ function _generate_block_inner(pool_name, expr, safe::Bool, source) $_validate_pool_return(_result, $(esc(pool_name))) end $rewind_call + $(_auto_compact_hook(pool_name)) _result end end @@ -830,6 +832,7 @@ function _generate_function_inner(pool_name, expr, safe::Bool, source) _result finally $rewind_call + $(_auto_compact_hook(pool_name)) end end else @@ -861,6 +864,7 @@ function _generate_function_inner(pool_name, expr, safe::Bool, source) $_validate_pool_return(_result, $(esc(pool_name))) end $rewind_call + $(_auto_compact_hook(pool_name)) _result end end @@ -1587,6 +1591,15 @@ 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-exit hook, generated after the normal rewind in every `@with_pool` +# exit. 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 AFTER the rewind, so `_current_depth == 1` ⇒ the outermost scope +# just closed (nothing borrowed → even active compaction is safe). +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/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/test/test_auto_compact.jl b/test/test_auto_compact.jl index 3a7713cc..fc325805 100644 --- a/test/test_auto_compact.jl +++ b/test/test_auto_compact.jl @@ -147,3 +147,50 @@ end _clear_registry!() end end + +@testset "auto_compact — Phase 3: scope-exit hook + registration gate" begin + + @testset "_maybe_auto_compact! fires only at _current_depth == 1" begin + pool = _inactive_bloated(1_000_000, 100) # inactive bloat, depth back to 1 + 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-exit 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 + + @testset "registration gate is off when AUTO_COMPACT is compiled out (default)" begin + if !AAP.AUTO_COMPACT + _clear_registry!() + fetch(Threads.@spawn get_task_local_pool()) # new task → creates a pool + @test _registry_len() == 0 # not auto-registered (gate off) + end + end +end From d88be839767ec926709e3147a8e9c449360daa79 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 15:45:24 -0700 Subject: [PATCH 05/14] feat(auto-compact): default-on + legacy shims + exports (phase 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make auto-compaction the default and complete the cross-version API. - AUTO_COMPACT Preference now defaults to `true` (opt-out): out of the box the scope-exit hook is compiled in (~0.6 ns/exit), task-local pools auto-register, and __init__ auto-starts the 30s timer. Set `auto_compact = false` to compile it out entirely (full DCE → zero hot-path cost). Removed a stray debug println. - src/legacy/auto_compact.jl (new): `const AUTO_COMPACT = false` + register_/_maybe_ no-ops + enable!/disable!/enabled no-op shims, so the SHARED macros.jl / task_local_pool.jl references resolve and the @with_pool hook DCEs on < 1.12. Fixes the `UndefVarError: AUTO_COMPACT` seen on the legacy path. - Export enable_auto_compact!, disable_auto_compact!, auto_compact_enabled, AUTO_COMPACT. - runtests.jl disables the auto-started timer up front for deterministic tests. - test_auto_compact.jl reworked for default-on: drops obsolete warn-asserts; adds __init__-autostart, auto-registration, and a deterministic end-to-end test that the macro hook compacts at a @with_pool boundary. New test_auto_compact_legacy.jl. CPU/modern path runs the real hook; < 1.12 is a portable no-op. --- src/AdaptiveArrayPools.jl | 2 + src/auto_compact.jl | 21 ++++--- src/legacy/auto_compact.jl | 49 +++++++++++++++ test/runtests.jl | 5 ++ test/test_auto_compact.jl | 102 ++++++++++++++++++++----------- test/test_auto_compact_legacy.jl | 32 ++++++++++ 6 files changed, 168 insertions(+), 43 deletions(-) create mode 100644 src/legacy/auto_compact.jl create mode 100644 test/test_auto_compact_legacy.jl diff --git a/src/AdaptiveArrayPools.jl b/src/AdaptiveArrayPools.jl index 79fc160e..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) @@ -45,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 index 2684f37c..ae7f53f1 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -8,8 +8,9 @@ # `_current_depth == 1` scope-exit safepoint, where nothing is borrowed. # # Tier 1 (compile-time): `AUTO_COMPACT` gates whether the scope-exit flag check is -# emitted (macros.jl) and whether pools register (task_local_pool.jl). Default off -# → zero hot-path cost. Tier 2 (runtime): `enable_auto_compact!` / `disable_auto_compact!`. +# emitted (macros.jl) and whether pools register (task_local_pool.jl). Default ON +# (~0.6 ns/exit); 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 @@ -17,19 +18,21 @@ using Preferences: @load_preference """ AUTO_COMPACT::Bool -Compile-time master switch (Tier 1) for auto-compaction. When `false` (default), the -scope-exit flag check is not emitted and pools are not registered — zero hot-path -cost. Set in `LocalPreferences.toml` (then restart) to compile it in: +Compile-time master switch (Tier 1) for auto-compaction. Default `true`: the scope-exit +flag check is emitted (a ~0.6 ns inlined check per `@with_pool` exit), 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 = true +auto_compact = false ``` -When on, `enable_auto_compact!` / `disable_auto_compact!` control the background -Timer at runtime (Tier 2). +`enable_auto_compact!` / `disable_auto_compact!` control the background Timer at runtime +(Tier 2). """ -const AUTO_COMPACT = @load_preference("auto_compact", false)::Bool +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. diff --git a/src/legacy/auto_compact.jl b/src/legacy/auto_compact.jl new file mode 100644 index 00000000..38760638 --- /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-exit 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-exit 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-exit 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/test/runtests.jl b/test/runtests.jl index 58ba1e97..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) @@ -57,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 index fc325805..8cf6eb03 100644 --- a/test/test_auto_compact.jl +++ b/test/test_auto_compact.jl @@ -7,15 +7,15 @@ import AdaptiveArrayPools as AAP # Auto-compact — timer-driven background capacity compaction. # Design: docs/plans/DESIGN_auto_compact.md # -# Architecture: 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 `_current_depth == 1` scope-exit safepoint. +# 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 `_current_depth == 1` scope-exit safepoint (the macro hook). # -# Phase 1 covers the data layer (flag field + registry + config + run), with no -# timer and no macro hook yet — every piece is exercised directly. +# 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. # ============================================================================== -# Allocated backing capacity (Memory length), mirroring test_compact.jl. _cap(v) = length(getfield(v, :ref).mem) # An inactive, bloated Float64 slot (cap = `big` high-water, last extent = `small`). @@ -29,8 +29,11 @@ end _clear_registry!() = lock(() -> empty!(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) -@testset "auto_compact — Phase 1: flag field + registry + run" begin +AAP.disable_auto_compact!() # stop the __init__-started timer for deterministic tests +@testset "auto_compact" begin + + # ── 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 @@ -44,7 +47,6 @@ _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPA p2 = AdaptiveArrayPool{0}() AAP.register_auto_compact!(p1) AAP.register_auto_compact!(p2) - # inject a guaranteed-dead entry to exercise pruning lock(() -> push!(AAP._AUTO_COMPACT_REGISTRY, WeakRef(nothing)), AAP._AUTO_COMPACT_LOCK) @test _registry_len() == 3 @@ -73,21 +75,15 @@ _registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPA @test AAP._AUTO_COMPACT_CONFIG.min_bytes == 2^20 @test AAP._AUTO_COMPACT_CONFIG.active == true end -end - -@testset "auto_compact — Phase 2: timer lifecycle + enable/disable" begin - # In the default test build the `auto_compact` Preference is off, so - # `enable_auto_compact!` emits a guidance @warn — asserted via @test_logs (which - # also captures it, keeping output clean). The timer still starts, so the - # sweep→flag mechanism is exercisable against a manually-registered pool. + # ── 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) - @test_logs (:warn,) match_mode = :any AAP.enable_auto_compact!( + AAP.enable_auto_compact!( interval = 0.1, factor = 20, shrink_to = 2.0, min_bytes = 4096, active = false, ) @test AAP.auto_compact_enabled() == true @@ -96,13 +92,12 @@ end @test AAP._AUTO_COMPACT_CONFIG.min_bytes == 4096 @test AAP._AUTO_COMPACT_CONFIG.active == false - sleep(0.35) # ~3 ticks @ 0.1s + sleep(0.35) # ~3 ticks @ 0.1s @test (@atomic pool._compact_requested) == true AAP.disable_auto_compact!() @test AAP.auto_compact_enabled() == false - # restore config defaults for the other testsets - AAP._AUTO_COMPACT_CONFIG.factor = 10.0 + 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 @@ -113,7 +108,7 @@ end _clear_registry!() pool = AdaptiveArrayPool{0}() AAP.register_auto_compact!(pool) - @test_logs (:warn,) match_mode = :any AAP.enable_auto_compact!(interval = 0.1) + AAP.enable_auto_compact!(interval = 0.1) AAP.disable_auto_compact!() @test AAP.auto_compact_enabled() == false @atomic pool._compact_requested = false @@ -126,9 +121,9 @@ end @testset "enable! replaces an existing timer (closes the old one)" begin AAP.disable_auto_compact!() - @test_logs (:warn,) match_mode = :any AAP.enable_auto_compact!(interval = 0.1) + AAP.enable_auto_compact!(interval = 0.1) t1 = AAP._AUTO_COMPACT_TIMER[] - @test_logs (:warn,) match_mode = :any AAP.enable_auto_compact!(interval = 0.2) + AAP.enable_auto_compact!(interval = 0.2) t2 = AAP._AUTO_COMPACT_TIMER[] @test t1 !== t2 @test !isopen(t1) # old timer closed @@ -140,18 +135,14 @@ end _clear_registry!() bogus = Ref(42) # heap object, NOT a pool lock(() -> push!(AAP._AUTO_COMPACT_REGISTRY, WeakRef(bogus)), AAP._AUTO_COMPACT_LOCK) - # the `::AdaptiveArrayPool` assertion in the sweep throws on the bogus entry; - # the _safe wrapper must swallow it so the Timer would keep running. - @test (AAP._safe_auto_compact_sweep!(nothing); true) + @test (AAP._safe_auto_compact_sweep!(nothing); true) # _safe wrapper swallows the throw @test bogus[] == 42 # keep `bogus` alive past the sweep _clear_registry!() end -end - -@testset "auto_compact — Phase 3: scope-exit hook + registration gate" begin + # ── Scope-exit hook ────────────────────────────────────────────────────── @testset "_maybe_auto_compact! fires only at _current_depth == 1" begin - pool = _inactive_bloated(1_000_000, 100) # inactive bloat, depth back to 1 + 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 @@ -186,11 +177,54 @@ end @test occursin("AUTO_COMPACT", s) # …gated by the compile-time const end - @testset "registration gate is off when AUTO_COMPACT is compiled out (default)" begin - if !AAP.AUTO_COMPACT - _clear_registry!() - fetch(Threads.@spawn get_task_local_pool()) # new task → creates a pool - @test _registry_len() == 0 # not auto-registered (gate off) + # ── 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 exit 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 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..85026c63 --- /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-exit 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 From 62a7896e45f2dfa8a279075659b7ebe238e1dfb8 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 16:03:26 -0700 Subject: [PATCH 06/14] fix(auto-compact): don't start the Timer during precompilation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit __init__ auto-started the background Timer whenever AUTO_COMPACT is on. When a CUDA/Metal extension is precompiled, its worker loads AdaptiveArrayPools and runs __init__ → the live Timer leaves an open libuv handle → the precompile worker hangs with "unfinished IO handles with timer events" (surfaced by the full suite, which precompiles both extensions). Guard the auto-start with `ccall(:jl_generating_output, Cint, ()) == 0` so no Timer is started during precompile/sysimage output, and atexit-close it for tidy shutdown. Normal processes still get default-on — verified: Metal + CUDA extensions now precompile cleanly and auto_compact is enabled=true at runtime. Also document in enable_auto_compact! that auto-compaction targets the task-local pool only and runs at @with_pool exits. --- src/auto_compact.jl | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/auto_compact.jl b/src/auto_compact.jl index ae7f53f1..2042b8c5 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -154,6 +154,13 @@ Requires the `auto_compact` compile-time Preference (`AUTO_COMPACT`) for the sco hook and pool auto-registration. With the Preference off this warns and the timer still runs but stays inert unless you `register_auto_compact!` pools manually. +!!! note "What gets compacted, and when" + Auto-compaction targets the **task-local pool** (`get_task_local_pool()`) only — the + pool `@with_pool` uses and the one that auto-registers. A hand-made `AdaptiveArrayPool()` + is never registered and never auto-compacted; call `compact!` on it directly. The work + runs at `@with_pool` scope exits (`_current_depth == 1`), so an idle task with no + `@with_pool` activity won't compact until it next exits a scope. + See also [`disable_auto_compact!`](@ref), [`auto_compact_enabled`](@ref). """ function enable_auto_compact!(; @@ -184,9 +191,15 @@ Whether the background auto-compact `Timer` is currently running. """ auto_compact_enabled() = _AUTO_COMPACT_TIMER[] !== nothing -# Auto-start when the Preference opts in. Dead-code-eliminated when `AUTO_COMPACT` is -# off (the default), so package load starts no timer unless the user set the Preference. +# 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__() - AUTO_COMPACT && enable_auto_compact!() + if AUTO_COMPACT && ccall(:jl_generating_output, Cint, ()) == 0 + enable_auto_compact!() + atexit(disable_auto_compact!) + end return nothing end From ffc6a943129ea3b5aea052e70e413e9546c5e05a Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 16:45:31 -0700 Subject: [PATCH 07/14] refactor(auto-compact): entry-hook + review hardening MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adopt a scope-ENTRY hook (check the flag before the @with_pool checkpoint, at depth==1) instead of the scope-exit hook. One hook point services a pending request when the NEXT @with_pool is entered, regardless of how the previous scope exited — normal, early return, break, continue, or exception — and never runs compaction inside a `finally` during unwind. This resolves the adversarial review's two must-fix findings (early-exit paths skipped the hook; finally-path compaction could mask a user exception) by construction. Additional review hardening: - Owner-side compaction wrapped in try/catch: a background-maintenance failure is logged and the cycle skipped; it never surfaces at the @with_pool boundary. - Atomic flag uses :monotonic — a one-way eventual-visibility signal with no cross-field ordering invariant; avoids seq-cst barriers on weak-ordered CPUs (e.g. Apple Silicon), the exact per-entry cost the hook minimizes. - Config snapshot + timer lifecycle guarded by the registry lock (no torn config read across threads; no leaked/double-closed Timer from concurrent enable/disable). - enable_auto_compact! warn gains maxlog=1; register_auto_compact! narrowed to ::AdaptiveArrayPool so a mis-registration fails at the call site. - Document the zero-alloc trade-off: a @with_pool that services a request allocates the right-sized backing, so it is not strictly zero-alloc; set auto_compact=false for a guaranteed-zero-alloc hot path. New tests: early return / break service the request at the next entry; @safe_with_pool fires the hook during exception unwind and preserves the exception. --- src/auto_compact.jl | 100 +++++++++++++++++++++++--------------- src/macros.jl | 19 ++++---- test/test_auto_compact.jl | 80 ++++++++++++++++++++++++++++-- 3 files changed, 146 insertions(+), 53 deletions(-) diff --git a/src/auto_compact.jl b/src/auto_compact.jl index 2042b8c5..41356c2c 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -4,12 +4,12 @@ # # 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 -# `_current_depth == 1` scope-exit safepoint, where nothing is borrowed. +# 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-exit flag check is +# 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/exit); set the Preference false to compile it OUT (full DCE → zero cost). +# (~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!`. # ============================================================================== @@ -18,8 +18,8 @@ using Preferences: @load_preference """ AUTO_COMPACT::Bool -Compile-time master switch (Tier 1) for auto-compaction. Default `true`: the scope-exit -flag check is emitted (a ~0.6 ns inlined check per `@with_pool` exit), new task-local +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: @@ -56,7 +56,7 @@ Add `pool` to the auto-compact registry (under lock) so the background Timer's s will flag it. Called from the task-local-pool slow path when `AUTO_COMPACT` is on (one entry per task-local pool). """ -function register_auto_compact!(pool) +function register_auto_compact!(pool::AdaptiveArrayPool) lock(() -> push!(_AUTO_COMPACT_REGISTRY, WeakRef(pool)), _AUTO_COMPACT_LOCK) return nothing end @@ -74,7 +74,7 @@ function _auto_compact_sweep!(_timer) pop!(_AUTO_COMPACT_REGISTRY) # dead → swap-remove else p = pool::AdaptiveArrayPool - @atomic p._compact_requested = true # request compaction + @atomic :monotonic p._compact_requested = true # request compaction i += 1 end end @@ -82,30 +82,39 @@ function _auto_compact_sweep!(_timer) return nothing end -# Owner-side execution: reset the flag BEFORE the work (so a concurrent sweep's -# re-request is not lost), then run the full configured `compact!`. Called at the -# `_current_depth == 1` safepoint by the macro hook (Phase 3) — the depth guard is the -# caller's. Returns nothing; the summary is discarded (no allocation). +# 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). function _run_auto_compact!(pool::AdaptiveArrayPool) - @atomic pool._compact_requested = false - cfg = _AUTO_COMPACT_CONFIG - compact!( - pool; - factor = cfg.factor, shrink_to = cfg.shrink_to, - min_bytes = cfg.min_bytes, active = cfg.active, - ) + @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-exit hook body, generated into every `@with_pool` exit (gated by AUTO_COMPACT). -# At the `_current_depth == 1` safepoint, run a flagged pool's compaction. `@inline` so -# the common case (flag clear) is a cheap inlined atomic-read + compare; the cold -# `_run_auto_compact!` stays a non-inlined call. The `::Any` fallback keeps the SHARED -# `@with_pool` codegen safe for non-CPU (GPU) pools — they no-op here, and a future phase -# adds their own methods. +# 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 non-CPU (GPU) pools — they no-op here; a future phase adds +# their own methods. `: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 pool._compact_requested) && pool._current_depth == 1 + if (@atomic :monotonic pool._compact_requested) && pool._current_depth == 1 _run_auto_compact!(pool) end return nothing @@ -127,13 +136,19 @@ function _safe_auto_compact_sweep!(timer) 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) - _stop_auto_compact_timer!() # replace any existing - _AUTO_COMPACT_TIMER[] = Timer(_safe_auto_compact_sweep!, interval; interval = interval) + 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 -function _stop_auto_compact_timer!() +_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) @@ -148,18 +163,21 @@ end 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 `_current_depth == 1` scope-exit safepoint. +`compact!` at its next `@with_pool` **entry** (the `_current_depth == 1` safepoint). -Requires the `auto_compact` compile-time Preference (`AUTO_COMPACT`) for the scope-exit +Requires the `auto_compact` compile-time Preference (`AUTO_COMPACT`) for the scope-entry hook and pool auto-registration. With the Preference off this warns and the timer still runs but stays inert unless you `register_auto_compact!` pools manually. -!!! note "What gets compacted, and when" +!!! note "What gets compacted, when, and the zero-alloc trade-off" Auto-compaction targets the **task-local pool** (`get_task_local_pool()`) only — the pool `@with_pool` uses and the one that auto-registers. A hand-made `AdaptiveArrayPool()` is never registered and never auto-compacted; call `compact!` on it directly. The work - runs at `@with_pool` scope exits (`_current_depth == 1`), so an idle task with no - `@with_pool` activity won't compact until it next exits a scope. + 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). """ @@ -167,12 +185,14 @@ 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 scope-exit hook is compiled out and pools won't auto-register. The timer runs but stays inert unless you register pools manually. Set the Preference and restart to activate." - cfg = _AUTO_COMPACT_CONFIG - cfg.factor = Float64(factor) - cfg.shrink_to = Float64(shrink_to) - cfg.min_bytes = min_bytes - cfg.active = active + AUTO_COMPACT || @warn "enable_auto_compact!: the `auto_compact` Preference is off — the scope-entry hook is compiled out and pools won't auto-register. The timer runs but stays inert unless you register pools manually. Set the Preference and restart to activate." 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 diff --git a/src/macros.jl b/src/macros.jl index 7a163424..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)) @@ -752,7 +753,6 @@ function _generate_block_inner(pool_name, expr, safe::Bool, source) _result finally $rewind_call - $(_auto_compact_hook(pool_name)) end end else @@ -769,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 @@ -784,7 +785,6 @@ function _generate_block_inner(pool_name, expr, safe::Bool, source) $_validate_pool_return(_result, $(esc(pool_name))) end $rewind_call - $(_auto_compact_hook(pool_name)) _result end end @@ -823,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)) @@ -832,7 +833,6 @@ function _generate_function_inner(pool_name, expr, safe::Bool, source) _result finally $rewind_call - $(_auto_compact_hook(pool_name)) end end else @@ -849,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 @@ -864,7 +865,6 @@ function _generate_function_inner(pool_name, expr, safe::Bool, source) $_validate_pool_return(_result, $(esc(pool_name))) end $rewind_call - $(_auto_compact_hook(pool_name)) _result end end @@ -1592,11 +1592,12 @@ 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-exit hook, generated after the normal rewind in every `@with_pool` -# exit. 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 AFTER the rewind, so `_current_depth == 1` ⇒ the outermost scope -# just closed (nothing borrowed → even active compaction is safe). +# 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)))) diff --git a/test/test_auto_compact.jl b/test/test_auto_compact.jl index 8cf6eb03..f948f07f 100644 --- a/test/test_auto_compact.jl +++ b/test/test_auto_compact.jl @@ -9,7 +9,7 @@ import AdaptiveArrayPools as AAP # # 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 `_current_depth == 1` scope-exit safepoint (the macro hook). +# 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 @@ -140,7 +140,7 @@ AAP.disable_auto_compact!() # stop the __init__-started timer for deterministi _clear_registry!() end - # ── Scope-exit hook ────────────────────────────────────────────────────── + # ── 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 @@ -166,7 +166,7 @@ AAP.disable_auto_compact!() # stop the __init__-started timer for deterministi @test AAP._maybe_auto_compact!("not a pool") === nothing end - @testset "macro wires the gated scope-exit hook into @with_pool" begin + @testset "macro wires the gated scope-entry hook into @with_pool" begin s = string( @macroexpand @with_pool pool begin x = acquire!(pool, Float64, 4) @@ -216,7 +216,7 @@ AAP.disable_auto_compact!() # stop the __init__-started timer for deterministi # 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 exit at depth 1 → hook fires + @with_pool p begin # scope ENTRY at depth 1 → hook fires acquire!(p, Float64, 4) end @@ -225,6 +225,78 @@ AAP.disable_auto_compact!() # stop the __init__-started timer for deterministi 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 From dcd3e8d1facd00b7fa973fb3c7883804eb623f35 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 17:16:20 -0700 Subject: [PATCH 08/14] (feat): Metal parity for auto-compact MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror the CPU auto-compact wiring onto the Metal backend so a `MetalAdaptiveArrayPool` is flagged by the single background Timer and runs the full `compact!` at its next `@with_pool :metal` entry, exactly as CPU pools do. Base module (shared across backends): - Broaden `register_auto_compact!`, the `_auto_compact_sweep!` flag write, and `_run_auto_compact!` from `AdaptiveArrayPool` to the `AbstractArrayPool` supertype. One registry, one sweep, one runner now serve every backend; `compact!` takes the same kwargs everywhere, so the runner is unchanged. Metal extension: - Add the `@atomic _compact_requested::Bool` field to `MetalAdaptiveArrayPool` (+ constructor) — the cross-thread handoff flag, parity with CPU. - New `_maybe_auto_compact!(::MetalAdaptiveArrayPool)`: services the flag only at `_current_depth == 1`, reusing the base `_run_auto_compact!`. The scope-entry hook is already emitted by the shared `@with_pool` codegen; before this it fell to the `::Any` no-op. - Register each task-local Metal pool (per device) for auto-compaction in the slow path, gated by the compile-time `AUTO_COMPACT` (DCE'd when off). Coexistence: a CPU pool and a Metal pool share the one CPU-side Timer. A single sweep flags both; each then compacts independently at its own scope boundary, and servicing one never touches the other's flag or buffers — pinned by a dedicated test. Full Metal suite + new `test/metal/test_auto_compact.jl` (30 tests) green on Apple Silicon. --- .../AdaptiveArrayPoolsMetalExt.jl | 1 + .../auto_compact.jl | 25 +++ .../task_local_pool.jl | 4 + ext/AdaptiveArrayPoolsMetalExt/types.jl | 6 + src/auto_compact.jl | 22 ++- test/metal/runtests.jl | 1 + test/metal/test_auto_compact.jl | 181 ++++++++++++++++++ 7 files changed, 232 insertions(+), 8 deletions(-) create mode 100644 ext/AdaptiveArrayPoolsMetalExt/auto_compact.jl create mode 100644 test/metal/test_auto_compact.jl 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/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/auto_compact.jl b/src/auto_compact.jl index 41356c2c..d67f93f1 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -54,9 +54,11 @@ const _AUTO_COMPACT_CONFIG = _AutoCompactConfig(10.0, 1.5, 2^20, true) 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). +(one entry per task-local pool). Accepts any `AbstractArrayPool` so the same registry +serves every backend (CPU `AdaptiveArrayPool`, GPU `MetalAdaptiveArrayPool`, …); each +carries its own `@atomic _compact_requested` flag and services it at its own scope entry. """ -function register_auto_compact!(pool::AdaptiveArrayPool) +function register_auto_compact!(pool::AbstractArrayPool) lock(() -> push!(_AUTO_COMPACT_REGISTRY, WeakRef(pool)), _AUTO_COMPACT_LOCK) return nothing end @@ -73,7 +75,7 @@ function _auto_compact_sweep!(_timer) _AUTO_COMPACT_REGISTRY[i] = _AUTO_COMPACT_REGISTRY[end] pop!(_AUTO_COMPACT_REGISTRY) # dead → swap-remove else - p = pool::AdaptiveArrayPool + p = pool::AbstractArrayPool # CPU or any GPU backend pool @atomic :monotonic p._compact_requested = true # request compaction i += 1 end @@ -87,8 +89,10 @@ end # 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). -function _run_auto_compact!(pool::AdaptiveArrayPool) +# 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 @@ -109,9 +113,11 @@ end # 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 non-CPU (GPU) pools — they no-op here; a future phase adds -# their own methods. `: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). +# `@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 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..aec0be45 --- /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 + +# Device capacity (elements) of a Metal backing buffer. +_mcap(v) = Int(getfield(v, :maxsize) ÷ sizeof(eltype(v))) +# CPU backing capacity (Memory length) — for the coexistence test. +_cpucap(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 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 + +AAP.disable_auto_compact!() # stop the __init__-started timer for deterministic tests + +@testset "Metal auto_compact" begin + + # ── 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 From bbc10436588ed810d72d4cffce893c369c7c8ad9 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 19:50:27 -0700 Subject: [PATCH 09/14] (feat): CUDA parity for auto-compact MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror the Metal auto-compact wiring onto the CUDA backend (the base-module support — one registry, one Timer, `AbstractArrayPool`-typed sweep/runner — was already broadened in the Metal commit, so this is purely the CUDA-side three edits plus tests). - Add the `@atomic _compact_requested::Bool` field to `CuAdaptiveArrayPool` (+ both construction paths: the inner ctor and the `_transfer_cuda_pool` migration ctor) — the cross-thread handoff flag, parity with CPU/Metal. - New `_maybe_auto_compact!(::CuAdaptiveArrayPool)`: services the flag only at `_current_depth == 1`, reusing the base `_run_auto_compact!`. The scope-entry hook is already emitted by the shared `@with_pool :cuda` codegen; before this it fell to the `::Any` no-op. - Register each task-local CUDA pool (per device) for auto-compaction in the slow path, gated by the compile-time `AUTO_COMPACT` (DCE'd when off). Tests: test/cuda/test_auto_compact.jl mirrors the Metal suite (field, registry sweep, scope-entry hook, end-to-end @with_pool :cuda, and a CPU+CUDA coexistence test pinning that one sweep flags both and each compacts independently at its own boundary). Validated locally as far as a non-NVIDIA host allows: the CUDA extension precompiles cleanly and reflection confirms the field placement, the concrete _maybe_auto_compact! method, and AbstractArrayPool dispatch. The GPU runtime assertions are validated on the remote CUDA box. --- .../AdaptiveArrayPoolsCUDAExt.jl | 3 + ext/AdaptiveArrayPoolsCUDAExt/auto_compact.jl | 25 +++ .../task_local_pool.jl | 4 + ext/AdaptiveArrayPoolsCUDAExt/types.jl | 11 +- test/cuda/runtests.jl | 1 + test/cuda/test_auto_compact.jl | 181 ++++++++++++++++++ 6 files changed, 223 insertions(+), 2 deletions(-) create mode 100644 ext/AdaptiveArrayPoolsCUDAExt/auto_compact.jl create mode 100644 test/cuda/test_auto_compact.jl 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/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/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..60bcb696 --- /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 + +# Device capacity (elements) of a CUDA backing buffer. +_ccap(v) = Int(getfield(v, :maxsize) ÷ sizeof(eltype(v))) +# CPU backing capacity (Memory length) — for the coexistence test. +_cpucap(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 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 + +AAP.disable_auto_compact!() # stop the __init__-started timer for deterministic tests + +@testset "CUDA auto_compact" begin + + # ── 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 From 0d9682aab266a51ef2439720a4487f7201623e7c Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 21:03:25 -0700 Subject: [PATCH 10/14] (test): scope auto-compact test helpers to avoid Main method clashes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Metal and CUDA `test_auto_compact.jl` files defined `_clear_registry!` / `_registry_len` (and `_cpucap`) as top-level functions in `Main`, colliding with the identically-named top-level helpers in the CPU `test_auto_compact.jl` when the full suite loads several backends' files together: WARNING: Method definition _clear_registry!() in module Main at test/test_auto_compact.jl:29 overwritten at test/cuda/test_auto_compact.jl:22. Move each GPU file's helpers INSIDE its `@testset ... begin` (a local scope, the same convention the `test_compact.jl` files already use), so they no longer define `Main` methods. The CPU file stays the sole top-level definer — no collision. Warning only; all tests passed before and after. Verified: 0 "overwritten" warnings when the CPU + Metal files load into one Main, both suites green. --- test/cuda/test_auto_compact.jl | 32 ++++++++++++++++---------------- test/metal/test_auto_compact.jl | 32 ++++++++++++++++---------------- 2 files changed, 32 insertions(+), 32 deletions(-) diff --git a/test/cuda/test_auto_compact.jl b/test/cuda/test_auto_compact.jl index 60bcb696..84d48dd7 100644 --- a/test/cuda/test_auto_compact.jl +++ b/test/cuda/test_auto_compact.jl @@ -14,26 +14,26 @@ import AdaptiveArrayPools as AAP -# Device capacity (elements) of a CUDA backing buffer. -_ccap(v) = Int(getfield(v, :maxsize) ÷ sizeof(eltype(v))) -# CPU backing capacity (Memory length) — for the coexistence test. -_cpucap(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 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 - 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}() diff --git a/test/metal/test_auto_compact.jl b/test/metal/test_auto_compact.jl index aec0be45..9602d67e 100644 --- a/test/metal/test_auto_compact.jl +++ b/test/metal/test_auto_compact.jl @@ -14,26 +14,26 @@ import AdaptiveArrayPools as AAP -# Device capacity (elements) of a Metal backing buffer. -_mcap(v) = Int(getfield(v, :maxsize) ÷ sizeof(eltype(v))) -# CPU backing capacity (Memory length) — for the coexistence test. -_cpucap(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 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 - 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}() From 803264b97405731b732892f9c5e4889c42fe7c74 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 21:13:14 -0700 Subject: [PATCH 11/14] (perf): GPU compact! copy-skip parity with CPU (skip dead inactive-slot copy) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CPU `_compact_slot!` takes a `copy_live::Bool` and copies the live `used` elements only for ACTIVE slots (`slot <= n_active`); an inactive slot's contents are dead (the next `acquire!` returns uninitialized memory the caller fills), so the copy is skipped — inactive compaction is allocate-only. The GPU backends only had the weaker `used > 0` skip (catches empty slots only), so an inactive slot that was used still paid a wasted GPU→GPU `copyto!` of `used` elements — exactly the auto-compact case, where every flagged slot is inactive. Bring Metal and CUDA `_compact_slot!` to full CPU parity: add the `copy_live` parameter (`copy_live && copyto!`), and have `_maybe_compact_slot!` pass `slot <= tp.n_active`. Behavior-preserving for correct usage (inactive contents are dead, never read) — it removes a device copy / kernel launch per inactive compacted slot. Tests: add the per-slot copy-decision testset to test/metal/test_compact.jl and test/cuda/test_compact.jl (mirror of the CPU one) — two bloated slots, hold only slot 1, `compact!(active=true)`: both shrink, the held slot's data survives, and the inactive (copy-skipped) slot still re-acquires correctly. Full Metal suite green on Apple Silicon; CUDA validated on the remote box. --- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 12 ++++++-- ext/AdaptiveArrayPoolsMetalExt/state.jl | 12 ++++++-- test/cuda/test_compact.jl | 41 +++++++++++++++++++++++++ test/metal/test_compact.jl | 41 +++++++++++++++++++++++++ 4 files changed, 100 insertions(+), 6 deletions(-) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index 1bf7ab67..4384de52 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 @@ -492,7 +497,8 @@ function AdaptiveArrayPools._maybe_compact_slot!(tp::CuTypedPool{T}, slot::Int, 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/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 5d5e5083..66277313 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 @@ -465,7 +470,8 @@ function AdaptiveArrayPools._maybe_compact_slot!(tp::MetalTypedPool{T, S}, slot: 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/test/cuda/test_compact.jl b/test/cuda/test_compact.jl index 8c845c11..018a881a 100644 --- a/test/cuda/test_compact.jl +++ b/test/cuda/test_compact.jl @@ -78,4 +78,45 @@ @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 + end diff --git a/test/metal/test_compact.jl b/test/metal/test_compact.jl index 4a2fd693..7ce6cb90 100644 --- a/test/metal/test_compact.jl +++ b/test/metal/test_compact.jl @@ -81,4 +81,45 @@ @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 + end From 25c4c8c1deb5d7afe281d6be169e66e05097e01c Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 21:23:46 -0700 Subject: [PATCH 12/14] (test): silence two confusing warnings in the auto-compact suite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two non-failing warnings in the full suite looked like errors: 1. `WARNING: Method definition _cap(Any) ... overwritten` — test_auto_compact.jl defined a top-level `_cap` helper in `Main`, colliding with the identically-named top-level helper in test_compact.jl. Same class as 0d9682a (GPU files), but CPU↔CPU, which that commit missed (test_compact.jl is an independent definer of `_cap`, not just test_auto_compact.jl). Scope the CPU helpers (`_cap`, `_clear_registry!`, `_registry_len`, `_inactive_bloated`) inside the testset. 2. The "auto-compact sweep failed (timer continues)" @warn is emitted by the DELIBERATE negative test that pushes a non-pool into the registry to prove the sweep wrapper catches the throw (so the Timer survives). The test passes; the warning just read as a failure. Wrap that one call in a NullLogger so the expected warning is silenced — the no-rethrow assertion is unchanged. Verified: 0 "overwritten" and 0 "sweep failed" lines when test_compact.jl + test_auto_compact.jl load together; compact! 105/105, auto_compact 54/54. --- test/test_auto_compact.jl | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/test/test_auto_compact.jl b/test/test_auto_compact.jl index f948f07f..2262c484 100644 --- a/test/test_auto_compact.jl +++ b/test/test_auto_compact.jl @@ -1,4 +1,5 @@ using Test +using Logging: with_logger, NullLogger using AdaptiveArrayPools using AdaptiveArrayPools: checkpoint!, rewind! import AdaptiveArrayPools as AAP @@ -16,23 +17,25 @@ import AdaptiveArrayPools as AAP # explicitly and the suite stays deterministic. # ============================================================================== -_cap(v) = length(getfield(v, :ref).mem) - -# 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 - -_clear_registry!() = lock(() -> empty!(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) -_registry_len() = lock(() -> length(AAP._AUTO_COMPACT_REGISTRY), AAP._AUTO_COMPACT_LOCK) - 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}() @@ -135,7 +138,12 @@ AAP.disable_auto_compact!() # stop the __init__-started timer for deterministi _clear_registry!() bogus = Ref(42) # heap object, NOT a pool lock(() -> push!(AAP._AUTO_COMPACT_REGISTRY, WeakRef(bogus)), AAP._AUTO_COMPACT_LOCK) - @test (AAP._safe_auto_compact_sweep!(nothing); true) # _safe wrapper swallows the throw + # 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 From b0ba1dad4dd0d95d26aa2332accdf0d48b7d6dbe Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 21:34:30 -0700 Subject: [PATCH 13/14] (fix): address Copilot PR review + test `Logging` dependency MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Copilot review on PR #46 (all valid): - register_auto_compact!(::AbstractArrayPool) could accept a flagless pool (DisabledPool, or a foreign AbstractArrayPool subtype) — registering one makes the sweep throw on its missing `_compact_requested` and abort the cycle, starving later pools. Guard with `hasfield(typeof(pool), :_compact_requested)` (const-folds per concrete type); add a regression test that a DisabledPool is silently skipped and a real pool registered alongside it still gets flagged. - Docstrings/warning for the `AUTO_COMPACT=false` case were misleading: with the Preference off the `@with_pool` hook is compiled out, so the timer's flags are never serviced — manual `register_auto_compact!` does NOT make it effective. Reword the `enable_auto_compact!` docstring + warning to say so, and soften "hand-made pool is never registered" to "not auto-registered (register explicitly)". - legacy/auto_compact.jl + test_auto_compact_legacy.jl said "scope-exit hook" — stale from the exit-hook era; the hook is at scope ENTRY now. Also fix a test-env load error (surfaced by the CUDA run): the warning-suppression helper used `using Logging`, which isn't a declared test dep. Source `with_logger` / `NullLogger` from `Base.CoreLogging` instead — built into Base, no dependency needed. CPU auto_compact suite green (55 tests). --- src/auto_compact.jl | 28 +++++++++++++++++++--------- src/legacy/auto_compact.jl | 6 +++--- test/test_auto_compact.jl | 16 +++++++++++++++- test/test_auto_compact_legacy.jl | 2 +- 4 files changed, 38 insertions(+), 14 deletions(-) diff --git a/src/auto_compact.jl b/src/auto_compact.jl index d67f93f1..25da8e8c 100644 --- a/src/auto_compact.jl +++ b/src/auto_compact.jl @@ -54,11 +54,17 @@ const _AUTO_COMPACT_CONFIG = _AutoCompactConfig(10.0, 1.5, 2^20, true) 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). Accepts any `AbstractArrayPool` so the same registry -serves every backend (CPU `AdaptiveArrayPool`, GPU `MetalAdaptiveArrayPool`, …); each -carries its own `@atomic _compact_requested` flag and services it at its own scope entry. +(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 @@ -172,13 +178,17 @@ Start (or restart) the background auto-compact `Timer` with the given config. Ev `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 this warns and the timer still -runs but stays inert unless you `register_auto_compact!` pools manually. +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()`) only — the - pool `@with_pool` uses and the one that auto-registers. A hand-made `AdaptiveArrayPool()` - is never registered and never auto-compacted; call `compact!` on it directly. The work + 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 @@ -191,7 +201,7 @@ 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 scope-entry hook is compiled out and pools won't auto-register. The timer runs but stays inert unless you register pools manually. Set the Preference and restart to activate." maxlog = 1 + 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) diff --git a/src/legacy/auto_compact.jl b/src/legacy/auto_compact.jl index 38760638..ea8e02f0 100644 --- a/src/legacy/auto_compact.jl +++ b/src/legacy/auto_compact.jl @@ -4,18 +4,18 @@ # 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-exit hook DCEs away (AUTO_COMPACT === false). The public +# 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-exit hook to nothing. +# 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-exit hook. +# Referenced (and DCE'd) by the @with_pool scope-entry hook. @inline _maybe_auto_compact!(::Any) = nothing """ diff --git a/test/test_auto_compact.jl b/test/test_auto_compact.jl index 2262c484..6305b9ac 100644 --- a/test/test_auto_compact.jl +++ b/test/test_auto_compact.jl @@ -1,5 +1,5 @@ using Test -using Logging: with_logger, NullLogger +using Base.CoreLogging: with_logger, NullLogger # from Base — no `Logging` dep needed using AdaptiveArrayPools using AdaptiveArrayPools: checkpoint!, rewind! import AdaptiveArrayPools as AAP @@ -61,6 +61,20 @@ AAP.disable_auto_compact!() # stop the __init__-started timer for deterministi _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 diff --git a/test/test_auto_compact_legacy.jl b/test/test_auto_compact_legacy.jl index 85026c63..7ae1f345 100644 --- a/test/test_auto_compact_legacy.jl +++ b/test/test_auto_compact_legacy.jl @@ -4,7 +4,7 @@ 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-exit hook DCE away, and the public enable/disable/enabled API +# 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 From 988c50fb87761f6cdf00b85569b9a636825328ac Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Mon, 29 Jun 2026 22:03:27 -0700 Subject: [PATCH 14/14] (fix): compact! must not undo RUNTIME_CHECK poison on inactive slots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex review (PR #46, P1): with RUNTIME_CHECK on (S/R=1), releasing a slot poisons its backing (NaN/typemax) AND shrinks the logical length to 0 (CPU `resize!(v,0)`; GPU `_resize_to_fit!(v,0)`, device memory retained), so an escaped `acquire_view!` view reads poison instead of live data. But `_maybe_compact_slot!` still saw such a slot as "bloated" (large retained capacity, nonzero recorded extent) and — with the inactive copy-skip (`copy_live=false`) — swapped in a fresh *uninitialized* buffer, UNDOING the poison: the escaped view would then read plausible memory instead of NaN. `compact!(; active=false)` (and auto-compact at depth 1, which sees every slot inactive) silently defeated escape detection. Fix: skip compacting a poisoned slot. Poison and the length-0 shrink are set together in `_invalidate_released_slots!`, and an unchecked slot keeps its logical length, so `length(v) == 0 && used > 0` is an exact marker for "poisoned". Add `length(v) == 0 && return 0` to `_maybe_compact_slot!` on CPU, Metal, and CUDA — no signature threading, and S=0 (the zero-overhead default) is unaffected (released slots there keep their length, so the guard never fires). The slot compacts normally once re-acquired. Active held slots (copy_live=true) are untouched. Tests: RUNTIME_CHECK=1 regression in test/test_compact.jl (proved the bug — current code shrank a poisoned 1M slot to 150) + test/metal/test_compact.jl + test/cuda/test_compact.jl. Full Metal suite green on Apple Silicon (incl. the R=1 case); CPU green; CUDA validated on the remote box. --- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 5 +++++ ext/AdaptiveArrayPoolsMetalExt/state.jl | 5 +++++ src/state.jl | 7 +++++++ test/cuda/test_compact.jl | 18 +++++++++++++++++ test/metal/test_compact.jl | 18 +++++++++++++++++ test/test_compact.jl | 27 +++++++++++++++++++++++++ 6 files changed, 80 insertions(+) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index 4384de52..7ff5714f 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -492,6 +492,11 @@ 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] diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 66277313..add26a27 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -465,6 +465,11 @@ 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] diff --git a/src/state.jl b/src/state.jl index a630a69b..e9c0fbfe 100644 --- a/src/state.jl +++ b/src/state.jl @@ -908,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 diff --git a/test/cuda/test_compact.jl b/test/cuda/test_compact.jl index 018a881a..d818b7e2 100644 --- a/test/cuda/test_compact.jl +++ b/test/cuda/test_compact.jl @@ -119,4 +119,22 @@ 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/test_compact.jl b/test/metal/test_compact.jl index 7ce6cb90..8f925814 100644 --- a/test/metal/test_compact.jl +++ b/test/metal/test_compact.jl @@ -122,4 +122,22 @@ 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/test_compact.jl b/test/test_compact.jl index 00cdd17f..7f99c361 100644 --- a/test/test_compact.jl +++ b/test/test_compact.jl @@ -442,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