From cff21b7bdda1e5f605c6be04e21f9b8b963cb674 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 22:04:05 -0700 Subject: [PATCH 01/10] (refactor): generalize touched-others stack helpers over AbstractArrayPool --- src/state.jl | 43 +++++++++++++++++++++++++++---------------- src/types.jl | 5 +++++ 2 files changed, 32 insertions(+), 16 deletions(-) diff --git a/src/state.jl b/src/state.jl index 82c3cbeb..f9feab0e 100644 --- a/src/state.jl +++ b/src/state.jl @@ -48,9 +48,10 @@ Also updates _current_depth and bitmask state for type touch tracking. pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) - # Push true when T is a fallback type (non-fixed-slot) so that - # _typed_lazy_rewind! iterates pool.others even if _acquire_impl! - # (which bypasses _record_type_touch!) is the only acquire path. + # Push true when T is a fallback type (non-fixed-slot). The flag no longer + # drives fallback rewind coverage (the touched-others stack below handles + # that); it only feeds _can_use_typed_path's fast-path check and the S>=1 + # pointer-overlap validation in debug.jl. push!(pool._touched_has_others, _fixed_slot_bit(T) == UInt16(0)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) if _fixed_slot_bit(T) == UInt16(0) @@ -356,7 +357,10 @@ end # T-independent rewind core: orphan cleanup + Case A/B restore. Returns the # pre-rewind n_active so the (S >= 1) caller can decide whether to invalidate. @inline function _rewind_state_core!(st, current_depth::Int) - # 1. Orphaned checkpoints from deeper scopes + # 1. Orphaned checkpoints from deeper scopes. These arise when a deeper + # scope checkpointed this pool but its own rewind was skipped (e.g. an + # exception unwound past it before rewind! ran) — see the exception-leak + # testset. while @inbounds(st._checkpoint_depths[end]) > current_depth pop!(st._checkpoint_depths) pop!(st._checkpoint_n_active) @@ -385,18 +389,25 @@ end # Touched-Others Stack (per-scope selective fallback checkpoint/rewind) # ============================================================================== -# Rewind and remove this depth's touched-fallback entries. At S = 0 the loop runs -# entirely on concrete PoolCheckpointState objects — zero dynamic dispatch; at -# S >= 1 it routes through the typed pools so released slots get invalidated. -# Fixed-only scopes exit via one isempty/top-depth check. -@inline function _drain_touched_others!(pool::AdaptiveArrayPool{S}, d::Int) where {S} +# Rewind and remove this depth's touched-fallback entries. With runtime checks +# off the loop runs entirely on concrete PoolCheckpointState objects — zero +# dynamic dispatch; with checks on it routes through the typed pools so released +# slots get invalidated. Fixed-only scopes exit via one isempty/top-depth check. +# Generic over AbstractArrayPool: CPU and GPU pools share this implementation +# (all carry the three parallel stacks; _runtime_check/_check_level are +# compile-time constants per pool type, so the branch still folds away). +@inline function _drain_touched_others!(pool::AbstractArrayPool, d::Int) depths = pool._touched_others_depths states = pool._touched_others_states while !isempty(depths) && @inbounds(depths[end]) == d pop!(depths) st = pop!(states) - if S >= 1 - _rewind_typed_pool!(pop!(pool._touched_others_pools), d, S) + if _runtime_check(pool) + # `st` is intentionally discarded here: `_touched_others_pools` is + # popped in lockstep with `states`, and `_rewind_typed_pool!` + # re-derives the checkpoint state from the popped `tp` via + # `_cp_state`, so the two never disagree. + _rewind_typed_pool!(pop!(pool._touched_others_pools), d, _check_level(pool)) else _rewind_state_core!(st, d) end @@ -404,14 +415,14 @@ end return nothing end -# Truncate-only variant for full rewind!(pool): its _others_values sweep already -# rewound every fallback pool, so draining again would double-pop checkpoints. -@inline function _truncate_touched_others!(pool::AdaptiveArrayPool{S}, d::Int) where {S} +# Truncate-only variant for full rewind!(pool): its full sweep already rewound +# every fallback pool, so draining again would double-pop checkpoints. +@inline function _truncate_touched_others!(pool::AbstractArrayPool, d::Int) depths = pool._touched_others_depths while !isempty(depths) && @inbounds(depths[end]) == d pop!(depths) pop!(pool._touched_others_states) - S >= 1 && pop!(pool._touched_others_pools) + _runtime_check(pool) && pop!(pool._touched_others_pools) end return nothing end @@ -420,7 +431,7 @@ end # depth-tagged entry, exactly once per (pool, depth). The checkpoint-depth guard # keeps it idempotent across the three producer paths and skips pools already # eagerly checkpointed by a full checkpoint!(pool) at this depth. -@inline function _touch_fallback_pool!(pool::AdaptiveArrayPool, tp::AbstractTypedPool, depth::Int) +@inline function _touch_fallback_pool!(pool::AbstractArrayPool, tp::AbstractTypedPool, depth::Int) st = _cp_state(tp)::PoolCheckpointState if @inbounds(st._checkpoint_depths[end]) != depth push!(st._checkpoint_n_active, st.n_active) diff --git a/src/types.jl b/src/types.jl index 1d6eb0e9..06c89e4f 100644 --- a/src/types.jl +++ b/src/types.jl @@ -507,6 +507,11 @@ Compile-time constant for `AdaptiveArrayPool{S}` — dead-code eliminated when ` @inline _runtime_check(::AdaptiveArrayPool{0}) = false @inline _runtime_check(::AdaptiveArrayPool) = true # S >= 1 +# Runtime-check level as an Int for backend-shared code that forwards it to +# _invalidate_released_slots! / _rewind_typed_pool!. A compile-time constant +# per concrete pool type — inlines to a literal, so DCE at level 0 is intact. +@inline _check_level(::AdaptiveArrayPool{S}) where {S} = S + """ _make_pool(level) -> AdaptiveArrayPool From 08489aa46c53e620af1db359408585ecf177687e Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 22:14:09 -0700 Subject: [PATCH 02/10] (refactor): adopt PoolCheckpointState in MetalTypedPool; add stack+memo fields --- .../AdaptiveArrayPoolsMetalExt.jl | 2 +- ext/AdaptiveArrayPoolsMetalExt/types.jl | 60 +++++++++++++++++-- test/metal/test_extension.jl | 17 ++++-- 3 files changed, 69 insertions(+), 10 deletions(-) diff --git a/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl b/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl index 309ebc12..8d8b59ff 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/AdaptiveArrayPoolsMetalExt.jl @@ -19,7 +19,7 @@ using Metal # On older Julia, the extension loads but provides no functionality. @static if VERSION >= v"1.12-" - using AdaptiveArrayPools: AbstractTypedPool, AbstractArrayPool + using AdaptiveArrayPools: AbstractTypedPool, AbstractArrayPool, PoolCheckpointState include("types.jl") include("dispatch.jl") diff --git a/ext/AdaptiveArrayPoolsMetalExt/types.jl b/ext/AdaptiveArrayPoolsMetalExt/types.jl index bc09c3c0..c8f74df1 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/types.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/types.jl @@ -39,10 +39,11 @@ mutable struct MetalTypedPool{T, S} <: AbstractTypedPool{T, MtlArray{T, 1, S}} # (`_slot_used`) to know how much of an over-allocated device buffer is in use. slot_extents::Vector{Int} - # --- State Management (1-based sentinel pattern) --- - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} + # --- State Management --- + # Checkpoint bookkeeping, extracted into a concrete shared struct (see CPU + # PoolCheckpointState docstring). `const`: the reference never changes after + # construction. Accessed as tp.n_active / tp._checkpoint_* via forwarding below. + const state::PoolCheckpointState # --- Auto-trim telemetry (parity with CPU TypedPool; see its docstring) --- # Peak `n_active` since the last auto-trim — the recent working-set width. Written on the @@ -56,11 +57,32 @@ function MetalTypedPool{T, S}() where {T, S} MtlArray{T, 1, S}[], # vectors Union{Nothing, Vector{Any}}[], # arr_wrappers (indexed by N) Int[], # slot_extents (parallel to vectors) - 0, [0], [0], # State (1-based sentinel) + PoolCheckpointState(), # state (1-based sentinel) 0, # _am_peak_n_active: no usage observed yet ) end +# Checkpoint-state property forwarding (mirror of CPU src/types.jl:294-307). +@inline function Base.getproperty(tp::MetalTypedPool, f::Symbol) + f === :n_active && return getfield(tp, :state).n_active + f === :_checkpoint_n_active && return getfield(tp, :state)._checkpoint_n_active + f === :_checkpoint_depths && return getfield(tp, :state)._checkpoint_depths + return getfield(tp, f) +end + +@inline function Base.setproperty!(tp::MetalTypedPool, f::Symbol, v) + f === :n_active && return setfield!(getfield(tp, :state), :n_active, convert(Int, v)) + return setfield!(tp, f, convert(fieldtype(typeof(tp), f), v)) +end + +Base.propertynames(tp::MetalTypedPool) = + (fieldnames(typeof(tp))..., :n_active, :_checkpoint_n_active, :_checkpoint_depths) + +# Route the generic checkpoint/rewind cores at the concrete state (zero-dispatch +# drain); without this, MetalTypedPool falls back to _cp_state(tp) = tp and +# _touch_fallback_pool!'s ::PoolCheckpointState assert throws. +@inline AdaptiveArrayPools._cp_state(tp::MetalTypedPool) = getfield(tp, :state) + # ============================================================================== # Metal Fixed Slot Configuration # ============================================================================== @@ -121,6 +143,20 @@ mutable struct MetalAdaptiveArrayPool{R, S} <: AbstractArrayPool _touched_type_masks::Vector{UInt16} # Per-depth: which fixed slots were touched + mode flags _touched_has_others::Vector{Bool} # Per-depth: any non-fixed-slot type touched? + # Touched-others tracking (depth-tagged, concrete) — mirror of CPU + # src/types.jl:434-453. Checkpoint variants push NOTHING; producers push one + # (state, depth) entry per first touch; rewind pops while the top tag matches. + # _touched_others_pools is populated only when R >= 1 (slot invalidation). + _touched_others_states::Vector{PoolCheckpointState} + _touched_others_depths::Vector{Int} + _touched_others_pools::Vector{Any} + + # Last-lookup memo for the fallback registry (mirror of CPU). Set on every + # slow-path lookup; cleared by empty! (identities die), preserved by + # reset!/trim!/compact! (identities survive). Task-local pool → no races. + _lookup_memo_type::Any + _lookup_memo_tp::Any + # Device tracking (safety) device_key::Any @@ -151,6 +187,11 @@ function MetalAdaptiveArrayPool{R, S}() where {R, S} 1, # _current_depth (1 = global scope) [UInt16(0)], # _touched_type_masks: sentinel (no bits set) [false], # _touched_has_others: sentinel (no others) + PoolCheckpointState[], # _touched_others_states: no fallback touches yet + Int[], # _touched_others_depths + Any[], # _touched_others_pools + nothing, # _lookup_memo_type + nothing, # _lookup_memo_tp Metal.device(), "", # _pending_callsite "", # _pending_return_site @@ -176,6 +217,15 @@ Return compile-time constant indicating whether runtime safety checks are enable @inline AdaptiveArrayPools._runtime_check(::MetalAdaptiveArrayPool{0}) = false @inline AdaptiveArrayPools._runtime_check(::MetalAdaptiveArrayPool) = true # R >= 1 +""" + _check_level(pool::MetalAdaptiveArrayPool) -> Int + +Runtime-check level as an Int (mirror of CPU `src/types.jl:513`), for +backend-shared code that forwards it to `_invalidate_released_slots!` / +`_rewind_typed_pool!`. Compile-time constant per concrete pool type. +""" +@inline AdaptiveArrayPools._check_level(::MetalAdaptiveArrayPool{R, S}) where {R, S} = R + """ _make_metal_pool(level) -> MetalAdaptiveArrayPool diff --git a/test/metal/test_extension.jl b/test/metal/test_extension.jl index 69abd6ef..a2bebc1f 100644 --- a/test/metal/test_extension.jl +++ b/test/metal/test_extension.jl @@ -5,10 +5,19 @@ @testset "MetalTypedPool structure" begin tp_fields = fieldnames(MetalTypedPool) @test :vectors in tp_fields - @test :n_active in tp_fields - @test :arr_wrappers in tp_fields - @test :_checkpoint_n_active in tp_fields - @test :_checkpoint_depths in tp_fields + @test :state in tp_fields # PoolCheckpointState (n_active + checkpoint vectors) + @test !(:n_active in tp_fields) # moved into state; reachable via property forwarding + # Forwarding round-trip (mirrors CPU src/types.jl:294-307) + tp = MetalTypedPool{Float32, Metal.PrivateStorage}() + @test tp.n_active === 0 + @test tp._checkpoint_n_active == [0] + @test tp._checkpoint_depths == [0] + tp.n_active = 2 + @test getfield(tp, :state).n_active === 2 + tp.n_active = Int32(1) # default-convert semantics preserved + @test tp.n_active === 1 + tp.n_active = 0 + @test :n_active in propertynames(tp) end @testset "MetalAdaptiveArrayPool structure" begin From b02decd861eafbb8252c2e6bb24c018af6744bac Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 22:42:45 -0700 Subject: [PATCH 03/10] =?UTF-8?q?(perf):=20Metal=20=E2=80=94=20depth-tagge?= =?UTF-8?q?d=20touched-others=20stack=20+=20fallback=20lookup=20memo?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ext/AdaptiveArrayPoolsMetalExt/acquire.jl | 9 +- ext/AdaptiveArrayPoolsMetalExt/dispatch.jl | 41 ++- ext/AdaptiveArrayPoolsMetalExt/state.jl | 106 ++++-- test/metal/runtests.jl | 1 + test/metal/test_extension.jl | 2 + test/metal/test_touched_others.jl | 386 +++++++++++++++++++++ 6 files changed, 495 insertions(+), 50 deletions(-) create mode 100644 test/metal/test_touched_others.jl diff --git a/ext/AdaptiveArrayPoolsMetalExt/acquire.jl b/ext/AdaptiveArrayPoolsMetalExt/acquire.jl index 2aa8a349..0e3d6b26 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/acquire.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/acquire.jl @@ -16,7 +16,7 @@ using AdaptiveArrayPools: get_view!, get_array!, allocate_vector, safe_prod, _record_type_touch!, _fixed_slot_bit, _checkpoint_typed_pool!, _store_arr_wrapper!, _check_pool_growth, _reshape_impl!, _acquire_impl!, _acquire_view_impl!, _maybe_record_borrow!, - _MODE_BITS_MASK + _MODE_BITS_MASK, _touch_fallback_pool! using Metal: unsafe_free! @@ -336,8 +336,13 @@ end end @inbounds pool._touched_type_masks[depth] = current_mask | b16 else - # Genuine others type (UInt8, Int8, etc.) — eagerly snapshotted at scope entry. + # Genuine others type (UInt8, Int8, etc.). @inbounds pool._touched_has_others[depth] = true + # First-touch lazy checkpoint for fallback types; depth == 1 (global + # scope) is exempt — matches get_typed_pool!'s gate. + if depth > 1 + _touch_fallback_pool!(pool, AdaptiveArrayPools.get_typed_pool!(pool, T), depth) + end end else current_mask = @inbounds pool._touched_type_masks[depth] diff --git a/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl b/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl index 879e76b2..d780eb7e 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/dispatch.jl @@ -37,17 +37,32 @@ const _METAL_FIXED_TYPES = Union{Float32, Float16, Int32, Int64, ComplexF32, Boo if T === Float64 || T === ComplexF64 throw(ArgumentError("Metal backend does not support $T")) end - return get!(p.others, T) do - tp = MetalTypedPool{T, Metal.PrivateStorage}() - # CRITICAL: Match CPU behavior - auto-checkpoint new pool if inside @with_pool scope - # Without this, rewind! would corrupt state for dynamically-created pools - if p._current_depth > 1 - push!(tp._checkpoint_n_active, 0) # n_active starts at 0 - push!(tp._checkpoint_depths, p._current_depth) - # Signal that a fallback type was touched so lazy/typed-lazy rewind - # iterates pool.others (same fix as CPU get_typed_pool!) - @inbounds p._touched_has_others[p._current_depth] = true - end - tp - end::MetalTypedPool{T, Metal.PrivateStorage} + # Memo fast path: same type as the previous slow-path lookup (mirror of CPU + # src/types.jl's get_typed_pool!; one pointer compare instead of an IdDict lookup). + p._lookup_memo_type === T && return p._lookup_memo_tp::MetalTypedPool{T, Metal.PrivateStorage} + tp = get(p.others, T, nothing) + if tp !== nothing + tp = tp::MetalTypedPool{T, Metal.PrivateStorage} + p._lookup_memo_type = T + p._lookup_memo_tp = tp + return tp + end + # New type — create, register, memoize, and first-touch checkpoint when + # inside a scope (depth > 1), pushing one depth-tagged stack entry. + new_tp = MetalTypedPool{T, Metal.PrivateStorage}() + p.others[T] = new_tp + p._lookup_memo_type = T + p._lookup_memo_tp = new_tp + if p._current_depth > 1 + st = getfield(new_tp, :state) + push!(st._checkpoint_n_active, 0) # n_active starts at 0 + push!(st._checkpoint_depths, p._current_depth) + push!(p._touched_others_states, st) + push!(p._touched_others_depths, p._current_depth) + AdaptiveArrayPools._runtime_check(p) && push!(p._touched_others_pools, new_tp) + # Signal that a fallback type was touched so lazy/typed-lazy rewind + # iterates the drain path (same fix as CPU get_typed_pool!) + @inbounds p._touched_has_others[p._current_depth] = true + end + return new_tp end diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index dd243f75..28b59733 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -8,7 +8,14 @@ using AdaptiveArrayPools: checkpoint!, rewind!, reset!, _checkpoint_typed_pool!, _rewind_typed_pool!, _has_bit, - _LAZY_MODE_BIT, _TYPED_LAZY_BIT, _TYPE_BITS_MASK + _LAZY_MODE_BIT, _TYPED_LAZY_BIT, _TYPE_BITS_MASK, + _touch_fallback_pool!, _drain_touched_others!, _truncate_touched_others! + +# Genuine fallback = lives in pool.others (stack-managed). NOT equivalent to +# _fixed_slot_bit(T) == 0: Float16 has bit 0 (bit-7 reassignment) but is a fixed +# struct field — routing it through the touched-others stack would double-rewind +# it against the lazy rewinds' Float16 special case (Case A then Case B). +@inline _metal_is_fallback_type(::Type{T}) where {T} = !(T <: _METAL_FIXED_TYPES) # ============================================================================== # Metal Fixed Slot Iteration @@ -56,8 +63,15 @@ end @inline function AdaptiveArrayPools.checkpoint!(pool::MetalAdaptiveArrayPool, ::Type{T}) where {T} pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) + # Flag push stays bit-based (feeds _can_use_typed_path/R>=1 validation only) — + # Float16 has bit 0 here even though it is routed as a fixed slot below. push!(pool._touched_has_others, AdaptiveArrayPools._fixed_slot_bit(T) == UInt16(0)) - _checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth) + if _metal_is_fallback_type(T) + _touch_fallback_pool!(pool, AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth) + else + # Fixed slots INCLUDING Float16: direct checkpoint, never stack-managed. + _checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth) + end return nothing end @@ -71,8 +85,17 @@ end push!(unique_indices, i) end end + # has_any_fallback keeps its current bit-based computation (flag semantics + # unchanged — Float16 contributes true here even though it is routed as a + # fixed slot below via _metal_is_fallback_type). has_any_fallback = any(i -> AdaptiveArrayPools._fixed_slot_bit(types[i].parameters[1]) == UInt16(0), unique_indices) - checkpoint_exprs = [:(_checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth)) for i in unique_indices] + checkpoint_exprs = map(unique_indices) do i + if !(types[i].parameters[1] <: _METAL_FIXED_TYPES) + :(_touch_fallback_pool!(pool, AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth)) + else + :(_checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth)) + end + end return quote pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) @@ -104,6 +127,9 @@ function AdaptiveArrayPools.rewind!(pool::MetalAdaptiveArrayPool{R, S}) where {R for tp in values(pool.others) _rewind_typed_pool!(tp, cur_depth, R) end + # Full sweep above already rewound every fallback pool — truncate-only (no + # re-rewind) to avoid double-popping the touched-others stack. + _truncate_touched_others!(pool, cur_depth) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) @@ -118,7 +144,13 @@ end reset!(AdaptiveArrayPools.get_typed_pool!(pool, T), R) return nothing end - _rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth, R) + # Fixed slots (INCLUDING Float16) rewind directly; genuine-fallback T was + # pushed onto the touched-others stack by checkpoint!(pool, T) and is + # covered by the drain below. + if !_metal_is_fallback_type(T) + _rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth, R) + end + _drain_touched_others!(pool, pool._current_depth) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -135,7 +167,11 @@ end push!(unique_indices, i) end end - rewind_exprs = [:(_rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth, R)) for i in reverse(unique_indices)] + # Fixed slots INCLUDING Float16 rewind directly; genuine-fallback types were + # pushed onto the touched-others stack by checkpoint!(pool, types...) and + # are covered by the drain below. + fixed_indices = [i for i in unique_indices if !(types[i].parameters[1] <: _METAL_FIXED_TYPES) == false] + rewind_exprs = [:(_rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth, R)) for i in reverse(fixed_indices)] reset_exprs = [:(reset!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), R)) for i in unique_indices] return quote if pool._current_depth == 1 @@ -143,6 +179,7 @@ end return nothing end $(rewind_exprs...) + _drain_touched_others!(pool, pool._current_depth) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -166,15 +203,11 @@ end pool._current_depth += 1 push!(pool._touched_type_masks, _LAZY_MODE_BIT) # lazy mode flag push!(pool._touched_has_others, false) - depth = pool._current_depth - # Eagerly checkpoint pre-existing others entries — same as CPU _lazy_checkpoint!. - # New types created during the scope start at n_active=0 (sentinel covers them, Case B safe). - # Pre-existing types need their count saved now so Case A fires correctly at rewind. - for p in values(pool.others) - _checkpoint_typed_pool!(p, depth) - @inbounds pool._touched_has_others[depth] = true - end - # Float16 uses lazy first-touch via bit 7 in _record_type_touch! — no eager checkpoint needed. + # Fallback (non-fixed-slot) pools are NOT eagerly checkpointed here: they are + # first-touch checkpointed via _touch_fallback_pool! (from _record_type_touch! + # or get_typed_pool!) and drained selectively at rewind via + # _drain_touched_others!, so only the fallback pools this scope actually + # touches pay any cost. Float16 uses its own lazy first-touch via bit 7. return nothing end @@ -188,11 +221,7 @@ end _has_bit(mask, Bool) && _rewind_typed_pool!(pool.bool, d, R) # Bit 7: Float16 (Metal reassignment — _fixed_slot_bit(Float16)==0, must use explicit bit check) mask & _metal_float16_bit() != 0 && _rewind_typed_pool!(pool.float16, d, R) - if @inbounds(pool._touched_has_others[d]) - for tp in values(pool.others) - _rewind_typed_pool!(tp, d, R) - end - end + _drain_touched_others!(pool, d) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -204,27 +233,25 @@ end # ============================================================================== # _typed_lazy_checkpoint!: typed checkpoint + set bit 14 for lazy extra-type tracking. -# Also eagerly snapshots pre-existing others entries (mirrors CPU fix for Issue #3). +# checkpoint!(pool, types...) already routes fallback types among `types` through +# _touch_fallback_pool! (one depth-tagged stack entry each); extra fallback types +# touched by helpers are first-touch checkpointed and stacked by +# _record_type_touch!'s genuine-fallback branch. Float16 uses lazy first-touch via +# bit 7 in _record_type_touch! — no eager checkpoint needed. @inline function AdaptiveArrayPools._typed_lazy_checkpoint!(pool::MetalAdaptiveArrayPool, types::Type...) checkpoint!(pool, types...) d = pool._current_depth @inbounds pool._touched_type_masks[d] |= _TYPED_LAZY_BIT - # Eagerly snapshot pre-existing others entries — same reasoning as _lazy_checkpoint!. - # Skip re-snapshot for entries already checkpointed at d by checkpoint!(pool, types...) - for p in values(pool.others) - if @inbounds(p._checkpoint_depths[end]) != d - _checkpoint_typed_pool!(p, d) - end - @inbounds pool._touched_has_others[d] = true - end - # Float16 uses lazy first-touch via bit 7 in _record_type_touch! — no eager checkpoint needed. return nothing end # _typed_lazy_rewind!: selective rewind of (tracked | touched) mask. # Uses direct field access with bit checks — foreach_fixed_slot is single-argument (no bit yield). # Bit 7: Float16 (Metal-specific; lazy-checkpointed on first touch by _record_type_touch!). -# has_others: genuine others types (UInt8, Int8, etc.) — eagerly checkpointed at scope entry. +# Genuine fallback types (UInt8, Int8, etc.) are drained selectively via +# _drain_touched_others! — the ONLY rewinder for typed-Float16 scopes stays the +# direct _checkpoint_depths[end] == d special case below (Float16 never gets a +# stack entry). @inline function AdaptiveArrayPools._typed_lazy_rewind!(pool::MetalAdaptiveArrayPool{R, S}, tracked_mask::UInt16) where {R, S} d = pool._current_depth touched = @inbounds(pool._touched_type_masks[d]) & _TYPE_BITS_MASK @@ -239,11 +266,7 @@ end if combined & _metal_float16_bit() != 0 || @inbounds(pool.float16._checkpoint_depths[end]) == d _rewind_typed_pool!(pool.float16, d, R) end - if @inbounds(pool._touched_has_others[d]) - for tp in values(pool.others) - _rewind_typed_pool!(tp, d, R) - end - end + _drain_touched_others!(pool, d) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -265,6 +288,12 @@ function AdaptiveArrayPools.reset!(pool::MetalAdaptiveArrayPool{R, S}) where {R, reset!(tp, R) end + # Reset touched-others tracking (transient scope state; memo intentionally + # survives — registered fallback identities are preserved by reset!). + empty!(pool._touched_others_states) + empty!(pool._touched_others_depths) + empty!(pool._touched_others_pools) + # Reset depth and bitmask sentinel state pool._current_depth = 1 empty!(pool._touched_type_masks) @@ -321,6 +350,13 @@ function Base.empty!(pool::MetalAdaptiveArrayPool) end empty!(pool.others) + # Memo points into the registry being cleared — drop it with the registry. + pool._lookup_memo_type = nothing + pool._lookup_memo_tp = nothing + empty!(pool._touched_others_states) + empty!(pool._touched_others_depths) + empty!(pool._touched_others_pools) + # Reset depth and bitmask sentinel state pool._current_depth = 1 empty!(pool._touched_type_masks) diff --git a/test/metal/runtests.jl b/test/metal/runtests.jl index 3ee5fbe6..b8ef09d9 100644 --- a/test/metal/runtests.jl +++ b/test/metal/runtests.jl @@ -41,6 +41,7 @@ else # Include all Metal test files include("test_extension.jl") + include("test_touched_others.jl") include("test_allocation.jl") include("test_display.jl") include("test_convenience.jl") diff --git a/test/metal/test_extension.jl b/test/metal/test_extension.jl index a2bebc1f..db2836e6 100644 --- a/test/metal/test_extension.jl +++ b/test/metal/test_extension.jl @@ -108,6 +108,8 @@ end @test tp._checkpoint_n_active == [0, 0] @test tp._checkpoint_depths == [0, 2] @test pool._touched_has_others == [false, true] # depth-2 marked as "has dynamic types" + @test pool._touched_others_depths == [2] + @test length(pool._touched_others_states) == 1 end end diff --git a/test/metal/test_touched_others.jl b/test/metal/test_touched_others.jl new file mode 100644 index 00000000..2101924c --- /dev/null +++ b/test/metal/test_touched_others.jl @@ -0,0 +1,386 @@ +# Metal port of test/test_touched_others.jl — depth-tagged touched-others stack +# and fallback lookup memo. Invariant: a fallback pool has an entry tagged with +# depth d in the depth-tagged stack (states/depths, kept in lockstep) ⟺ it was +# first touched at depth d — except under full checkpoint!(pool), whose eager +# sweep pairs with full rewind!(pool)'s sweep (stack stays empty, truncate-only). +# +# Fallback exercise types (per task-3 brief): UInt16, UInt8, Int8, UInt32, Int16, +# UInt64, Int128, UInt128 — NOT Float64/ComplexF64 (rejected by the Metal backend). +# Float16 is a Metal-only divergence: a FIXED struct field with `_fixed_slot_bit == +# 0`, so it must NEVER be routed through the touched-others stack (see the +# dedicated "Float16 bit-7" testsets below). + +using AdaptiveArrayPools: _lazy_checkpoint!, _lazy_rewind!, + _typed_lazy_checkpoint!, _typed_lazy_rewind!, + _tracked_mask_for_types, _can_use_typed_path + +# Metal.jl's hardware backend rejects Int128/UInt128 ("Metal does not support +# Int128 values"), so the 8-type pollution-regression testset below substitutes +# two distinct isbits struct types for those two. Named distinctly from the CPU +# test suite's TOFooA/B/C (both files load into the same top-level module during +# a full-suite run) to avoid a struct-redefinition clash. +struct MetalFallbackStructA + x::Float32 +end +struct MetalFallbackStructB + x::Float32 +end + +@testset "touched-others: fields & lifecycle" begin + pool = MetalAdaptiveArrayPool() + @test isempty(pool._touched_others_states) && isempty(pool._touched_others_depths) && isempty(pool._touched_others_pools) + + # reset! clears transient scope state, keeps registry + acquire!(pool, UInt16, 4) + reset!(pool) + @test length(pool._touched_others_states) == 0 + @test isempty(pool._touched_others_depths) + @test haskey(pool.others, UInt16) # registry kept + + # empty! clears everything + acquire!(pool, UInt16, 4) + empty!(pool) + @test length(pool._touched_others_states) == 0 + @test isempty(pool._touched_others_depths) + @test !haskey(pool.others, UInt16) +end + +@testset "touched-others: checkpoint/rewind plumbing balance" begin + pool = MetalAdaptiveArrayPool() + + # lazy pair + _lazy_checkpoint!(pool) + @test isempty(pool._touched_others_depths) + _lazy_rewind!(pool) + @test isempty(pool._touched_others_depths) + + # typed single pair (fixed-slot type) + checkpoint!(pool, Float32) + @test isempty(pool._touched_others_depths) + rewind!(pool, Float32) + @test isempty(pool._touched_others_depths) + + # typed multi pair + checkpoint!(pool, Float32, Int32) + rewind!(pool, Float32, Int32) + @test isempty(pool._touched_others_depths) + + # full pair + checkpoint!(pool) + rewind!(pool) + @test isempty(pool._touched_others_depths) + + # typed-lazy pair + _typed_lazy_checkpoint!(pool, Float32) + _typed_lazy_rewind!(pool, _tracked_mask_for_types(Float32)) + @test isempty(pool._touched_others_depths) + + # nesting + _lazy_checkpoint!(pool) + checkpoint!(pool, Float32) + @test isempty(pool._touched_others_depths) + rewind!(pool, Float32) + _lazy_rewind!(pool) + @test isempty(pool._touched_others_depths) +end + +@testset "touched-others: no eager checkpoint on lazy entry (pollution regression)" begin + pool = MetalAdaptiveArrayPool() + # Register all 8 fallback exercise types at global scope, then reset counters. + # (Int128/UInt128 substituted with distinct isbits structs — Metal.jl's hardware + # backend rejects 128-bit integer element types.) + fallback_types = (UInt16, UInt8, Int8, UInt32, Int16, UInt64, MetalFallbackStructA, MetalFallbackStructB) + for T in fallback_types + acquire!(pool, T, 4) + end + reset!(pool) + tps = Dict(T => get_typed_pool!(pool, T) for T in fallback_types) + + _lazy_checkpoint!(pool) + # THE regression assertion: unrelated registered fallbacks are NOT touched + for T in fallback_types[2:end] + @test tps[T]._checkpoint_depths == [0] + end + + acquire!(pool, UInt16, 8) + @test tps[UInt16]._checkpoint_depths[end] == 2 # first-touch checkpoint at depth 2 + @test length(pool._touched_others_depths) == 1 + @test pool._touched_others_states[end] === tps[UInt16].state + @test tps[UInt16].n_active == 1 + + # Re-acquire same type: no duplicate stack entry + acquire!(pool, UInt16, 8) + @test length(pool._touched_others_depths) == 1 + + _lazy_rewind!(pool) + @test tps[UInt16].n_active == 0 + @test isempty(pool._touched_others_depths) + for T in fallback_types[2:end] + @test tps[T]._checkpoint_depths == [0] # still never visited + end +end + +@testset "touched-others: nested scopes, same fallback type at two depths" begin + pool = MetalAdaptiveArrayPool() + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, UInt16, 4) + @test tpA.n_active == 1 + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, UInt16, 4) + acquire!(pool, UInt16, 4) + @test tpA.n_active == 3 + @test length(pool._touched_others_depths) == 2 # one entry per depth + + _lazy_rewind!(pool) # exit depth 3 + @test tpA.n_active == 1 + + _lazy_rewind!(pool) # exit depth 2 + @test tpA.n_active == 0 +end + +@testset "touched-others: depth tags are exact and monotone" begin + pool = MetalAdaptiveArrayPool() + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + @test isempty(pool._touched_others_depths) # entry pushes nothing + acquire!(pool, UInt16, 4) + @test pool._touched_others_depths == [2] + @test pool._touched_others_states[end] === tpA.state + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, UInt8, 4) # different fallback type, new depth + @test pool._touched_others_depths == [2, 3] + @test issorted(pool._touched_others_depths) # monotone invariant + + _lazy_rewind!(pool) # drains ONLY the ==3 entries + @test pool._touched_others_depths == [2] + @test tpA.n_active == 1 + @test get_typed_pool!(pool, UInt8).n_active == 0 + + _lazy_rewind!(pool) + @test isempty(pool._touched_others_depths) + @test tpA.n_active == 0 +end + +@testset "touched-others: nested scope NOT touching outer's fallback" begin + pool = MetalAdaptiveArrayPool() + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, UInt16, 4) + _lazy_checkpoint!(pool) # depth 3: does not touch UInt16 + acquire!(pool, Float32, 16) # Float32 fixed-slot work only + @test length(pool._touched_others_depths) == 1 # no new fallback entry + _lazy_rewind!(pool) + @test tpA.n_active == 1 # outer's array untouched + _lazy_rewind!(pool) + @test tpA.n_active == 0 +end + +@testset "touched-others: typed scope with helper touching a fallback (typed-lazy)" begin + pool = MetalAdaptiveArrayPool() + tpB = get_typed_pool!(pool, UInt8) + + _typed_lazy_checkpoint!(pool, Float32) + acquire!(pool, Float32, 8) # tracked fixed-slot work + acquire!(pool, UInt8, 4) # untracked helper-style fallback touch + @test pool._touched_has_others[end] == true + @test pool._touched_others_states[end] === tpB.state + @test !_can_use_typed_path(pool, _tracked_mask_for_types(Float32)) + _typed_lazy_rewind!(pool, _tracked_mask_for_types(Float32)) + @test tpB.n_active == 0 +end + +@testset "touched-others: tracked fallback type via typed checkpoint!" begin + pool = MetalAdaptiveArrayPool() + + checkpoint!(pool, UInt16) # fallback T tracked by macro + tpA = get_typed_pool!(pool, UInt16) + @test pool._touched_others_states[end] === tpA.state # pushed at checkpoint + acquire!(pool, UInt16, 4) # public-API acquire: no double push + @test count(st -> st === tpA.state, pool._touched_others_states) == 1 + # macro exit path for has_others=true is _typed_lazy_rewind! + _typed_lazy_rewind!(pool, _tracked_mask_for_types(UInt16)) + @test tpA.n_active == 0 + @test isempty(pool._touched_others_depths) +end + +@testset "touched-others: new type registered mid-scope" begin + pool = MetalAdaptiveArrayPool() + _lazy_checkpoint!(pool) + acquire!(pool, Int8, 4) # first-ever registration, in-scope + tpC = get_typed_pool!(pool, Int8) + @test pool._touched_others_states[end] === tpC.state + @test count(st -> st === tpC.state, pool._touched_others_states) == 1 + _lazy_rewind!(pool) + @test tpC.n_active == 0 +end + +@testset "touched-others: full checkpoint!/rewind! pairing unchanged" begin + pool = MetalAdaptiveArrayPool() + acquire!(pool, UInt16, 4) + reset!(pool) + tpA = get_typed_pool!(pool, UInt16) + + checkpoint!(pool) # eager: checkpoints ALL others + @test tpA._checkpoint_depths[end] == 2 + acquire!(pool, UInt16, 4) + @test length(pool._touched_others_depths) == 0 # guard saw existing depth-2 entry + rewind!(pool) + @test tpA.n_active == 0 + @test isempty(pool._touched_others_depths) +end + +@testset "touched-others: similar! records fallback touch" begin + pool = MetalAdaptiveArrayPool() + src = MtlArray(UInt16[1, 2]) + _lazy_checkpoint!(pool) + similar!(pool, src) + tpA = get_typed_pool!(pool, UInt16) + @test pool._touched_others_states[end] === tpA.state + _lazy_rewind!(pool) + @test tpA.n_active == 0 +end + +@testset "touched-others: @with_pool integration + exception-leak recovery" begin + # The task-local Metal pool is a process-wide-per-device singleton shared with + # every other test file. Start from a known-clean baseline so the absolute-depth + # assertions below are independent of test file execution order. + empty!(get_task_local_metal_pool()) + + # Integration through the real macro (task-local pool) + f_leaf(n) = @with_pool :metal p begin + q = acquire!(p, UInt16, n) + length(q) + end + @test f_leaf(8) == 8 + tl = get_task_local_metal_pool() + @test get_typed_pool!(tl, UInt16).n_active == 0 + @test isempty(tl._touched_others_states) + + # Inner scope throws, outer catches: outer exit must clean up leaked state + function f_outer() + @with_pool :metal p begin + acquire!(p, UInt8, 4) + try + @with_pool :metal p2 begin + acquire!(p2, Int8, 4) + error("boom") + end + catch + end + 1 + end + end + @test f_outer() == 1 + @test tl._current_depth == 1 + @test get_typed_pool!(tl, UInt8).n_active == 0 + @test get_typed_pool!(tl, Int8).n_active == 0 + @test isempty(tl._touched_others_states) + @test isempty(tl._touched_others_depths) + empty!(tl) # leave the task-local pool clean for other test files +end + +@testset "fallback lookup memo: fields and lifecycle" begin + pool = MetalAdaptiveArrayPool() + tp = get_typed_pool!(pool, UInt16) # slow-path lookup + @test pool._lookup_memo_type === UInt16 + @test pool._lookup_memo_tp === tp + + empty!(pool) + @test pool._lookup_memo_type === nothing + @test pool._lookup_memo_tp === nothing + + tp2 = get_typed_pool!(pool, UInt16) # re-register + @test pool._lookup_memo_type === UInt16 + @test pool._lookup_memo_tp === tp2 + + reset!(pool) # keeps registry AND memo + @test pool._lookup_memo_type === UInt16 + @test pool._lookup_memo_tp === tp2 +end + +@testset "fallback lookup memo: identity and invalidation" begin + pool = MetalAdaptiveArrayPool() + tp1 = get_typed_pool!(pool, UInt16) + @test get_typed_pool!(pool, UInt16) === tp1 # repeat lookup: same pool + @test get_typed_pool!(pool, UInt8) !== tp1 # different type: different pool + @test get_typed_pool!(pool, UInt16) === tp1 # alternating types stay correct + + reset!(pool) # keeps registry → memo may stay + @test get_typed_pool!(pool, UInt16) === tp1 + + empty!(pool) # kills registry → memo MUST die + tp2 = get_typed_pool!(pool, UInt16) + @test tp2 !== tp1 # stale-memo regression guard + @test tp2 === pool.others[UInt16] + + # end-to-end: acquire after empty! must use the fresh pool + acquire!(pool, UInt16, 4) + @test tp2.n_active == 1 && tp1.n_active == 0 + reset!(pool) +end + +# ============================================================================== +# Float16 bit-7 non-interaction (GPU-only divergence — CPU has no such field) +# ============================================================================== +# Float16 is a FIXED struct field (pool.float16) but `_fixed_slot_bit(Float16) == +# 0` (Float16 is absent from the base module's generic fixed-slot bit table). +# Routing MUST key off `_metal_is_fallback_type(T) = !(T <: _METAL_FIXED_TYPES)`, +# never `_fixed_slot_bit(T) == 0` — Float16 goes through the direct +# `_checkpoint_typed_pool!`/`_rewind_typed_pool!` path like other fixed slots and +# must NEVER get a touched-others stack entry. + +@testset "Float16 bit-7: lazy scope acquiring Float16 pushes no stack entry" begin + pool = MetalAdaptiveArrayPool() + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, Float16, 4) + @test isempty(pool._touched_others_depths) + @test pool.float16.n_active == 1 + _lazy_rewind!(pool) + @test pool.float16.n_active == 0 + @test isempty(pool._touched_others_depths) +end + +@testset "Float16 bit-7: typed inner scope alongside live parent Float16 arrays" begin + empty!(get_task_local_metal_pool()) + pool = get_task_local_metal_pool() + + _lazy_checkpoint!(pool) # depth 2 (parent, lazy) + v_outer = acquire!(pool, Float16, 4) + v_outer .= Float16(3) + @test pool.float16.n_active == 1 + + result = @with_pool :metal pool2 begin + w = acquire!(pool2, Float16, 8) + w .= Float16(9) + sum(w) + end + @test result == Float16(9) * 8 + @test isempty(pool._touched_others_depths) # Float16 never stack-managed + @test pool.float16.n_active == 1 # parent's slot survives the inner exit + @test all(Array(v_outer) .== Float16(3)) # parent's array contents survive + + _lazy_rewind!(pool) + @test pool.float16.n_active == 0 + empty!(get_task_local_metal_pool()) +end + +@testset "Float16 bit-7: mixed scope Float16 + fallback pushes exactly one entry" begin + pool = MetalAdaptiveArrayPool() + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, Float16, 4) + acquire!(pool, UInt16, 4) + @test length(pool._touched_others_depths) == 1 + @test pool._touched_others_depths == [2] + tpA = get_typed_pool!(pool, UInt16) + @test pool._touched_others_states[end] === tpA.state + _lazy_rewind!(pool) + @test pool.float16.n_active == 0 + @test tpA.n_active == 0 + @test isempty(pool._touched_others_depths) +end From 7e8184e1517c84a718e573b9d9aa3dc9f295dfd9 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 22:55:54 -0700 Subject: [PATCH 04/10] (test): exercise typed Float16 scope in touched-others regression guard --- ext/AdaptiveArrayPoolsMetalExt/state.jl | 2 +- test/metal/test_touched_others.jl | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/ext/AdaptiveArrayPoolsMetalExt/state.jl b/ext/AdaptiveArrayPoolsMetalExt/state.jl index 28b59733..5b9ac2fa 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/state.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/state.jl @@ -170,7 +170,7 @@ end # Fixed slots INCLUDING Float16 rewind directly; genuine-fallback types were # pushed onto the touched-others stack by checkpoint!(pool, types...) and # are covered by the drain below. - fixed_indices = [i for i in unique_indices if !(types[i].parameters[1] <: _METAL_FIXED_TYPES) == false] + fixed_indices = [i for i in unique_indices if types[i].parameters[1] <: _METAL_FIXED_TYPES] rewind_exprs = [:(_rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth, R)) for i in reverse(fixed_indices)] reset_exprs = [:(reset!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), R)) for i in unique_indices] return quote diff --git a/test/metal/test_touched_others.jl b/test/metal/test_touched_others.jl index 2101924c..591de4b3 100644 --- a/test/metal/test_touched_others.jl +++ b/test/metal/test_touched_others.jl @@ -349,14 +349,21 @@ end @testset "Float16 bit-7: typed inner scope alongside live parent Float16 arrays" begin empty!(get_task_local_metal_pool()) pool = get_task_local_metal_pool() + depth_before = pool._current_depth _lazy_checkpoint!(pool) # depth 2 (parent, lazy) v_outer = acquire!(pool, Float16, 4) v_outer .= Float16(3) @test pool.float16.n_active == 1 + # acquire!(pool2, Float16, 8) below is a static-type call, so the macro + # statically resolves Float16 and drives the TYPED path + # (_typed_lazy_checkpoint!/_typed_lazy_rewind!) rather than the untyped + # _lazy_checkpoint!/_lazy_rewind! path — no separate typed-macro syntax + # exists or is needed to opt in. result = @with_pool :metal pool2 begin w = acquire!(pool2, Float16, 8) + @test isempty(pool._touched_others_depths) # Float16 never stack-managed, even in typed scope w .= Float16(9) sum(w) end @@ -367,6 +374,7 @@ end _lazy_rewind!(pool) @test pool.float16.n_active == 0 + @test pool._current_depth == depth_before # depth balance: no leak across the testset empty!(get_task_local_metal_pool()) end From ad2bdd8b4493ad956ffc84d085e33a37c7915291 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 23:38:08 -0700 Subject: [PATCH 05/10] =?UTF-8?q?(perf):=20CUDA=20=E2=80=94=20mirror=20dep?= =?UTF-8?q?th-tagged=20stack,=20PoolCheckpointState,=20lookup=20memo?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../AdaptiveArrayPoolsCUDAExt.jl | 2 +- ext/AdaptiveArrayPoolsCUDAExt/acquire.jl | 9 +- ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl | 41 +- ext/AdaptiveArrayPoolsCUDAExt/state.jl | 107 +++-- ext/AdaptiveArrayPoolsCUDAExt/types.jl | 100 +++-- test/cuda/runtests.jl | 1 + test/cuda/test_extension.jl | 19 +- test/cuda/test_touched_others.jl | 394 ++++++++++++++++++ test/test_state.jl | 20 +- 9 files changed, 591 insertions(+), 102 deletions(-) create mode 100644 test/cuda/test_touched_others.jl diff --git a/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl b/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl index f8602e9b..f55d0b00 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl @@ -16,7 +16,7 @@ using CUDA # On older Julia, the extension loads but provides no functionality. @static if VERSION >= v"1.12-" - using AdaptiveArrayPools: AbstractTypedPool, AbstractArrayPool + using AdaptiveArrayPools: AbstractTypedPool, AbstractArrayPool, PoolCheckpointState # Type definitions include("types.jl") diff --git a/ext/AdaptiveArrayPoolsCUDAExt/acquire.jl b/ext/AdaptiveArrayPoolsCUDAExt/acquire.jl index 8ca46d0f..f313e316 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/acquire.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/acquire.jl @@ -34,7 +34,7 @@ using AdaptiveArrayPools: get_view!, get_array!, allocate_vector, safe_prod, _record_type_touch!, _fixed_slot_bit, _checkpoint_typed_pool!, _store_arr_wrapper!, _check_pool_growth, _reshape_impl!, _acquire_impl!, _acquire_view_impl!, _maybe_record_borrow!, - _MODE_BITS_MASK + _MODE_BITS_MASK, _touch_fallback_pool! using CUDA: unsafe_free! @@ -339,8 +339,13 @@ end end @inbounds pool._touched_type_masks[depth] = current_mask | b16 else - # Genuine others type (UInt8, Int8, etc.) — eagerly snapshotted at scope entry. + # Genuine others type (UInt8, Int8, etc.). @inbounds pool._touched_has_others[depth] = true + # First-touch lazy checkpoint for fallback types; depth == 1 (global + # scope) is exempt — matches get_typed_pool!'s gate. + if depth > 1 + _touch_fallback_pool!(pool, AdaptiveArrayPools.get_typed_pool!(pool, T), depth) + end end else current_mask = @inbounds pool._touched_type_masks[depth] diff --git a/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl b/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl index 968c84c4..948c3d55 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl @@ -33,17 +33,32 @@ const _CUDA_FIXED_TYPES = Union{Float32, Float64, Float16, Int32, Int64, Complex # Slow path: rare types via IdDict (with checkpoint correction!) @inline function AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{T}) where {T} - return get!(p.others, T) do - tp = CuTypedPool{T}() - # CRITICAL: Match CPU behavior - auto-checkpoint new pool if inside @with_pool scope - # Without this, rewind! would corrupt state for dynamically-created pools - if p._current_depth > 1 - push!(tp._checkpoint_n_active, 0) # n_active starts at 0 - push!(tp._checkpoint_depths, p._current_depth) - # Signal that a fallback type was touched so lazy/typed-lazy rewind - # iterates pool.others (same fix as CPU get_typed_pool!) - @inbounds p._touched_has_others[p._current_depth] = true - end - tp - end::CuTypedPool{T} + # Memo fast path: same type as the previous slow-path lookup (mirror of CPU + # src/types.jl's get_typed_pool!; one pointer compare instead of an IdDict lookup). + p._lookup_memo_type === T && return p._lookup_memo_tp::CuTypedPool{T} + tp = get(p.others, T, nothing) + if tp !== nothing + tp = tp::CuTypedPool{T} + p._lookup_memo_type = T + p._lookup_memo_tp = tp + return tp + end + # New type — create, register, memoize, and first-touch checkpoint when + # inside a scope (depth > 1), pushing one depth-tagged stack entry. + new_tp = CuTypedPool{T}() + p.others[T] = new_tp + p._lookup_memo_type = T + p._lookup_memo_tp = new_tp + if p._current_depth > 1 + st = getfield(new_tp, :state) + push!(st._checkpoint_n_active, 0) # n_active starts at 0 + push!(st._checkpoint_depths, p._current_depth) + push!(p._touched_others_states, st) + push!(p._touched_others_depths, p._current_depth) + AdaptiveArrayPools._runtime_check(p) && push!(p._touched_others_pools, new_tp) + # Signal that a fallback type was touched so lazy/typed-lazy rewind + # iterates the drain path (same fix as CPU get_typed_pool!) + @inbounds p._touched_has_others[p._current_depth] = true + end + return new_tp end diff --git a/ext/AdaptiveArrayPoolsCUDAExt/state.jl b/ext/AdaptiveArrayPoolsCUDAExt/state.jl index b465f8db..8c9c28aa 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/state.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/state.jl @@ -8,7 +8,14 @@ using AdaptiveArrayPools: checkpoint!, rewind!, reset!, _checkpoint_typed_pool!, _rewind_typed_pool!, _has_bit, - _LAZY_MODE_BIT, _TYPED_LAZY_BIT, _TYPE_BITS_MASK + _LAZY_MODE_BIT, _TYPED_LAZY_BIT, _TYPE_BITS_MASK, + _touch_fallback_pool!, _drain_touched_others!, _truncate_touched_others! + +# Genuine fallback = lives in pool.others (stack-managed). NOT equivalent to +# _fixed_slot_bit(T) == 0: Float16 has bit 0 (bit-7 reassignment) but is a fixed +# struct field — routing it through the touched-others stack would double-rewind +# it against the lazy rewinds' Float16 special case (Case A then Case B). +@inline _cuda_is_fallback_type(::Type{T}) where {T} = !(T <: _CUDA_FIXED_TYPES) # ============================================================================== # GPU Fixed Slot Iteration @@ -56,8 +63,15 @@ end @inline function AdaptiveArrayPools.checkpoint!(pool::CuAdaptiveArrayPool, ::Type{T}) where {T} pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) + # Flag push stays bit-based (feeds _can_use_typed_path/S>=1 validation only) — + # Float16 has bit 0 here even though it is routed as a fixed slot below. push!(pool._touched_has_others, AdaptiveArrayPools._fixed_slot_bit(T) == UInt16(0)) - _checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth) + if _cuda_is_fallback_type(T) + _touch_fallback_pool!(pool, AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth) + else + # Fixed slots INCLUDING Float16: direct checkpoint, never stack-managed. + _checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth) + end return nothing end @@ -71,8 +85,17 @@ end push!(unique_indices, i) end end + # has_any_fallback keeps its current bit-based computation (flag semantics + # unchanged — Float16 contributes true here even though it is routed as a + # fixed slot below via _cuda_is_fallback_type). has_any_fallback = any(i -> AdaptiveArrayPools._fixed_slot_bit(types[i].parameters[1]) == UInt16(0), unique_indices) - checkpoint_exprs = [:(_checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth)) for i in unique_indices] + checkpoint_exprs = map(unique_indices) do i + if !(types[i].parameters[1] <: _CUDA_FIXED_TYPES) + :(_touch_fallback_pool!(pool, AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth)) + else + :(_checkpoint_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth)) + end + end return quote pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) @@ -104,6 +127,9 @@ function AdaptiveArrayPools.rewind!(pool::CuAdaptiveArrayPool{S}) where {S} for tp in values(pool.others) _rewind_typed_pool!(tp, cur_depth, S) end + # Full sweep above already rewound every fallback pool — truncate-only (no + # re-rewind) to avoid double-popping the touched-others stack. + _truncate_touched_others!(pool, cur_depth) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) @@ -118,7 +144,13 @@ end reset!(AdaptiveArrayPools.get_typed_pool!(pool, T), S) return nothing end - _rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth, S) + # Fixed slots (INCLUDING Float16) rewind directly; genuine-fallback T was + # pushed onto the touched-others stack by checkpoint!(pool, T) and is + # covered by the drain below. + if !_cuda_is_fallback_type(T) + _rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, T), pool._current_depth, S) + end + _drain_touched_others!(pool, pool._current_depth) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -135,7 +167,11 @@ end push!(unique_indices, i) end end - rewind_exprs = [:(_rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth, S)) for i in reverse(unique_indices)] + # Fixed slots INCLUDING Float16 rewind directly; genuine-fallback types were + # pushed onto the touched-others stack by checkpoint!(pool, types...) and + # are covered by the drain below. + fixed_indices = [i for i in unique_indices if types[i].parameters[1] <: _CUDA_FIXED_TYPES] + rewind_exprs = [:(_rewind_typed_pool!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), pool._current_depth, S)) for i in reverse(fixed_indices)] reset_exprs = [:(reset!(AdaptiveArrayPools.get_typed_pool!(pool, types[$i]), S)) for i in unique_indices] return quote if pool._current_depth == 1 @@ -143,6 +179,7 @@ end return nothing end $(rewind_exprs...) + _drain_touched_others!(pool, pool._current_depth) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -167,15 +204,11 @@ end pool._current_depth += 1 push!(pool._touched_type_masks, _LAZY_MODE_BIT) # lazy mode flag push!(pool._touched_has_others, false) - depth = pool._current_depth - # Eagerly checkpoint pre-existing others entries — same as CPU _lazy_checkpoint!. - # New types created during the scope start at n_active=0 (sentinel covers them, Case B safe). - # Pre-existing types need their count saved now so Case A fires correctly at rewind. - for p in values(pool.others) - _checkpoint_typed_pool!(p, depth) - @inbounds pool._touched_has_others[depth] = true - end - # Float16 uses lazy first-touch via bit 7 in _record_type_touch! — no eager checkpoint needed. + # Fallback (non-fixed-slot) pools are NOT eagerly checkpointed here: they are + # first-touch checkpointed via _touch_fallback_pool! (from _record_type_touch! + # or get_typed_pool!) and drained selectively at rewind via + # _drain_touched_others!, so only the fallback pools this scope actually + # touches pay any cost. Float16 uses its own lazy first-touch via bit 7. return nothing end @@ -191,11 +224,7 @@ end _has_bit(mask, Bool) && _rewind_typed_pool!(pool.bool, d, S) # Bit 7: Float16 (CUDA reassignment — _fixed_slot_bit(Float16)==0, must use explicit bit check) mask & _cuda_float16_bit() != 0 && _rewind_typed_pool!(pool.float16, d, S) - if @inbounds(pool._touched_has_others[d]) - for tp in values(pool.others) - _rewind_typed_pool!(tp, d, S) - end - end + _drain_touched_others!(pool, d) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -207,28 +236,25 @@ end # ============================================================================== # _typed_lazy_checkpoint!: typed checkpoint + set bit 14 for lazy extra-type tracking. -# Also eagerly snapshots pre-existing others entries (mirrors CPU fix for Issue #3). +# checkpoint!(pool, types...) already routes fallback types among `types` through +# _touch_fallback_pool! (one depth-tagged stack entry each); extra fallback types +# touched by helpers are first-touch checkpointed and stacked by +# _record_type_touch!'s genuine-fallback branch. Float16 uses lazy first-touch via +# bit 7 in _record_type_touch! — no eager checkpoint needed. @inline function AdaptiveArrayPools._typed_lazy_checkpoint!(pool::CuAdaptiveArrayPool, types::Type...) checkpoint!(pool, types...) d = pool._current_depth @inbounds pool._touched_type_masks[d] |= _TYPED_LAZY_BIT - # Eagerly snapshot pre-existing others entries — same reasoning as _lazy_checkpoint!. - # Skip re-snapshot for entries already checkpointed at d by checkpoint!(pool, types...) - # (e.g. Float16 in types... was just checkpointed above — avoid double-push). - for p in values(pool.others) - if @inbounds(p._checkpoint_depths[end]) != d - _checkpoint_typed_pool!(p, d) - end - @inbounds pool._touched_has_others[d] = true - end - # Float16 uses lazy first-touch via bit 7 in _record_type_touch! — no eager checkpoint needed. return nothing end # _typed_lazy_rewind!: selective rewind of (tracked | touched) mask. # Uses direct field access with bit checks — foreach_fixed_slot is single-argument (no bit yield). # Bit 7: Float16 (CUDA-specific; lazy-checkpointed on first touch by _record_type_touch!). -# has_others: genuine others types (UInt8, Int8, etc.) — eagerly checkpointed at scope entry. +# Genuine fallback types (UInt8, Int8, etc.) are drained selectively via +# _drain_touched_others! — the ONLY rewinder for typed-Float16 scopes stays the +# direct _checkpoint_depths[end] == d special case below (Float16 never gets a +# stack entry). @inline function AdaptiveArrayPools._typed_lazy_rewind!(pool::CuAdaptiveArrayPool{S}, tracked_mask::UInt16) where {S} d = pool._current_depth touched = @inbounds(pool._touched_type_masks[d]) & _TYPE_BITS_MASK @@ -250,11 +276,7 @@ end if combined & _cuda_float16_bit() != 0 || @inbounds(pool.float16._checkpoint_depths[end]) == d _rewind_typed_pool!(pool.float16, d, S) end - if @inbounds(pool._touched_has_others[d]) - for tp in values(pool.others) - _rewind_typed_pool!(tp, d, S) - end - end + _drain_touched_others!(pool, d) pop!(pool._touched_type_masks) pop!(pool._touched_has_others) pool._current_depth -= 1 @@ -276,6 +298,12 @@ function AdaptiveArrayPools.reset!(pool::CuAdaptiveArrayPool{S}) where {S} reset!(tp, S) end + # Reset touched-others tracking (transient scope state; memo intentionally + # survives — registered fallback identities are preserved by reset!). + empty!(pool._touched_others_states) + empty!(pool._touched_others_depths) + empty!(pool._touched_others_pools) + # Reset depth and bitmask sentinel state pool._current_depth = 1 empty!(pool._touched_type_masks) @@ -339,6 +367,13 @@ function Base.empty!(pool::CuAdaptiveArrayPool) end empty!(pool.others) + # Memo points into the registry being cleared — drop it with the registry. + pool._lookup_memo_type = nothing + pool._lookup_memo_tp = nothing + empty!(pool._touched_others_states) + empty!(pool._touched_others_depths) + empty!(pool._touched_others_pools) + # Reset depth and bitmask sentinel state pool._current_depth = 1 empty!(pool._touched_type_masks) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/types.jl b/ext/AdaptiveArrayPoolsCUDAExt/types.jl index e885b7ea..71183a44 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/types.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/types.jl @@ -41,10 +41,11 @@ mutable struct CuTypedPool{T} <: AbstractTypedPool{T, CuVector{T}} # (`_slot_used`) to know how much of an over-allocated device buffer is in use. slot_extents::Vector{Int} - # --- State Management (1-based sentinel pattern) --- - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} + # --- State Management --- + # Checkpoint bookkeeping, extracted into a concrete shared struct (see CPU + # PoolCheckpointState docstring). `const`: the reference never changes after + # construction. Accessed as tp.n_active / tp._checkpoint_* via forwarding below. + const state::PoolCheckpointState # --- Auto-trim telemetry (parity with CPU TypedPool; see its docstring) --- # Peak `n_active` since the last auto-trim — the recent working-set width. Written on the @@ -58,11 +59,32 @@ function CuTypedPool{T}() where {T} CuVector{T}[], # vectors Union{Nothing, Vector{Any}}[], # arr_wrappers (indexed by N) Int[], # slot_extents (parallel to vectors) - 0, [0], [0], # State (1-based sentinel) + PoolCheckpointState(), # state (1-based sentinel) 0, # _am_peak_n_active: no usage observed yet ) end +# Checkpoint-state property forwarding (mirror of CPU src/types.jl:294-307). +@inline function Base.getproperty(tp::CuTypedPool, f::Symbol) + f === :n_active && return getfield(tp, :state).n_active + f === :_checkpoint_n_active && return getfield(tp, :state)._checkpoint_n_active + f === :_checkpoint_depths && return getfield(tp, :state)._checkpoint_depths + return getfield(tp, f) +end + +@inline function Base.setproperty!(tp::CuTypedPool, f::Symbol, v) + f === :n_active && return setfield!(getfield(tp, :state), :n_active, convert(Int, v)) + return setfield!(tp, f, convert(fieldtype(typeof(tp), f), v)) +end + +Base.propertynames(tp::CuTypedPool) = + (fieldnames(typeof(tp))..., :n_active, :_checkpoint_n_active, :_checkpoint_depths) + +# Route the generic checkpoint/rewind cores at the concrete state (zero-dispatch +# drain); without this, CuTypedPool falls back to _cp_state(tp) = tp and +# _touch_fallback_pool!'s ::PoolCheckpointState assert throws. +@inline AdaptiveArrayPools._cp_state(tp::CuTypedPool) = getfield(tp, :state) + # ============================================================================== # GPU Fixed Slot Configuration # ============================================================================== @@ -125,6 +147,20 @@ mutable struct CuAdaptiveArrayPool{S} <: AbstractArrayPool _touched_type_masks::Vector{UInt16} # Per-depth: which fixed slots were touched + mode flags _touched_has_others::Vector{Bool} # Per-depth: any non-fixed-slot type touched? + # Touched-others tracking (depth-tagged, concrete) — mirror of CPU + # src/types.jl:434-453. Checkpoint variants push NOTHING; producers push one + # (state, depth) entry per first touch; rewind pops while the top tag matches. + # _touched_others_pools is populated only when S >= 1 (slot invalidation). + _touched_others_states::Vector{PoolCheckpointState} + _touched_others_depths::Vector{Int} + _touched_others_pools::Vector{Any} + + # Last-lookup memo for the fallback registry (mirror of CPU). Set on every + # slow-path lookup; cleared by empty! (identities die), preserved by + # reset!/trim!/compact! (identities survive). Task-local pool → no races. + _lookup_memo_type::Any + _lookup_memo_tp::Any + # Device tracking (safety) device_id::Int @@ -158,6 +194,11 @@ function CuAdaptiveArrayPool{S}() where {S} 1, # _current_depth (1 = global scope) [UInt16(0)], # _touched_type_masks: sentinel (no bits set) [false], # _touched_has_others: sentinel (no others) + PoolCheckpointState[], # _touched_others_states: no fallback touches yet + Int[], # _touched_others_depths + Any[], # _touched_others_pools + nothing, # _lookup_memo_type + nothing, # _lookup_memo_tp CUDA.deviceid(dev), "", # _pending_callsite "", # _pending_return_site @@ -183,6 +224,15 @@ Return compile-time constant indicating whether runtime safety checks are enable @inline AdaptiveArrayPools._runtime_check(::CuAdaptiveArrayPool{0}) = false @inline AdaptiveArrayPools._runtime_check(::CuAdaptiveArrayPool) = true # S >= 1 +""" + _check_level(pool::CuAdaptiveArrayPool) -> Int + +Runtime-check level as an Int (mirror of CPU `src/types.jl:513`), for +backend-shared code that forwards it to `_invalidate_released_slots!` / +`_rewind_typed_pool!`. Compile-time constant per concrete pool type. +""" +@inline AdaptiveArrayPools._check_level(::CuAdaptiveArrayPool{S}) where {S} = S + """ _make_cuda_pool(level) -> CuAdaptiveArrayPool @@ -195,40 +245,12 @@ _make_cuda_pool(runtime_check::Bool) = _make_cuda_pool(Int(runtime_check)) return CuAdaptiveArrayPool{1}() end -""" - _make_cuda_pool(level, old::CuAdaptiveArrayPool) -> CuAdaptiveArrayPool - -Create a new CUDA pool, transferring cached arrays and scope state from `old`. -Only reference copies — no memory allocation for underlying GPU buffers. - -Transferred: all CuTypedPool slots, `others`, depth & touch tracking, device_id. -Reset: `_pending_callsite/return_site` (transient macro state), - `_borrow_log` (created fresh when S >= 1). -""" -_make_cuda_pool(runtime_check::Bool, old::CuAdaptiveArrayPool) = _make_cuda_pool(Int(runtime_check), old) -@noinline function _make_cuda_pool(level::Int, old::CuAdaptiveArrayPool) - level == 0 && return _transfer_cuda_pool(Val(0), old) - return _transfer_cuda_pool(Val(1), old) -end - -"""Transfer cached arrays and scope state from `old` pool into a new `CuAdaptiveArrayPool{V}`.""" -function _transfer_cuda_pool(::Val{V}, old::CuAdaptiveArrayPool) where {V} - return CuAdaptiveArrayPool{V}( - old.float32, old.float64, old.float16, - old.int32, old.int64, - old.complexf32, old.complexf64, old.bool, - old.others, - old._current_depth, - old._touched_type_masks, - old._touched_has_others, - old.device_id, - "", # _pending_callsite: reset - "", # _pending_return_site: reset - V >= 1 ? IdDict{Any, String}() : nothing, # _borrow_log - false, # _compact_requested: reset on migration - false # _trim_requested: reset on migration - ) -end +# NOTE: a former 2-arg `_make_cuda_pool(level, old)` overload migrated a pool +# across S while transferring caches by reference. Removed (no production +# callers): the touched-others stack's shape is S-dependent (pools entries +# exist only at S >= 1), so a mid-scope migration would desync the next +# rewind. If runtime S switching is ever needed, reintroduce it with an +# explicit global-scope guard. """Human-readable runtime check label.""" function _cuda_check_label(s::Int) diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl index 17dc64d6..d19b3f91 100644 --- a/test/cuda/runtests.jl +++ b/test/cuda/runtests.jl @@ -40,6 +40,7 @@ else # Include all CUDA test files include("test_extension.jl") + include("test_touched_others.jl") include("test_allocation.jl") include("test_nway_cache.jl") include("test_display.jl") diff --git a/test/cuda/test_extension.jl b/test/cuda/test_extension.jl index 247ff373..7d876b3e 100644 --- a/test/cuda/test_extension.jl +++ b/test/cuda/test_extension.jl @@ -5,12 +5,21 @@ @testset "CuTypedPool structure" begin tp_fields = fieldnames(CuTypedPool) @test :vectors in tp_fields - @test :n_active in tp_fields + @test :state in tp_fields # PoolCheckpointState (n_active + checkpoint vectors) + @test !(:n_active in tp_fields) # moved into state; reachable via property forwarding # arr_wrappers (setfield!-based wrapper reuse, replaces N-way cache) @test :arr_wrappers in tp_fields - # State management - @test :_checkpoint_n_active in tp_fields - @test :_checkpoint_depths in tp_fields + # Forwarding round-trip (mirrors CPU src/types.jl:294-307) + tp = CuTypedPool{Float32}() + @test tp.n_active === 0 + @test tp._checkpoint_n_active == [0] + @test tp._checkpoint_depths == [0] + tp.n_active = 2 + @test getfield(tp, :state).n_active === 2 + tp.n_active = Int32(1) # default-convert semantics preserved + @test tp.n_active === 1 + tp.n_active = 0 + @test :n_active in propertynames(tp) end @testset "CuAdaptiveArrayPool structure" begin @@ -89,6 +98,8 @@ end @test tp._checkpoint_n_active == [0, 0] @test tp._checkpoint_depths == [0, 2] @test pool._touched_has_others == [false, true] # depth-2 marked as "has dynamic types" + @test pool._touched_others_depths == [2] + @test length(pool._touched_others_states) == 1 end end diff --git a/test/cuda/test_touched_others.jl b/test/cuda/test_touched_others.jl new file mode 100644 index 00000000..a4cf8143 --- /dev/null +++ b/test/cuda/test_touched_others.jl @@ -0,0 +1,394 @@ +# CUDA port of test/test_touched_others.jl — depth-tagged touched-others stack +# and fallback lookup memo. Invariant: a fallback pool has an entry tagged with +# depth d in the depth-tagged stack (states/depths, kept in lockstep) ⟺ it was +# first touched at depth d — except under full checkpoint!(pool), whose eager +# sweep pairs with full rewind!(pool)'s sweep (stack stays empty, truncate-only). +# +# Fallback exercise types (per task-3 brief): UInt16, UInt8, Int8, UInt32, Int16, +# UInt64, Int128, UInt128 — CUDA would accept Int128/UInt128 as CuArray element +# types, but two isbits structs are substituted anyway (same as Metal) to keep +# the two GPU test files diffable against each other. +# Float16 is a fixed struct field with `_fixed_slot_bit == 0` on CUDA too (bit-7 +# reassignment, same as Metal), so it must NEVER be routed through the +# touched-others stack (see the dedicated "Float16 bit-7" testsets below). + +using AdaptiveArrayPools: _lazy_checkpoint!, _lazy_rewind!, + _typed_lazy_checkpoint!, _typed_lazy_rewind!, + _tracked_mask_for_types, _can_use_typed_path + +# Named distinctly from the CPU test suite's TOFooA/B/C (both files load into +# the same top-level module during a full-suite run) to avoid a +# struct-redefinition clash. +struct CudaFallbackStructA + x::Float32 +end +struct CudaFallbackStructB + x::Float32 +end + +@testset "touched-others: fields & lifecycle" begin + pool = CuAdaptiveArrayPool() + @test isempty(pool._touched_others_states) && isempty(pool._touched_others_depths) && isempty(pool._touched_others_pools) + + # reset! clears transient scope state, keeps registry + acquire!(pool, UInt16, 4) + reset!(pool) + @test length(pool._touched_others_states) == 0 + @test isempty(pool._touched_others_depths) + @test haskey(pool.others, UInt16) # registry kept + + # empty! clears everything + acquire!(pool, UInt16, 4) + empty!(pool) + @test length(pool._touched_others_states) == 0 + @test isempty(pool._touched_others_depths) + @test !haskey(pool.others, UInt16) +end + +@testset "touched-others: checkpoint/rewind plumbing balance" begin + pool = CuAdaptiveArrayPool() + + # lazy pair + _lazy_checkpoint!(pool) + @test isempty(pool._touched_others_depths) + _lazy_rewind!(pool) + @test isempty(pool._touched_others_depths) + + # typed single pair (fixed-slot type) + checkpoint!(pool, Float32) + @test isempty(pool._touched_others_depths) + rewind!(pool, Float32) + @test isempty(pool._touched_others_depths) + + # typed multi pair + checkpoint!(pool, Float32, Int32) + rewind!(pool, Float32, Int32) + @test isempty(pool._touched_others_depths) + + # full pair + checkpoint!(pool) + rewind!(pool) + @test isempty(pool._touched_others_depths) + + # typed-lazy pair + _typed_lazy_checkpoint!(pool, Float32) + _typed_lazy_rewind!(pool, _tracked_mask_for_types(Float32)) + @test isempty(pool._touched_others_depths) + + # nesting + _lazy_checkpoint!(pool) + checkpoint!(pool, Float32) + @test isempty(pool._touched_others_depths) + rewind!(pool, Float32) + _lazy_rewind!(pool) + @test isempty(pool._touched_others_depths) +end + +@testset "touched-others: no eager checkpoint on lazy entry (pollution regression)" begin + pool = CuAdaptiveArrayPool() + # Register all 8 fallback exercise types at global scope, then reset counters. + # (Int128/UInt128 substituted with distinct isbits structs — kept for + # diffability against the Metal test file, which must substitute them.) + fallback_types = (UInt16, UInt8, Int8, UInt32, Int16, UInt64, CudaFallbackStructA, CudaFallbackStructB) + for T in fallback_types + acquire!(pool, T, 4) + end + reset!(pool) + tps = Dict(T => get_typed_pool!(pool, T) for T in fallback_types) + + _lazy_checkpoint!(pool) + # THE regression assertion: unrelated registered fallbacks are NOT touched + for T in fallback_types[2:end] + @test tps[T]._checkpoint_depths == [0] + end + + acquire!(pool, UInt16, 8) + @test tps[UInt16]._checkpoint_depths[end] == 2 # first-touch checkpoint at depth 2 + @test length(pool._touched_others_depths) == 1 + @test pool._touched_others_states[end] === tps[UInt16].state + @test tps[UInt16].n_active == 1 + + # Re-acquire same type: no duplicate stack entry + acquire!(pool, UInt16, 8) + @test length(pool._touched_others_depths) == 1 + + _lazy_rewind!(pool) + @test tps[UInt16].n_active == 0 + @test isempty(pool._touched_others_depths) + for T in fallback_types[2:end] + @test tps[T]._checkpoint_depths == [0] # still never visited + end +end + +@testset "touched-others: nested scopes, same fallback type at two depths" begin + pool = CuAdaptiveArrayPool() + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, UInt16, 4) + @test tpA.n_active == 1 + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, UInt16, 4) + acquire!(pool, UInt16, 4) + @test tpA.n_active == 3 + @test length(pool._touched_others_depths) == 2 # one entry per depth + + _lazy_rewind!(pool) # exit depth 3 + @test tpA.n_active == 1 + + _lazy_rewind!(pool) # exit depth 2 + @test tpA.n_active == 0 +end + +@testset "touched-others: depth tags are exact and monotone" begin + pool = CuAdaptiveArrayPool() + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + @test isempty(pool._touched_others_depths) # entry pushes nothing + acquire!(pool, UInt16, 4) + @test pool._touched_others_depths == [2] + @test pool._touched_others_states[end] === tpA.state + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, UInt8, 4) # different fallback type, new depth + @test pool._touched_others_depths == [2, 3] + @test issorted(pool._touched_others_depths) # monotone invariant + + _lazy_rewind!(pool) # drains ONLY the ==3 entries + @test pool._touched_others_depths == [2] + @test tpA.n_active == 1 + @test get_typed_pool!(pool, UInt8).n_active == 0 + + _lazy_rewind!(pool) + @test isempty(pool._touched_others_depths) + @test tpA.n_active == 0 +end + +@testset "touched-others: nested scope NOT touching outer's fallback" begin + pool = CuAdaptiveArrayPool() + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, UInt16, 4) + _lazy_checkpoint!(pool) # depth 3: does not touch UInt16 + acquire!(pool, Float32, 16) # Float32 fixed-slot work only + @test length(pool._touched_others_depths) == 1 # no new fallback entry + _lazy_rewind!(pool) + @test tpA.n_active == 1 # outer's array untouched + _lazy_rewind!(pool) + @test tpA.n_active == 0 +end + +@testset "touched-others: typed scope with helper touching a fallback (typed-lazy)" begin + pool = CuAdaptiveArrayPool() + tpB = get_typed_pool!(pool, UInt8) + + _typed_lazy_checkpoint!(pool, Float32) + acquire!(pool, Float32, 8) # tracked fixed-slot work + acquire!(pool, UInt8, 4) # untracked helper-style fallback touch + @test pool._touched_has_others[end] == true + @test pool._touched_others_states[end] === tpB.state + @test !_can_use_typed_path(pool, _tracked_mask_for_types(Float32)) + _typed_lazy_rewind!(pool, _tracked_mask_for_types(Float32)) + @test tpB.n_active == 0 +end + +@testset "touched-others: tracked fallback type via typed checkpoint!" begin + pool = CuAdaptiveArrayPool() + + checkpoint!(pool, UInt16) # fallback T tracked by macro + tpA = get_typed_pool!(pool, UInt16) + @test pool._touched_others_states[end] === tpA.state # pushed at checkpoint + acquire!(pool, UInt16, 4) # public-API acquire: no double push + @test count(st -> st === tpA.state, pool._touched_others_states) == 1 + # macro exit path for has_others=true is _typed_lazy_rewind! + _typed_lazy_rewind!(pool, _tracked_mask_for_types(UInt16)) + @test tpA.n_active == 0 + @test isempty(pool._touched_others_depths) +end + +@testset "touched-others: new type registered mid-scope" begin + pool = CuAdaptiveArrayPool() + _lazy_checkpoint!(pool) + acquire!(pool, Int8, 4) # first-ever registration, in-scope + tpC = get_typed_pool!(pool, Int8) + @test pool._touched_others_states[end] === tpC.state + @test count(st -> st === tpC.state, pool._touched_others_states) == 1 + _lazy_rewind!(pool) + @test tpC.n_active == 0 +end + +@testset "touched-others: full checkpoint!/rewind! pairing unchanged" begin + pool = CuAdaptiveArrayPool() + acquire!(pool, UInt16, 4) + reset!(pool) + tpA = get_typed_pool!(pool, UInt16) + + checkpoint!(pool) # eager: checkpoints ALL others + @test tpA._checkpoint_depths[end] == 2 + acquire!(pool, UInt16, 4) + @test length(pool._touched_others_depths) == 0 # guard saw existing depth-2 entry + rewind!(pool) + @test tpA.n_active == 0 + @test isempty(pool._touched_others_depths) +end + +@testset "touched-others: similar! records fallback touch" begin + pool = CuAdaptiveArrayPool() + src = CuArray(UInt16[1, 2]) + _lazy_checkpoint!(pool) + similar!(pool, src) + tpA = get_typed_pool!(pool, UInt16) + @test pool._touched_others_states[end] === tpA.state + _lazy_rewind!(pool) + @test tpA.n_active == 0 +end + +@testset "touched-others: @with_pool integration + exception-leak recovery" begin + # The task-local CUDA pool is a process-wide-per-device singleton shared with + # every other test file. Start from a known-clean baseline so the absolute-depth + # assertions below are independent of test file execution order. + empty!(get_task_local_cuda_pool()) + + # Integration through the real macro (task-local pool) + f_leaf(n) = @with_pool :cuda p begin + q = acquire!(p, UInt16, n) + length(q) + end + @test f_leaf(8) == 8 + tl = get_task_local_cuda_pool() + @test get_typed_pool!(tl, UInt16).n_active == 0 + @test isempty(tl._touched_others_states) + + # Inner scope throws, outer catches: outer exit must clean up leaked state + function f_outer() + @with_pool :cuda p begin + acquire!(p, UInt8, 4) + try + @with_pool :cuda p2 begin + acquire!(p2, Int8, 4) + error("boom") + end + catch + end + 1 + end + end + @test f_outer() == 1 + @test tl._current_depth == 1 + @test get_typed_pool!(tl, UInt8).n_active == 0 + @test get_typed_pool!(tl, Int8).n_active == 0 + @test isempty(tl._touched_others_states) + @test isempty(tl._touched_others_depths) + empty!(tl) # leave the task-local pool clean for other test files +end + +@testset "fallback lookup memo: fields and lifecycle" begin + pool = CuAdaptiveArrayPool() + tp = get_typed_pool!(pool, UInt16) # slow-path lookup + @test pool._lookup_memo_type === UInt16 + @test pool._lookup_memo_tp === tp + + empty!(pool) + @test pool._lookup_memo_type === nothing + @test pool._lookup_memo_tp === nothing + + tp2 = get_typed_pool!(pool, UInt16) # re-register + @test pool._lookup_memo_type === UInt16 + @test pool._lookup_memo_tp === tp2 + + reset!(pool) # keeps registry AND memo + @test pool._lookup_memo_type === UInt16 + @test pool._lookup_memo_tp === tp2 +end + +@testset "fallback lookup memo: identity and invalidation" begin + pool = CuAdaptiveArrayPool() + tp1 = get_typed_pool!(pool, UInt16) + @test get_typed_pool!(pool, UInt16) === tp1 # repeat lookup: same pool + @test get_typed_pool!(pool, UInt8) !== tp1 # different type: different pool + @test get_typed_pool!(pool, UInt16) === tp1 # alternating types stay correct + + reset!(pool) # keeps registry → memo may stay + @test get_typed_pool!(pool, UInt16) === tp1 + + empty!(pool) # kills registry → memo MUST die + tp2 = get_typed_pool!(pool, UInt16) + @test tp2 !== tp1 # stale-memo regression guard + @test tp2 === pool.others[UInt16] + + # end-to-end: acquire after empty! must use the fresh pool + acquire!(pool, UInt16, 4) + @test tp2.n_active == 1 && tp1.n_active == 0 + reset!(pool) +end + +# ============================================================================== +# Float16 bit-7 non-interaction (GPU-only divergence — CPU has no such field) +# ============================================================================== +# Float16 is a FIXED struct field (pool.float16) but `_fixed_slot_bit(Float16) == +# 0` (Float16 is absent from the base module's generic fixed-slot bit table). +# Routing MUST key off `_cuda_is_fallback_type(T) = !(T <: _CUDA_FIXED_TYPES)`, +# never `_fixed_slot_bit(T) == 0` — Float16 goes through the direct +# `_checkpoint_typed_pool!`/`_rewind_typed_pool!` path like other fixed slots and +# must NEVER get a touched-others stack entry. + +@testset "Float16 bit-7: lazy scope acquiring Float16 pushes no stack entry" begin + pool = CuAdaptiveArrayPool() + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, Float16, 4) + @test isempty(pool._touched_others_depths) + @test pool.float16.n_active == 1 + _lazy_rewind!(pool) + @test pool.float16.n_active == 0 + @test isempty(pool._touched_others_depths) +end + +@testset "Float16 bit-7: typed inner scope alongside live parent Float16 arrays" begin + empty!(get_task_local_cuda_pool()) + pool = get_task_local_cuda_pool() + depth_before = pool._current_depth + + _lazy_checkpoint!(pool) # depth 2 (parent, lazy) + v_outer = acquire!(pool, Float16, 4) + v_outer .= Float16(3) + @test pool.float16.n_active == 1 + + # acquire!(pool2, Float16, 8) below is a static-type call, so the macro + # statically resolves Float16 and drives the TYPED path + # (_typed_lazy_checkpoint!/_typed_lazy_rewind!) rather than the untyped + # _lazy_checkpoint!/_lazy_rewind! path — no separate typed-macro syntax + # exists or is needed to opt in. + result = @with_pool :cuda pool2 begin + w = acquire!(pool2, Float16, 8) + @test isempty(pool._touched_others_depths) # Float16 never stack-managed, even in typed scope + w .= Float16(9) + sum(w) + end + @test result == Float16(9) * 8 + @test isempty(pool._touched_others_depths) # Float16 never stack-managed + @test pool.float16.n_active == 1 # parent's slot survives the inner exit + @test all(Array(v_outer) .== Float16(3)) # parent's array contents survive + + _lazy_rewind!(pool) + @test pool.float16.n_active == 0 + @test pool._current_depth == depth_before # depth balance: no leak across the testset + empty!(get_task_local_cuda_pool()) +end + +@testset "Float16 bit-7: mixed scope Float16 + fallback pushes exactly one entry" begin + pool = CuAdaptiveArrayPool() + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, Float16, 4) + acquire!(pool, UInt16, 4) + @test length(pool._touched_others_depths) == 1 + @test pool._touched_others_depths == [2] + tpA = get_typed_pool!(pool, UInt16) + @test pool._touched_others_states[end] === tpA.state + _lazy_rewind!(pool) + @test pool.float16.n_active == 0 + @test tpA.n_active == 0 + @test isempty(pool._touched_others_depths) +end diff --git a/test/test_state.jl b/test/test_state.jl index 8efd0a6b..f1622ea1 100644 --- a/test/test_state.jl +++ b/test/test_state.jl @@ -2417,9 +2417,12 @@ import AdaptiveArrayPools: _typed_lazy_checkpoint!, _typed_lazy_rewind!, _tracke end @testset "Issue #5: CUDA _typed_lazy_checkpoint! parity" begin - # Bug: CUDA version is missing two features present in CPU version: - # 1. Double-checkpoint guard: `_checkpoint_depths[end] != d` - # 2. has_others flag: `_touched_has_others[d] = true` + # Parity with the CPU touched-others-stack architecture (PR #51): + # _typed_lazy_checkpoint! must delegate to checkpoint!(pool, types...) + # (which routes fallback types through _touch_fallback_pool! — the + # double-checkpoint guard and has_others flag live THERE now) and set + # _TYPED_LAZY_BIT. It must NOT eagerly snapshot pool.others — that + # re-introduces the cross-scope pollution PR #51 removed. cuda_state_path = joinpath(@__DIR__, "..", "ext", "AdaptiveArrayPoolsCUDAExt", "state.jl") if isfile(cuda_state_path) code = read(cuda_state_path, String) @@ -2431,11 +2434,14 @@ import AdaptiveArrayPools: _typed_lazy_checkpoint!, _typed_lazy_rewind!, _tracke if func_match !== nothing func_body = func_match.match - # Must have double-checkpoint guard (like CPU version) - @test contains(func_body, "_checkpoint_depths[end]") + # Must delegate to checkpoint!(pool, types...) (like CPU version) + @test contains(func_body, "checkpoint!(pool, types...)") - # Must set _touched_has_others flag (like CPU version) - @test contains(func_body, "_touched_has_others") + # Must set the typed-lazy mode bit (like CPU version) + @test contains(func_body, "_TYPED_LAZY_BIT") + + # Must NOT eagerly snapshot pool.others (pollution regression) + @test !contains(func_body, "values(pool.others)") end else @warn "CUDA extension not found, skipping parity test" From 32035a938f113f99d0b6319b5087f069bc4c8ecd Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 10 Jul 2026 00:38:28 -0700 Subject: [PATCH 06/10] (test): cover S>=1 lazy-drain lockstep on GPU pools; retarget stale parity guard --- test/cuda/test_touched_others.jl | 32 ++++++++++++++++++++++++++++++- test/metal/test_touched_others.jl | 32 ++++++++++++++++++++++++++++++- test/test_state.jl | 19 +++++++++--------- 3 files changed, 72 insertions(+), 11 deletions(-) diff --git a/test/cuda/test_touched_others.jl b/test/cuda/test_touched_others.jl index a4cf8143..feb1cd3f 100644 --- a/test/cuda/test_touched_others.jl +++ b/test/cuda/test_touched_others.jl @@ -4,7 +4,7 @@ # first touched at depth d — except under full checkpoint!(pool), whose eager # sweep pairs with full rewind!(pool)'s sweep (stack stays empty, truncate-only). # -# Fallback exercise types (per task-3 brief): UInt16, UInt8, Int8, UInt32, Int16, +# Fallback exercise types: UInt16, UInt8, Int8, UInt32, Int16, # UInt64, Int128, UInt128 — CUDA would accept Int128/UInt128 as CuArray element # types, but two isbits structs are substituted anyway (same as Metal) to keep # the two GPU test files diffable against each other. @@ -285,6 +285,36 @@ end empty!(tl) # leave the task-local pool clean for other test files end +# ============================================================================== +# S=1: runtime-check-gated pools vector + invalidation +# ============================================================================== +# Every touched-others testset above drives CuAdaptiveArrayPool() at the +# default RUNTIME_CHECK level (S=0, no preference flip), so none of them exercise +# the `_runtime_check(pool)` branch in `_drain_touched_others!`/`_touch_fallback_pool!` +# that pushes/pops `_touched_others_pools` and invalidates a released fallback +# slot. Construct an S=1 pool directly to close that gap. + +@testset "touched-others: S=1 lazy-drain lockstep (pools vector populated, invalidation engaged)" begin + pool = ext._make_cuda_pool(1) + _lazy_checkpoint!(pool) # depth 2 + v = acquire!(pool, UInt16, 8) # genuine fallback acquire + tp = get_typed_pool!(pool, UInt16) + @test length(pool._touched_others_pools) == 1 # S=1-only: pools vector populated + @test pool._touched_others_depths == [2] + @test pool._touched_others_pools[end] === tp + + _lazy_rewind!(pool) + @test isempty(pool._touched_others_states) + @test isempty(pool._touched_others_depths) + @test isempty(pool._touched_others_pools) + @test tp.n_active == 0 + + # S=1 invalidation actually engaged: the previously-acquired wrapper's dims + # were zeroed (same assertion pattern as "arr_wrappers invalidated on rewind" + # in test_cuda_safety.jl). + @test all(==(0), size(v)) +end + @testset "fallback lookup memo: fields and lifecycle" begin pool = CuAdaptiveArrayPool() tp = get_typed_pool!(pool, UInt16) # slow-path lookup diff --git a/test/metal/test_touched_others.jl b/test/metal/test_touched_others.jl index 591de4b3..5d6a7a8a 100644 --- a/test/metal/test_touched_others.jl +++ b/test/metal/test_touched_others.jl @@ -4,7 +4,7 @@ # first touched at depth d — except under full checkpoint!(pool), whose eager # sweep pairs with full rewind!(pool)'s sweep (stack stays empty, truncate-only). # -# Fallback exercise types (per task-3 brief): UInt16, UInt8, Int8, UInt32, Int16, +# Fallback exercise types: UInt16, UInt8, Int8, UInt32, Int16, # UInt64, Int128, UInt128 — NOT Float64/ComplexF64 (rejected by the Metal backend). # Float16 is a Metal-only divergence: a FIXED struct field with `_fixed_slot_bit == # 0`, so it must NEVER be routed through the touched-others stack (see the @@ -285,6 +285,36 @@ end empty!(tl) # leave the task-local pool clean for other test files end +# ============================================================================== +# R=1: runtime-check-gated pools vector + invalidation +# ============================================================================== +# Every touched-others testset above drives MetalAdaptiveArrayPool() at the +# default RUNTIME_CHECK level (R=0, no preference flip), so none of them exercise +# the `_runtime_check(pool)` branch in `_drain_touched_others!`/`_touch_fallback_pool!` +# that pushes/pops `_touched_others_pools` and invalidates a released fallback +# slot. Construct an R=1 pool directly to close that gap. + +@testset "touched-others: R=1 lazy-drain lockstep (pools vector populated, invalidation engaged)" begin + pool = ext._make_metal_pool(1) + _lazy_checkpoint!(pool) # depth 2 + v = acquire!(pool, UInt16, 8) # genuine fallback acquire + tp = get_typed_pool!(pool, UInt16) + @test length(pool._touched_others_pools) == 1 # R=1-only: pools vector populated + @test pool._touched_others_depths == [2] + @test pool._touched_others_pools[end] === tp + + _lazy_rewind!(pool) + @test isempty(pool._touched_others_states) + @test isempty(pool._touched_others_depths) + @test isempty(pool._touched_others_pools) + @test tp.n_active == 0 + + # R=1 invalidation actually engaged: the previously-acquired wrapper's dims + # were zeroed (same assertion pattern as "arr_wrappers invalidated on rewind" + # in test_metal_safety.jl). + @test all(==(0), size(v)) +end + @testset "fallback lookup memo: fields and lifecycle" begin pool = MetalAdaptiveArrayPool() tp = get_typed_pool!(pool, UInt16) # slow-path lookup diff --git a/test/test_state.jl b/test/test_state.jl index f1622ea1..0cd06bd8 100644 --- a/test/test_state.jl +++ b/test/test_state.jl @@ -2390,10 +2390,14 @@ import AdaptiveArrayPools: _typed_lazy_checkpoint!, _typed_lazy_rewind!, _tracke end end - @testset "Issue #4: CUDA _lazy_checkpoint! parity (has_others flag)" begin - # Bug: CUDA _lazy_checkpoint! eagerly checkpoints pool.others but - # does NOT set _touched_has_others = true, same as CPU Issue #1. - # Verify via source code inspection (no GPU needed). + @testset "Issue #4: CUDA _lazy_checkpoint! parity" begin + # Parity with the CPU touched-others-stack architecture (PR #51): + # _lazy_checkpoint! must NOT eagerly checkpoint pool.others — fallback + # pools are first-touch checkpointed via _touch_fallback_pool! (from + # _record_type_touch!/get_typed_pool!) and drained selectively at rewind + # via _drain_touched_others!, so only the fallback pools a scope actually + # touches pay any cost. Eagerly snapshotting pool.others here would + # re-introduce the cross-scope pollution PR #51 removed. cuda_state_path = joinpath(@__DIR__, "..", "ext", "AdaptiveArrayPoolsCUDAExt", "state.jl") if isfile(cuda_state_path) code = read(cuda_state_path, String) @@ -2405,11 +2409,8 @@ import AdaptiveArrayPools: _typed_lazy_checkpoint!, _typed_lazy_rewind!, _tracke @test func_match !== nothing if func_match !== nothing func_body = func_match.match - # If it eagerly checkpoints others (has `for p in values(pool.others)`), - # then it MUST also set _touched_has_others[...] = true within the loop - if contains(func_body, "values(pool.others)") - @test occursin(r"_touched_has_others\[.*\]\s*=\s*true", func_body) - end + # Must NOT eagerly snapshot pool.others (pollution regression guard) + @test !contains(func_body, "values(pool.others)") end else @warn "CUDA extension not found, skipping parity test" From 28e186360dd1f2dd81fc03b6f66d7c22ef600b27 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 10 Jul 2026 10:41:26 -0700 Subject: [PATCH 07/10] (test): mixed-tuple and R=1 unwind coverage for touched-others; comment hygiene --- ext/AdaptiveArrayPoolsCUDAExt/types.jl | 15 ++-- ext/AdaptiveArrayPoolsMetalExt/types.jl | 15 ++-- test/cuda/test_touched_others.jl | 97 +++++++++++++++++++++++++ test/metal/test_touched_others.jl | 97 +++++++++++++++++++++++++ test/test_state.jl | 10 ++- 5 files changed, 218 insertions(+), 16 deletions(-) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/types.jl b/ext/AdaptiveArrayPoolsCUDAExt/types.jl index 71183a44..a8915350 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/types.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/types.jl @@ -64,7 +64,8 @@ function CuTypedPool{T}() where {T} ) end -# Checkpoint-state property forwarding (mirror of CPU src/types.jl:294-307). +# Checkpoint-state property forwarding (mirror of the CPU checkpoint-state property +# forwarding on TypedPool/BitTypedPool in src/types.jl). @inline function Base.getproperty(tp::CuTypedPool, f::Symbol) f === :n_active && return getfield(tp, :state).n_active f === :_checkpoint_n_active && return getfield(tp, :state)._checkpoint_n_active @@ -147,8 +148,9 @@ mutable struct CuAdaptiveArrayPool{S} <: AbstractArrayPool _touched_type_masks::Vector{UInt16} # Per-depth: which fixed slots were touched + mode flags _touched_has_others::Vector{Bool} # Per-depth: any non-fixed-slot type touched? - # Touched-others tracking (depth-tagged, concrete) — mirror of CPU - # src/types.jl:434-453. Checkpoint variants push NOTHING; producers push one + # Touched-others tracking (depth-tagged, concrete) — mirror of the CPU + # touched-others tracking fields on AdaptiveArrayPool in src/types.jl. + # Checkpoint variants push NOTHING; producers push one # (state, depth) entry per first touch; rewind pops while the top tag matches. # _touched_others_pools is populated only when S >= 1 (slot invalidation). _touched_others_states::Vector{PoolCheckpointState} @@ -227,9 +229,10 @@ Return compile-time constant indicating whether runtime safety checks are enable """ _check_level(pool::CuAdaptiveArrayPool) -> Int -Runtime-check level as an Int (mirror of CPU `src/types.jl:513`), for -backend-shared code that forwards it to `_invalidate_released_slots!` / -`_rewind_typed_pool!`. Compile-time constant per concrete pool type. +Runtime-check level as an Int (mirror of the CPU `_check_level` for +`AdaptiveArrayPool` in src/types.jl), for backend-shared code that forwards it to +`_invalidate_released_slots!` / `_rewind_typed_pool!`. Compile-time constant per +concrete pool type. """ @inline AdaptiveArrayPools._check_level(::CuAdaptiveArrayPool{S}) where {S} = S diff --git a/ext/AdaptiveArrayPoolsMetalExt/types.jl b/ext/AdaptiveArrayPoolsMetalExt/types.jl index c8f74df1..9be3b696 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/types.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/types.jl @@ -62,7 +62,8 @@ function MetalTypedPool{T, S}() where {T, S} ) end -# Checkpoint-state property forwarding (mirror of CPU src/types.jl:294-307). +# Checkpoint-state property forwarding (mirror of the CPU checkpoint-state property +# forwarding on TypedPool/BitTypedPool in src/types.jl). @inline function Base.getproperty(tp::MetalTypedPool, f::Symbol) f === :n_active && return getfield(tp, :state).n_active f === :_checkpoint_n_active && return getfield(tp, :state)._checkpoint_n_active @@ -143,8 +144,9 @@ mutable struct MetalAdaptiveArrayPool{R, S} <: AbstractArrayPool _touched_type_masks::Vector{UInt16} # Per-depth: which fixed slots were touched + mode flags _touched_has_others::Vector{Bool} # Per-depth: any non-fixed-slot type touched? - # Touched-others tracking (depth-tagged, concrete) — mirror of CPU - # src/types.jl:434-453. Checkpoint variants push NOTHING; producers push one + # Touched-others tracking (depth-tagged, concrete) — mirror of the CPU + # touched-others tracking fields on AdaptiveArrayPool in src/types.jl. + # Checkpoint variants push NOTHING; producers push one # (state, depth) entry per first touch; rewind pops while the top tag matches. # _touched_others_pools is populated only when R >= 1 (slot invalidation). _touched_others_states::Vector{PoolCheckpointState} @@ -220,9 +222,10 @@ Return compile-time constant indicating whether runtime safety checks are enable """ _check_level(pool::MetalAdaptiveArrayPool) -> Int -Runtime-check level as an Int (mirror of CPU `src/types.jl:513`), for -backend-shared code that forwards it to `_invalidate_released_slots!` / -`_rewind_typed_pool!`. Compile-time constant per concrete pool type. +Runtime-check level as an Int (mirror of the CPU `_check_level` for +`AdaptiveArrayPool` in src/types.jl), for backend-shared code that forwards it to +`_invalidate_released_slots!` / `_rewind_typed_pool!`. Compile-time constant per +concrete pool type. """ @inline AdaptiveArrayPools._check_level(::MetalAdaptiveArrayPool{R, S}) where {R, S} = R diff --git a/test/cuda/test_touched_others.jl b/test/cuda/test_touched_others.jl index feb1cd3f..597c1482 100644 --- a/test/cuda/test_touched_others.jl +++ b/test/cuda/test_touched_others.jl @@ -209,6 +209,23 @@ end @test isempty(pool._touched_others_depths) end +@testset "touched-others: mixed fixed+fallback tuple via typed checkpoint!" begin + pool = CuAdaptiveArrayPool() + + checkpoint!(pool, Float32, UInt16) # mixed tuple: one fixed slot, one fallback + acquire!(pool, Float32, 4) # fixed slot: checkpointed directly, no stack entry + acquire!(pool, UInt16, 4) # fallback: already pushed at checkpoint! time + tpU = get_typed_pool!(pool, UInt16) + @test length(pool._touched_others_depths) == 1 # exactly the UInt16 entry + @test pool._touched_others_states[end] === tpU.state + @test pool.float32._checkpoint_depths[end] == pool._current_depth # fixed slot checkpointed directly + + rewind!(pool, Float32, UInt16) + @test isempty(pool._touched_others_depths) + @test pool.float32.n_active == 0 + @test tpU.n_active == 0 +end + @testset "touched-others: new type registered mid-scope" begin pool = CuAdaptiveArrayPool() _lazy_checkpoint!(pool) @@ -315,6 +332,86 @@ end @test all(==(0), size(v)) end +@testset "touched-others: S=1 exception-leak recovery (lockstep across throw)" begin + # The "@with_pool integration + exception-leak recovery" testset above only + # exercises the two-stack (states/depths) lockstep at the default S=0. Construct + # an S=1 pool directly (same pattern as the S=1 lockstep testset above) to verify + # the third parallel stack (_touched_others_pools) stays in lockstep through a + # leaked-scope exception unwind too. + pool = ext._make_cuda_pool(1) + + _lazy_checkpoint!(pool) # outer scope, depth 2 + acquire!(pool, UInt8, 4) + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test pool._touched_others_depths == [2] + + # Inner scope checkpoints, touches a fallback type, then throws BEFORE its own + # rewind runs — mirrors @with_pool's direct-rewind fast path, which has no + # try/finally at its own scope (only an OUTER caller's entry-depth guard cleans + # up a leaked nested scope; see _generate_block_inner in src/macros.jl). + function _leaky_inner!(p) + _lazy_checkpoint!(p) # inner scope, depth 3 + acquire!(p, Int8, 4) + error("boom") + end + + try + _leaky_inner!(pool) + catch + end + + # Leaked: still at the inner depth, both fallback entries live, in lockstep + @test pool._current_depth == 3 + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 2 + @test pool._touched_others_depths == [2, 3] + + # Outer recovers exactly like @with_pool's entry-depth guard: the generic + # (eager) rewind! — not _lazy_rewind! — since the caller does not know which + # mode the leaked scope's checkpoint used. + rewind!(pool) # exit depth 3 (leaked inner) + @test pool._current_depth == 2 + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test pool._touched_others_depths == [2] + @test get_typed_pool!(pool, Int8).n_active == 0 + + _lazy_rewind!(pool) # outer's own scope exit + @test pool._current_depth == 1 + @test isempty(pool._touched_others_states) && isempty(pool._touched_others_depths) && isempty(pool._touched_others_pools) + @test get_typed_pool!(pool, UInt8).n_active == 0 +end + +@testset "touched-others: S=1 nested-depth lockstep (fallback touches at depths 2 and 3)" begin + # Mirrors "touched-others: depth tags are exact and monotone" above, but at S=1, + # with the three-stack lockstep asserted at every observation point. + pool = ext._make_cuda_pool(1) + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + @test isempty(pool._touched_others_depths) + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) + acquire!(pool, UInt16, 4) + @test pool._touched_others_depths == [2] + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test pool._touched_others_pools[end] === tpA + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, UInt8, 4) # different fallback type, new depth + @test pool._touched_others_depths == [2, 3] + @test issorted(pool._touched_others_depths) # monotone invariant + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 2 + + _lazy_rewind!(pool) # rewind one level: drains ONLY the ==3 entries + @test pool._touched_others_depths == [2] + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test tpA.n_active == 1 + @test get_typed_pool!(pool, UInt8).n_active == 0 + + _lazy_rewind!(pool) # rewind the remaining level + @test isempty(pool._touched_others_depths) + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) + @test tpA.n_active == 0 +end + @testset "fallback lookup memo: fields and lifecycle" begin pool = CuAdaptiveArrayPool() tp = get_typed_pool!(pool, UInt16) # slow-path lookup diff --git a/test/metal/test_touched_others.jl b/test/metal/test_touched_others.jl index 5d6a7a8a..6206a1d8 100644 --- a/test/metal/test_touched_others.jl +++ b/test/metal/test_touched_others.jl @@ -209,6 +209,23 @@ end @test isempty(pool._touched_others_depths) end +@testset "touched-others: mixed fixed+fallback tuple via typed checkpoint!" begin + pool = MetalAdaptiveArrayPool() + + checkpoint!(pool, Float32, UInt16) # mixed tuple: one fixed slot, one fallback + acquire!(pool, Float32, 4) # fixed slot: checkpointed directly, no stack entry + acquire!(pool, UInt16, 4) # fallback: already pushed at checkpoint! time + tpU = get_typed_pool!(pool, UInt16) + @test length(pool._touched_others_depths) == 1 # exactly the UInt16 entry + @test pool._touched_others_states[end] === tpU.state + @test pool.float32._checkpoint_depths[end] == pool._current_depth # fixed slot checkpointed directly + + rewind!(pool, Float32, UInt16) + @test isempty(pool._touched_others_depths) + @test pool.float32.n_active == 0 + @test tpU.n_active == 0 +end + @testset "touched-others: new type registered mid-scope" begin pool = MetalAdaptiveArrayPool() _lazy_checkpoint!(pool) @@ -315,6 +332,86 @@ end @test all(==(0), size(v)) end +@testset "touched-others: R=1 exception-leak recovery (lockstep across throw)" begin + # The "@with_pool integration + exception-leak recovery" testset above only + # exercises the two-stack (states/depths) lockstep at the default R=0. Construct + # an R=1 pool directly (same pattern as the R=1 lockstep testset above) to verify + # the third parallel stack (_touched_others_pools) stays in lockstep through a + # leaked-scope exception unwind too. + pool = ext._make_metal_pool(1) + + _lazy_checkpoint!(pool) # outer scope, depth 2 + acquire!(pool, UInt8, 4) + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test pool._touched_others_depths == [2] + + # Inner scope checkpoints, touches a fallback type, then throws BEFORE its own + # rewind runs — mirrors @with_pool's direct-rewind fast path, which has no + # try/finally at its own scope (only an OUTER caller's entry-depth guard cleans + # up a leaked nested scope; see _generate_block_inner in src/macros.jl). + function _leaky_inner!(p) + _lazy_checkpoint!(p) # inner scope, depth 3 + acquire!(p, Int8, 4) + error("boom") + end + + try + _leaky_inner!(pool) + catch + end + + # Leaked: still at the inner depth, both fallback entries live, in lockstep + @test pool._current_depth == 3 + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 2 + @test pool._touched_others_depths == [2, 3] + + # Outer recovers exactly like @with_pool's entry-depth guard: the generic + # (eager) rewind! — not _lazy_rewind! — since the caller does not know which + # mode the leaked scope's checkpoint used. + rewind!(pool) # exit depth 3 (leaked inner) + @test pool._current_depth == 2 + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test pool._touched_others_depths == [2] + @test get_typed_pool!(pool, Int8).n_active == 0 + + _lazy_rewind!(pool) # outer's own scope exit + @test pool._current_depth == 1 + @test isempty(pool._touched_others_states) && isempty(pool._touched_others_depths) && isempty(pool._touched_others_pools) + @test get_typed_pool!(pool, UInt8).n_active == 0 +end + +@testset "touched-others: R=1 nested-depth lockstep (fallback touches at depths 2 and 3)" begin + # Mirrors "touched-others: depth tags are exact and monotone" above, but at R=1, + # with the three-stack lockstep asserted at every observation point. + pool = ext._make_metal_pool(1) + tpA = get_typed_pool!(pool, UInt16) + + _lazy_checkpoint!(pool) # depth 2 + @test isempty(pool._touched_others_depths) + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) + acquire!(pool, UInt16, 4) + @test pool._touched_others_depths == [2] + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test pool._touched_others_pools[end] === tpA + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, UInt8, 4) # different fallback type, new depth + @test pool._touched_others_depths == [2, 3] + @test issorted(pool._touched_others_depths) # monotone invariant + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 2 + + _lazy_rewind!(pool) # rewind one level: drains ONLY the ==3 entries + @test pool._touched_others_depths == [2] + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) == 1 + @test tpA.n_active == 1 + @test get_typed_pool!(pool, UInt8).n_active == 0 + + _lazy_rewind!(pool) # rewind the remaining level + @test isempty(pool._touched_others_depths) + @test length(pool._touched_others_states) == length(pool._touched_others_depths) == length(pool._touched_others_pools) + @test tpA.n_active == 0 +end + @testset "fallback lookup memo: fields and lifecycle" begin pool = MetalAdaptiveArrayPool() tp = get_typed_pool!(pool, UInt16) # slow-path lookup diff --git a/test/test_state.jl b/test/test_state.jl index 0cd06bd8..ac17f346 100644 --- a/test/test_state.jl +++ b/test/test_state.jl @@ -2391,13 +2391,14 @@ import AdaptiveArrayPools: _typed_lazy_checkpoint!, _typed_lazy_rewind!, _tracke end @testset "Issue #4: CUDA _lazy_checkpoint! parity" begin - # Parity with the CPU touched-others-stack architecture (PR #51): + # Parity with the CPU touched-others-stack architecture (the constant-cost + # touched-others architecture, depth-tagged stack): # _lazy_checkpoint! must NOT eagerly checkpoint pool.others — fallback # pools are first-touch checkpointed via _touch_fallback_pool! (from # _record_type_touch!/get_typed_pool!) and drained selectively at rewind # via _drain_touched_others!, so only the fallback pools a scope actually # touches pay any cost. Eagerly snapshotting pool.others here would - # re-introduce the cross-scope pollution PR #51 removed. + # re-introduce the cross-scope pollution that architecture removed. cuda_state_path = joinpath(@__DIR__, "..", "ext", "AdaptiveArrayPoolsCUDAExt", "state.jl") if isfile(cuda_state_path) code = read(cuda_state_path, String) @@ -2418,12 +2419,13 @@ import AdaptiveArrayPools: _typed_lazy_checkpoint!, _typed_lazy_rewind!, _tracke end @testset "Issue #5: CUDA _typed_lazy_checkpoint! parity" begin - # Parity with the CPU touched-others-stack architecture (PR #51): + # Parity with the CPU touched-others-stack architecture (the constant-cost + # touched-others architecture, depth-tagged stack): # _typed_lazy_checkpoint! must delegate to checkpoint!(pool, types...) # (which routes fallback types through _touch_fallback_pool! — the # double-checkpoint guard and has_others flag live THERE now) and set # _TYPED_LAZY_BIT. It must NOT eagerly snapshot pool.others — that - # re-introduces the cross-scope pollution PR #51 removed. + # re-introduces the cross-scope pollution that architecture removed. cuda_state_path = joinpath(@__DIR__, "..", "ext", "AdaptiveArrayPoolsCUDAExt", "state.jl") if isfile(cuda_state_path) code = read(cuda_state_path, String) From a85cfc6beabbbf0c1398ea6597069fb4f1929169 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 10 Jul 2026 10:52:02 -0700 Subject: [PATCH 08/10] =?UTF-8?q?(fix):=20GPU=20poison-fill=20is=20best-ef?= =?UTF-8?q?fort=20=E2=80=94=20never=20throw=20during=20rewind?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the CPU _poison_fill! contract: custom isbits element types without zero(T) skip the poison pass instead of throwing MethodError mid-rewind. Surfaced by the touched-others struct fallback types under runtime_check=1. --- ext/AdaptiveArrayPoolsCUDAExt/debug.jl | 9 ++++++++- ext/AdaptiveArrayPoolsMetalExt/debug.jl | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/ext/AdaptiveArrayPoolsCUDAExt/debug.jl b/ext/AdaptiveArrayPoolsCUDAExt/debug.jl index 2d9961cf..39324dcc 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/debug.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/debug.jl @@ -36,7 +36,14 @@ Fill a CuVector with a detectable sentinel value (NaN for floats, typemax for in @noinline to avoid inlining GPU kernel launch overhead into hot rewind paths. """ @noinline function _cuda_poison_fill!(v::CuVector{T}) where {T} - length(v) > 0 && CUDA.fill!(v, _cuda_poison_value(T)) + length(v) > 0 || return nothing + # Mirror the CPU _poison_fill! contract: poisoning is best-effort and must + # not throw during rewind — custom isbits structs without zero(T) simply + # skip the poison pass (invalidation still shrinks the logical length). + try + CUDA.fill!(v, _cuda_poison_value(T)) + catch + end return nothing end diff --git a/ext/AdaptiveArrayPoolsMetalExt/debug.jl b/ext/AdaptiveArrayPoolsMetalExt/debug.jl index 6edecb72..a1ddc707 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/debug.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/debug.jl @@ -36,7 +36,14 @@ Fill a MtlArray with a detectable sentinel value (NaN for floats, typemax for in @noinline to avoid inlining GPU kernel launch overhead into hot rewind paths. """ @noinline function _metal_poison_fill!(v::MtlArray{T, 1}) where {T} - length(v) > 0 && Metal.fill!(v, _metal_poison_value(T)) + length(v) > 0 || return nothing + # Mirror the CPU _poison_fill! contract: poisoning is best-effort and must + # not throw during rewind — custom isbits structs without zero(T) simply + # skip the poison pass (invalidation still shrinks the logical length). + try + Metal.fill!(v, _metal_poison_value(T)) + catch + end return nothing end From b53867d2084abe5caaf6c38e083aa8f481527338 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 10 Jul 2026 10:52:02 -0700 Subject: [PATCH 09/10] (test): auto-manage end-to-end asserts are S-adaptive; scopes return nothing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At S=1, rewind's invalidation poisons released slots and compact! deliberately defers their shrink to preserve the poison (see _maybe_compact_slot!) — the capacity asserts now encode both semantics. Blocks also stop returning pool-backed arrays (PoolEscapeError under runtime_check=1). --- test/test_auto_manage.jl | 45 +++++++++++++++++++++++++++++----------- 1 file changed, 33 insertions(+), 12 deletions(-) diff --git a/test/test_auto_manage.jl b/test/test_auto_manage.jl index d182e266..73d78dbc 100644 --- a/test/test_auto_manage.jl +++ b/test/test_auto_manage.jl @@ -282,10 +282,10 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic empty!(pool) # clean slate AAP.register_auto_manage!(pool) @with_pool p begin # grow slot to 1M high-water - x = acquire!(p, Float64, 1_000_000); x .= 0.0 + x = acquire!(p, Float64, 1_000_000); x .= 0.0; nothing end @with_pool p begin # reuse small → bloated, inactive - x = acquire!(p, Float64, 100); x .= 0.0 + x = acquire!(p, Float64, 100); x .= 0.0; nothing end cap0 = _cap(pool.float64.vectors[1]) @test cap0 >= 1_000_000 @@ -296,10 +296,19 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic @atomic pool._compact_requested = true @with_pool p begin # scope ENTRY at depth 1 → hook fires acquire!(p, Float64, 4) + nothing end @test (@atomic pool._compact_requested) == false # hook consumed the request - @test _cap(pool.float64.vectors[1]) < cap0 # auto-managed, no manual compact! + if RUNTIME_CHECK == 0 + @test _cap(pool.float64.vectors[1]) < cap0 # auto-managed, no manual compact! + else + # S=1: rewind's invalidation poisoned the released slot (logical length 0); + # compact! deliberately skips poisoned slots so escaped views keep reading + # sentinels (see _maybe_compact_slot!). The hook still fired (flag consumed + # above); the capacity shrink is deferred until the slot is re-acquired. + @test _cap(pool.float64.vectors[1]) == cap0 + end empty!(pool) _clear_registry!() end @@ -313,10 +322,10 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic empty!(pool) # early RETURN @with_pool p begin - x = acquire!(p, Float64, 1_000_000); x .= 0.0 + x = acquire!(p, Float64, 1_000_000); x .= 0.0; nothing end @with_pool p begin - x = acquire!(p, Float64, 100); x .= 0.0 + x = acquire!(p, Float64, 100); x .= 0.0; nothing end cap0 = _cap(pool.float64.vectors[1]) @atomic pool._compact_requested = true @@ -326,14 +335,18 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic end @test ret() == 42 @test (@atomic pool._compact_requested) == false # serviced at the next scope ENTRY - @test _cap(pool.float64.vectors[1]) < cap0 + if RUNTIME_CHECK == 0 + @test _cap(pool.float64.vectors[1]) < cap0 + else + @test _cap(pool.float64.vectors[1]) == cap0 # S=1: poisoned slot, shrink deferred + end empty!(pool) # early BREAK @with_pool p begin - x = acquire!(p, Float64, 1_000_000); x .= 0.0 + x = acquire!(p, Float64, 1_000_000); x .= 0.0; nothing end @with_pool p begin - x = acquire!(p, Float64, 100); x .= 0.0 + x = acquire!(p, Float64, 100); x .= 0.0; nothing end cap1 = _cap(pool.float64.vectors[1]) @atomic pool._compact_requested = true @@ -344,7 +357,11 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic end end @test (@atomic pool._compact_requested) == false # serviced at the scope ENTRY - @test _cap(pool.float64.vectors[1]) < cap1 + if RUNTIME_CHECK == 0 + @test _cap(pool.float64.vectors[1]) < cap1 + else + @test _cap(pool.float64.vectors[1]) == cap1 # S=1: poisoned slot, shrink deferred + end empty!(pool) end @@ -353,10 +370,10 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic AAP.disable_auto_manage!() empty!(pool) @with_pool p begin - x = acquire!(p, Float64, 1_000_000); x .= 0.0 + x = acquire!(p, Float64, 1_000_000); x .= 0.0; nothing end @with_pool p begin - x = acquire!(p, Float64, 100); x .= 0.0 + x = acquire!(p, Float64, 100); x .= 0.0; nothing end cap0 = _cap(pool.float64.vectors[1]) @atomic pool._compact_requested = true @@ -372,7 +389,11 @@ AAP.disable_auto_manage!() # stop the __init__-started timer for deterministic end @test threw @test (@atomic pool._compact_requested) == false # hook fired at scope ENTRY (before the throw) - @test _cap(pool.float64.vectors[1]) < cap0 + if RUNTIME_CHECK == 0 + @test _cap(pool.float64.vectors[1]) < cap0 + else + @test _cap(pool.float64.vectors[1]) == cap0 # S=1: poisoned slot, shrink deferred + end empty!(pool) end end From 760d0fd8067b90d205a5fd6a38724a6492d55932 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 10 Jul 2026 11:06:36 -0700 Subject: [PATCH 10/10] (test): GPU allocation tests never return pool-backed arrays from scopes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same S=1 hygiene as the auto-manage tests: a @with_pool block whose last expression is a pool-backed array throws PoolEscapeError under runtime_check=1. All blocks now end with nothing; the @allocated / pointer comparisons are unaffected (they never used the scope return value). NOT run locally per user request — verify in the next full suite run. --- test/cuda/test_allocation.jl | 32 +++++++++++++++++++------------- test/metal/test_allocation.jl | 21 ++++++++++++++------- 2 files changed, 33 insertions(+), 20 deletions(-) diff --git a/test/cuda/test_allocation.jl b/test/cuda/test_allocation.jl index fa6ded44..e5bd84a0 100644 --- a/test/cuda/test_allocation.jl +++ b/test/cuda/test_allocation.jl @@ -10,14 +10,14 @@ # First acquire - populates pool @with_pool :cuda p begin v = acquire!(p, Float32, 100) - v .= 1.0f0 + v .= 1.0f0; nothing end # Second acquire (same size) - should reuse alloc = CUDA.@allocated begin @with_pool :cuda p begin v = acquire!(p, Float32, 100) - v .= 2.0f0 + v .= 2.0f0; nothing end end @@ -34,6 +34,7 @@ acquire!(p, Float32, 100) acquire!(p, Float32, 200) acquire!(p, Float32, 300) + nothing end # Second pass should reuse all @@ -42,7 +43,7 @@ v1 = acquire!(p, Float32, 100) v2 = acquire!(p, Float32, 200) v3 = acquire!(p, Float32, 300) - v1 .= 1.0f0; v2 .= 2.0f0; v3 .= 3.0f0 + v1 .= 1.0f0; v2 .= 2.0f0; v3 .= 3.0f0; nothing end end @@ -56,14 +57,14 @@ # Warmup with 2D array @with_pool :cuda p begin A = acquire!(p, Float64, 10, 10) - A .= 1.0 + A .= 1.0; nothing end # Reuse check alloc = CUDA.@allocated begin @with_pool :cuda p begin A = acquire!(p, Float64, 10, 10) - A .= 2.0 + A .= 2.0; nothing end end @@ -77,13 +78,13 @@ # Warmup with 3D array @with_pool :cuda p begin T = acquire!(p, Float32, 5, 5, 4) - T .= 1.0f0 + T .= 1.0f0; nothing end alloc = CUDA.@allocated begin @with_pool :cuda p begin T = acquire!(p, Float32, 5, 5, 4) - T .= 2.0f0 + T .= 2.0f0; nothing end end @@ -119,13 +120,13 @@ # Warmup @with_pool :cuda p begin A = acquire!(p, Float64, 10, 10) - A .= 1.0 + A .= 1.0; nothing end alloc = CUDA.@allocated begin @with_pool :cuda p begin A = acquire!(p, Float64, 10, 10) - A .= 2.0 + A .= 2.0; nothing end end @@ -141,6 +142,7 @@ # Warmup pool @with_pool :cuda p begin acquire!(p, Float32, N) + nothing end # Measure pooled allocation @@ -149,7 +151,7 @@ for _ in 1:ITERS @with_pool :cuda p begin v = acquire!(p, Float32, N) - v .= 1.0f0 + v .= 1.0f0; nothing end end end @@ -247,6 +249,7 @@ end function _test_cuda_nd_alloc!() @with_pool :cuda p begin acquire!(p, Float64, 10, 10) + nothing end end @@ -266,6 +269,7 @@ end function _test_cuda_nd_cached_alloc!() @with_pool :cuda p begin acquire!(p, Float64, 10, 10) + nothing end end @@ -285,6 +289,7 @@ end function _test_cuda_1d_alloc!() @with_pool :cuda p begin acquire!(p, Float64, 100) + nothing end end @@ -310,6 +315,7 @@ end acquire!(p, Float32, 100) acquire!(p, Float64, 100) acquire!(p, Int32, 100) + nothing end # Reuse all types @@ -318,7 +324,7 @@ end v32 = acquire!(p, Float32, 100) v64 = acquire!(p, Float64, 100) vi32 = acquire!(p, Int32, 100) - v32 .= 1.0f0; v64 .= 2.0; vi32 .= 3 + v32 .= 1.0f0; v64 .= 2.0; vi32 .= 3; nothing end end @@ -332,13 +338,13 @@ end # Warmup @with_pool :cuda p begin v = acquire!(p, Float16, 100) - v .= Float16(1.0) + v .= Float16(1.0); nothing end alloc = CUDA.@allocated begin @with_pool :cuda p begin v = acquire!(p, Float16, 100) - v .= Float16(2.0) + v .= Float16(2.0); nothing end end diff --git a/test/metal/test_allocation.jl b/test/metal/test_allocation.jl index 8db866ef..1b7805ff 100644 --- a/test/metal/test_allocation.jl +++ b/test/metal/test_allocation.jl @@ -10,14 +10,14 @@ # First acquire - populates pool @with_pool :metal p begin v = acquire!(p, Float32, 100) - v .= 1.0f0 + v .= 1.0f0; nothing end # Second acquire (same size) - should reuse GPU memory alloc = Metal.@allocated begin @with_pool :metal p begin v = acquire!(p, Float32, 100) - v .= 2.0f0 + v .= 2.0f0; nothing end end @@ -34,6 +34,7 @@ acquire!(p, Float32, 100) acquire!(p, Float32, 200) acquire!(p, Float32, 300) + nothing end # Second pass should reuse all GPU memory @@ -42,7 +43,7 @@ v1 = acquire!(p, Float32, 100) v2 = acquire!(p, Float32, 200) v3 = acquire!(p, Float32, 300) - v1 .= 1.0f0; v2 .= 2.0f0; v3 .= 3.0f0 + v1 .= 1.0f0; v2 .= 2.0f0; v3 .= 3.0f0; nothing end end @@ -57,14 +58,14 @@ # Warmup with 2D array @with_pool :metal p begin A = acquire!(p, Float32, 10, 10) - A .= 1.0f0 + A .= 1.0f0; nothing end # Reuse check — GPU allocation only alloc = Metal.@allocated begin @with_pool :metal p begin A = acquire!(p, Float32, 10, 10) - A .= 2.0f0 + A .= 2.0f0; nothing end end @@ -78,13 +79,13 @@ # Warmup with 3D array @with_pool :metal p begin T = acquire!(p, Float32, 5, 5, 4) - T .= 1.0f0 + T .= 1.0f0; nothing end alloc = Metal.@allocated begin @with_pool :metal p begin T = acquire!(p, Float32, 5, 5, 4) - T .= 2.0f0 + T .= 2.0f0; nothing end end @@ -180,6 +181,7 @@ end function _test_metal_nd_alloc!() @with_pool :metal p begin acquire!(p, Float32, 10, 10) + nothing end end @@ -198,6 +200,7 @@ end function _test_metal_1d_alloc!() @with_pool :metal p begin acquire!(p, Float32, 100) + nothing end end @@ -222,6 +225,7 @@ end acquire!(p, Float32, 100) acquire!(p, Int32, 100) acquire!(p, Float16, 100) + nothing end # Reuse all types — check GPU allocation only @@ -231,6 +235,7 @@ end vi32 = acquire!(p, Int32, 100) v16 = acquire!(p, Float16, 100) v32 .= 1.0f0; vi32 .= 3; v16 .= Float16(4.0) + nothing end end @@ -246,12 +251,14 @@ end @with_pool :metal p begin v = acquire!(p, Float16, 100) v .= Float16(1.0) + nothing end alloc = Metal.@allocated begin @with_pool :metal p begin v = acquire!(p, Float16, 100) v .= Float16(2.0) + nothing end end