(perf): GPU parity — depth-tagged touched-others stack, PoolCheckpointState, lookup memo (CUDA/Metal)

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")

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]

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
 

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

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,14 +167,19 @@ 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
             $(reset_exprs...)
             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)