From 7d18e6bab5b018366c0c30402e0b4c72b8aac8f5 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 13:26:13 -0700 Subject: [PATCH 01/10] Remove outdated design documents for selective memory reclamation, untracked acquire handling, and vector resize memory behavior - Deleted `selective_memory_reclamation.md` as the design has been implemented and is no longer needed. - Removed `untracked_acquire_design.md` which outlined a previous approach to handling untracked acquisitions in AdaptiveArrayPools. - Eliminated `vector_resize_memory_behavior.md` that discussed the memory behavior of vector resizing, as the design decisions have been finalized. --- .gitignore | 2 + docs/design/capacity_compaction.md | 338 ----- docs/design/cuda_extension_design.md | 1267 ----------------- docs/design/fixed_slots_codegen_design.md | 364 ----- docs/design/hybrid_api_design.md | 527 ------- .../macro-linenumbernode-improvement.md | 369 ----- docs/design/nd_array_approach_comparison.md | 440 ------ docs/design/selective_memory_reclamation.md | 417 ------ docs/design/untracked_acquire_design.md | 614 -------- docs/design/vector_resize_memory_behavior.md | 247 ---- 10 files changed, 2 insertions(+), 4583 deletions(-) delete mode 100644 docs/design/capacity_compaction.md delete mode 100644 docs/design/cuda_extension_design.md delete mode 100644 docs/design/fixed_slots_codegen_design.md delete mode 100644 docs/design/hybrid_api_design.md delete mode 100644 docs/design/macro-linenumbernode-improvement.md delete mode 100644 docs/design/nd_array_approach_comparison.md delete mode 100644 docs/design/selective_memory_reclamation.md delete mode 100644 docs/design/untracked_acquire_design.md delete mode 100644 docs/design/vector_resize_memory_behavior.md diff --git a/.gitignore b/.gitignore index 62091ad2..463e1e29 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +claudedocs + *.jl.cov *.jl.*.cov *.jl.mem diff --git a/docs/design/capacity_compaction.md b/docs/design/capacity_compaction.md deleted file mode 100644 index e19a6e15..00000000 --- a/docs/design/capacity_compaction.md +++ /dev/null @@ -1,338 +0,0 @@ -# Capacity Compaction for AdaptiveArrayPools (`compact!`) - -- **Date:** 2026-06-28 -- **Status:** Implemented (CPU + Metal + CUDA), Julia 1.12+. Shipped on top of the - `trim!` work. The Timer-driven auto mode (§9) remains future work. -- **Scope:** Manual `compact!` API that shrinks over-allocated backing buffers in - place. CPU + Metal + CUDA parity; an optional Timer-driven auto mode is deferred. -- **Author:** design session - ---- - -## 1. Motivation - -`trim!` (the sibling feature) drops the pool's references to **inactive** slots — -it reduces the *number of retained slots*. It deliberately never touches active -slots, so it cannot help when the bloated buffer is one the caller is still -holding. - -The remaining gap: a slot's backing buffer keeps the **high-water-mark capacity** -it ever needed. A rare large input grows a slot to (say) 1,000,000 elements; every -later `acquire!` on that slot only uses 100, but the buffer stays at 1,000,000 -because `resize!` down never frees capacity (verified: `resize!(v, 100)` leaves -`capacity == 1_000_000`). That over-allocation is invisible to `trim!` whenever -the slot is active. - -`compact!` fills that gap: - -> For buffers whose allocated capacity is far larger than the size currently in -> use, shrink the capacity back down to a small multiple of the used size — -> reclaiming the outlier bloat while keeping a right-sized buffer for reuse. - -This is a **cold-path** operation: called rarely (manually, or eventually on a -slow timer), not in the hot `acquire!`/`rewind!` loop. The design therefore -favors simplicity and safety over maximal reclamation. - -## 2. Relationship to `trim!` — two orthogonal axes - -| | `trim!` (exists) | `compact!` (this doc) | -| --- | --- | --- | -| Reduces | number of retained slots (structural) | over-allocated capacity of a slot's buffer | -| Touches | inactive slots only | retained buffers, **including active ones** | -| Safety | always safe (no live references) | must re-sync live wrappers (see §7) | -| Verb family | `reset!` / `empty!` / `trim!` | adds `compact!` | - -They are independent: `trim!` removes the tail; `compact!` compacts what remains. -Keeping them as **separate functions** (rather than a `trim!(; shrink_factor=…)` -kwarg) keeps each operation single-purpose and preserves `trim!`'s "always safe" -guarantee. The opt-in is layered: calling `compact!` opts into capacity compaction -of inactive slots (Tier 1, safe); `compact!(…; active=true)` further opts into -compacting active slots (Tier 2, §7). - -## 3. Core concept: capacity vs logical length - -For a backing vector `v = tp.vectors[slot]`: - -- **used** = the slot's **current logical extent**, recorded per-slot in - `tp.slot_extents[slot]` by `_claim_slot!` (`_slot_used`). This is the size of the - most recent claim and covers **both** acquisition paths — `acquire!` (Array - wrapper) *and* `acquire_view!` (which returns an **uncached** `SubArray`/ - `ReshapedArray`). **Not** `length(v)`: `_claim_slot!` only ever *grows* the backing - (`if length(vec) < n; resize!`, `acquire.jl:73`), so `length(v)` is the **high-water - mark**. An earlier design read `used` from the cached wrappers only; that - **under-counted a live view's extent** (views cache no wrapper) and let - `compact!(active=true)` shrink a backing below a held view → OOB. `slot_extents` - fixes this (and, being a `Vector{Int}`, makes `_slot_used` return a concrete `Int`, - keeping the gate arithmetic allocation-free). (Verified by TDD.) -- **capacity** = allocated buffer length (`length(getfield(v, :ref).mem)` on 1.12; - wrapped in `_slot_capacity`). - -`compact!` shrinks `capacity` toward `used`; it never changes `used` (held wrappers -keep their `:size`, and `target ≥ used` so every wrapper stays in bounds). - -A slot is a `compact!` candidate when **both**: - -``` -capacity ≥ factor × used # ratio gate: "≥10× bloated" -(capacity − target) × sizeof(T) ≥ min_bytes # absolute gate: real reclaim - where target = ceil(shrink_to × used) -``` - -The absolute byte gate prevents pointless reallocation of small buffers (e.g. a -1000-capacity / 50-used buffer is 20× bloated but only reclaims ~7.6 KB). - -## 4. API (as implemented) - -Mirror `trim!`'s surface for consistency: - -```julia -compact!(pool::AdaptiveArrayPool; - factor::Real = 10, # shrink only if capacity ≥ factor × used - shrink_to::Real = 1.5, # new capacity = ceil(shrink_to × used) - min_bytes::Int = 2^20, # …and only if ≥ this many bytes are reclaimed - active::Bool = true, # default: also compact ACTIVE slots (§7); false = inactive-only - force_gc::Bool = false) -compact!(pool::AdaptiveArrayPool, ::Type{T}; kwargs...) # single type -compact!(pool::AdaptiveArrayPool, types::Type...; kwargs...) # varargs (mirrors trim!) -compact!(; kwargs...) # task-local CPU pool -``` - -GPU extensions, `DisabledPool`, and the legacy (< 1.12) path mirror `trim!`'s -overload set exactly (legacy = warn-once no-op). Returns the same shape of -summary as `trim!`, plus a compaction-specific field: - -```julia -(; slots_compacted, bytes_reclaimed, gc_triggered) -``` - -(Reuse the `_trim_summary` pattern — a single `_compact_summary` helper — so every -overload returns one concrete, inferrable `NamedTuple`. See the `trim!` PR's -type-stability work for why this matters.) - -## 5. Shrink mechanism — in-place backing shrink (preserve object identity) - -**Decision (revised):** shrink the backing vector **in place** — keep the same -`Vector` object and swap only its internal `Memory` for a new, smaller one. Do -**not** replace the slot object (`tp.vectors[slot] = nv`). - -Use an **explicit `Memory` swap**, not `sizehint!`. `sizehint!` is documented as an -*advisory hint* — it is not guaranteed to shrink the buffer (the amount is -implementation/version-defined), so it cannot be the mechanism for a reclamation -API. The explicit swap is deterministic: - -```julia -target = ceil(Int, shrink_to * used) # target ≥ used -nv = Vector{T}(undef, target) # new small buffer (guaranteed size) -copyto!(nv, 1, v, 1, used) # preserve the live `used` elements -setfield!(v, :size, (used,)) # logical length stays `used` -setfield!(v, :ref, getfield(nv, :ref)) # v keeps identity; its Memory is now nv's -# then re-sync every cached wrapper for this slot (see §7) -``` - -This is the same `setfield!(:ref/:size)` the pool already uses for wrapper reuse, -so it is 1.12+-only (legacy < 1.12 → no-op, like `trim!`). - -Why keep the object and swap `:ref`, rather than "allocate a new vector and assign -the slot" (the earlier draft)? A smoke test -(`scratchpad/manual_refswap.jl`) confirmed the difference for live -`acquire_view!` views: - -- **Replace object** (`tp.vectors[slot] = nv`): a `SubArray`'s `parent` still points - at the *old* vector object → the view does **not** follow → broken. -- **Swap `:ref` in place**: the vector object keeps its identity and only its `:ref` - changes; a `SubArray` reads through `parent` on every access, so it **follows the - new buffer automatically** (verified: identity kept, `capacity 1M→5`, view reads - the new buffer, old `Memory` GC-collected). - -So **both** wrappers (re-synced, §7) **and** views (follow `parent`) survive. - -A bounds note: a view is only valid up to the slot's current `used` length, so -shrinking capacity toward `used` never invalidates a view that was valid at the -current size (a stale wider view was already out of bounds before compaction). - -**Does the old buffer actually get freed?** Yes — verified. The §5 `Memory` swap -re-points `v.ref` at a new smaller `Memory`; a `WeakRef` to the old `Memory` is dead -after `GC.gc()`, confirming it is reclaimed. One catch: -each `Array` **wrapper carries its own `:ref`** still pointing at the old `Memory`, -which pins it. So full reclamation requires **both** the in-place shrink **and** the -wrapper re-sync (§7.1) — skipping the re-sync leaves the wrapper stale *and* keeps -the old buffer alive. (`acquire_view!` views do not pin the old buffer: they read -through `parent`, which now points at the new one.) - -Note this is a different operation from `trim!`. `trim!` shrinks the *outer slot -array* (`resize!(tp.vectors, n_active)`), dropping whole inactive backing-`Vector` -objects — it reduces the slot **count** and returns those buffers 100%. `compact!` -keeps the slot and shrinks its buffer's **capacity** in place. They are -complementary (remove unused slots vs. right-size the buffers of slots still in -use), not two spellings of one thing. - -## 6. Modular layering (for the future auto-manager) - -The reusable unit is the per-slot primitive; everything else composes it: - -``` -_maybe_compact_slot!(tp, slot, factor, shrink_to, min_bytes) -> reclaimed::Int - │ one slot: gate check → (if candidate) in-place shrink + wrapper re-sync - │ returns bytes reclaimed (0 if skipped) - ▼ -compact!(pool; …) -> summary - │ scan all typed pools' retained slots, call the primitive, accumulate - ▼ -auto-manager (Phase C): - Timer fires → sets a "compaction requested" flag (does NOT touch buffers) - → the next pool touch point on the user's task runs the deferred compaction -``` - -The auto-manager writes **no new reclamation logic** — it only *schedules* a call -to `compact!`/`_maybe_compact_slot!`, which then runs synchronously on the user's -task at the next touch point (§7.3). The `factor`/`min_bytes` gates make it -naturally "leave the common case alone, only touch true outliers." - -That touch point may have `n_active > 0`, so the auto path **can** compact active -slots (Tier 2), not just inactive ones — safely, because the in-place swap + -re-sync runs single-threaded at a moment the user is inside a pool call rather than -mid-array-access. (This is the whole point of cooperative deferral: it lets the -timer reach the long-lived *active* outliers the user cares about, without the data -race a direct background mutation would cause.) - -## 7. Safety - -### 7.1 Two tiers — and why all the risk lives in Tier 2 - -The entire safety story splits cleanly by slot state: - -- **Tier 1 — inactive slots (default).** No live user reference exists. Shrinking - is unconditionally safe; the only bookkeeping is clearing the slot's cached - wrappers (they get rebuilt on the next `acquire!`). A plain `compact!(pool)` does - only this. -- **Tier 2 — active slots (opt-in, `active=true`).** The user is currently holding - the array from this slot, so the swap must re-sync the wrapper (§7.1) and is only - done at a point where the user is not mid-access — either an explicit mid-scope - manual call, or the auto path's cooperative deferral (§7.3). With the in-place - strategy (§5) both wrappers and views stay valid, so the gating concern here is - **concurrency, not references**. - -**Why Tier 2 is possible at all** (the non-obvious part): the array `acquire!` -returns is **not a copy** — it *is* the pool's own cached wrapper object, the same -one stored in `tp.arr_wrappers[N][slot]`. (Verified: `a === arr_wrappers[1][1]`.) -So the pool is not blind to "what the user holds" — it handed out its own objects -and still references every one of them. The wrapper is a thin handle that only -*points* at the backing buffer via a mutable `:ref`; the data lives in the buffer. -So after the in-place backing shrink (§5), the pool re-points the wrapper: - -```julia -# v = tp.vectors[slot]; after the §5 in-place Memory swap -for wrappers in tp.arr_wrappers # one per dimensionality N - w = wrappers === nothing ? continue : wrappers[slot] - w === nothing && continue - setfield!(w, :ref, getfield(v, :ref)) # re-point to v's new (smaller) Memory - # :size unchanged — logical dims are preserved -end -``` - -The wrapper carries its **own** `:ref` (a copy taken at creation), so it does not -follow `v`'s in-place change on its own and must be re-synced. The user holds the -*same* `w` object, so repointing its `:ref` redirects the user's array — no -reassignment in user code. This is the **same `setfield!(:ref)` the pool already -performs on every `acquire!`**, so the mechanism is load-bearing, not new. - -### 7.2 Views survive too — under the in-place strategy - -With the §5 in-place shrink, escaped `acquire_view!` results are **no longer a -hazard**. A `SubArray` / `ReshapedArray` holds its `parent` — the backing **vector -object** — and reads through it on every access. Because in-place shrink keeps that -object's identity and only swaps its `:ref`, the view transparently follows to the -new buffer (verified in the smoke test). Both reference kinds survive: - -| held reference | how it survives an in-place compaction | -| --- | --- | -| `acquire!` / `reshape!` `Array` wrapper | pool re-syncs its own `:ref` (§7.1) | -| `acquire_view!` `SubArray` / `ReshapedArray` | follows `parent`'s in-place `:ref` swap automatically | - -This is exactly why §5 chose in-place over replace-object: the *replace-object* -alternative (`tp.vectors[slot] = nv`) strands views on the old object (stale read + -pinned old buffer). Staying in-place is a **correctness requirement** for Tier 2, -not a stylistic preference. - -The only residual view caveat is bounds: never shrink capacity below the slot's -current `used` length — which the policy never does (`target ≥ used`). - -### 7.3 Concurrency — the real constraint for auto (Phase C) - -References are solved (§7.1–7.2). The genuine wall for **active** compaction is -concurrency: a `Timer` callback that reallocates a buffer while the user's task is -reading/writing it is a data race no reference trick can fix — lost writes, -use-after-free on the freed old `Memory`, a torn `:ref`. This is independent of -wrappers vs views. - -**The only safe auto design is cooperative deferral:** - -- The `Timer` does **not** touch buffers. It sets a "compaction requested" flag - (plus the policy). -- The actual shrink runs **synchronously on the user's task** at the next pool - touch point (`acquire!` / `checkpoint!` / `rewind!`, or an explicit - `maybe_compact!(pool)` the user can sprinkle into a long compute). At that instant - the user is inside a pool call, not mid-array-access, so the in-place swap + - re-sync is single-threaded-safe — and it can compact **active** slots, not just - inactive ones. -- Honest limitation: a buffer held across a long compute with **no** intervening - pool call cannot be compacted until the next touch point. Reclamation is delayed, - never unsafe. With `Threads.nthreads() > 1`, stricter synchronization is required - if pool calls can race across threads. - -So the user's goal — "a timer that reclaims oversized **active** outliers" — is -achievable: the timer **schedules** the compaction for the next safe touch point -rather than performing it directly. - -## 8. Backend parity - -- **CPU (`Vector{T}`):** the explicit `Memory` swap of §5. CPU's native `resize!` - *never* reallocates on shrink (capacity is retained forever), so the explicit - swap is the only way to actually return the memory. -- **CUDA / Metal:** backing is `CuVector` / `MtlArray`. Same shape — allocate a - smaller **device** buffer, copy `used` elements device-to-device, swap the - `:data`/DataRef + `:maxsize` + `:dims` (keeping the array object's identity, as on - CPU), then re-sync the GPU wrappers. Mind the GPU refcount rules the `trim!` work - established (Metal passes the `DataRef` directly; CUDA's constructor takes - ownership). Crucially, this must **not** route through the pool's - `_resize_to_fit!`, which deliberately *preserves* the GPU allocation - (`maxsize`) for reuse — `compact!` is exactly the op that overrides that and frees - it. Metal is hardware-testable on Apple Silicon; CUDA via CI. - - Note the symmetry: CUDA.jl's native `resize!` (v5.10+) already reallocates a - device buffer when it shrinks below `capacity ÷ 4` (`acquire.jl:18`). That is the - same idea as `compact!`'s `factor` gate — "realloc once `used` drops below a - fraction of capacity" — just with a fixed 25% threshold instead of our - configurable `factor` (e.g. `factor=10` ⇒ realloc below 10%). `compact!` brings - that capacity-aware behavior to CPU (which has none) and makes the threshold a - policy knob on both backends. -- **Legacy (< 1.12):** defined no-op that warns once and returns a zero summary, - exactly like legacy `trim!` (the `setfield!(:ref)` mechanism is 1.12+-only). - -## 9. Phasing - -1. **Phase A — CPU `compact!`** + the `_maybe_compact_slot!` primitive, full TDD: - capacity actually reclaimed, `used` data preserved, live wrapper still valid - after compaction, gates respected (ratio + min_bytes), no-op when nothing is - bloated, type-stable summary, varargs/per-type/DisabledPool/legacy parity. -2. **Phase B — Metal + CUDA parity** (Metal validated on hardware). -3. **Phase C — auto mode (deferred; out of scope for the first PR).** When taken - up: an external policy decides *when* and flips a "compaction requested" flag; - the actual compaction runs at a safe pool-scope boundary — concretely, the end - of a `@with_pool` block (the `rewind!`/scope-finalization the macro already - emits), not an arbitrary instant. That point is single-threaded-safe and may have - `n_active > 0`, so it can also reach active outliers (§7.3). Driving it from a - `Timer` is one possible policy source, but the **trigger is decoupled from the - execution point** — the timer only flags; the macro boundary executes. - -## 10. Open questions / decisions still to make - -- **Default `min_bytes`** — `2^20` (1 MiB) is a starting guess; tune against a - benchmark before settling. -- **Default `factor` / `shrink_to`** — `factor=10`, `shrink_to=1.5` are starting - guesses; CUDA.jl's native 25% (`÷4`) is a reference point for `factor`. -- **Whether `compact!` should also drop fully-inactive slots** (i.e. subsume - `trim!`) or stay strictly capacity-only and compose with `trim!`. Current lean: - stay capacity-only; users call `trim!(pool); compact!(pool)` for both axes. -- **Per-type policy** — whether `factor`/`shrink_to` should ever be settable - per element type, or stay global per call. Default: global per call (simpler). diff --git a/docs/design/cuda_extension_design.md b/docs/design/cuda_extension_design.md deleted file mode 100644 index f7eb7ee7..00000000 --- a/docs/design/cuda_extension_design.md +++ /dev/null @@ -1,1267 +0,0 @@ -# AdaptiveArrayPools.jl CUDA Extension Design - -> **Update (v0.2.2, feat/new_array_nd)**: The CPU path now uses `setfield!`-based wrapper -> reuse (Julia 1.11+) instead of the N-way cache for `unsafe_acquire!`. The **CUDA extension -> still uses the N-way set-associative cache** described in this document, since `CuArray` -> does not support `setfield!`-based field mutation. `CACHE_WAYS` and `set_cache_ways!` are -> now only relevant for the CUDA backend (and Julia 1.10 legacy CPU path). - -## 1. Executive Summary - -This document outlines the design for extending AdaptiveArrayPools.jl to support GPU arrays via CUDA.jl. The design prioritizes: - -1. **Zero overhead when CUDA not loaded** - Pure CPU code path unchanged -2. **Maximum code reuse** - Generic functions with minimal dispatch points -3. **Extensibility** - Abstract type hierarchy for future GPU backends -4. **Separate namespaces** - Independent CPU and GPU pools per task - -### Key Design Decision: Parametric Abstract Types - -Instead of duplicating code in the extension, we use parametric abstract types and generic functions. The extension only needs to define: -- Type definitions (~50 lines) -- One allocation method (~3 lines) -- Task-local getter (~15 lines) - -**Total extension code: ~70 lines** (vs ~300 lines with full duplication) - ---- - -## 2. Current Architecture Analysis - -### 2.1 Core Type Structure - -```julia -# Current: Concrete types only -mutable struct TypedPool{T} - vectors::Vector{Vector{T}} - views::Vector{SubArray{T,1,Vector{T},...}} - view_lengths::Vector{Int} - nd_arrays::Vector{Any} - nd_dims::Vector{Any} - nd_ptrs::Vector{UInt} - nd_next_way::Vector{Int} - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} -end -``` - -### 2.2 Code Reuse Analysis - -| Component | Operates On | GPU-Specific? | -|-----------|-------------|---------------| -| `get_view!` | vectors, n_active, cache | Only allocation | -| `get_nd_view!` | calls get_view!, reshape | **No** | -| `get_nd_array!` | calls get_view!, unsafe_wrap | Only wrap call | -| `checkpoint!` | n_active, checkpoint stacks | **No** | -| `rewind!` | n_active, checkpoint stacks | **No** | -| `reset!` | n_active, checkpoint stacks | **No** | -| `empty!` | all fields | Clear vectors only | - -**Key insight**: 95%+ of logic is type-agnostic. Only allocation/wrapping differs. - ---- - -## 3. Proposed Architecture - -### 3.1 Type Hierarchy - -``` - AbstractTypedPool{T, V<:AbstractVector{T}} - │ - ┌──────────────┼──────────────┐ - │ │ │ - TypedPool{T} CuTypedPool{T} [Future: ROCTypedPool{T}] - V = Vector{T} V = CuVector{T} V = ROCArray{T,1} - - - AbstractArrayPool - │ - ┌──────────────┼──────────────┐ - │ │ │ - AdaptiveArrayPool CuAdaptiveArrayPool [Future: ROCArrayPool] -``` - -### 3.2 Abstract Type Definitions - -```julia -# src/types.jl - New additions - -""" - AbstractTypedPool{T, V<:AbstractVector{T}} - -Abstract base for type-specific memory pools. -`T` is the element type, `V` is the storage vector type. - -Concrete subtypes must have these fields: -- vectors::Vector{V} -- views, view_lengths, nd_* fields -- n_active::Int -- _checkpoint_n_active::Vector{Int} -- _checkpoint_depths::Vector{Int} -""" -abstract type AbstractTypedPool{T, V<:AbstractVector{T}} end - -""" - AbstractArrayPool - -Abstract base for multi-type array pools. - -Concrete subtypes must have these fields: -- Fixed slot fields (type-specific TypedPools) -- others::IdDict{DataType,Any} -- _current_depth::Int -- _untracked_flags::Vector{Bool} -""" -abstract type AbstractArrayPool end - -# Storage type accessor (for generic code) -storage_type(::AbstractTypedPool{T,V}) where {T,V} = V -storage_type(::Type{<:AbstractTypedPool{T,V}}) where {T,V} = V -``` - -### 3.3 Concrete Types - -#### CPU (existing, modified to inherit) - -```julia -# src/types.jl - -mutable struct TypedPool{T} <: AbstractTypedPool{T, Vector{T}} - # Storage - vectors::Vector{Vector{T}} - views::Vector{SubArray{T,1,Vector{T},Tuple{UnitRange{Int64}},true}} - view_lengths::Vector{Int} - - # N-D cache - nd_arrays::Vector{Any} - nd_dims::Vector{Any} - nd_ptrs::Vector{UInt} - nd_next_way::Vector{Int} - - # State - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} -end - -mutable struct AdaptiveArrayPool <: AbstractArrayPool - # Fixed slots (CPU types) - float64::TypedPool{Float64} - float32::TypedPool{Float32} - int64::TypedPool{Int64} - int32::TypedPool{Int32} - complexf64::TypedPool{ComplexF64} - complexf32::TypedPool{ComplexF32} - bool::TypedPool{Bool} - - others::IdDict{DataType,Any} - _current_depth::Int - _untracked_flags::Vector{Bool} -end -``` - -#### GPU (extension - minimal definitions) - -> **[AI Review: Float16 & Device Safety]** -> 1. **Float16 Support**: Added `Float16` to fixed slots. This is critical for modern AI/ML workloads on GPU. -> 2. **Device Awareness**: Added `device_id::Int` to `CuAdaptiveArrayPool`. This is crucial for multi-GPU setups. A pool created on Device 0 cannot be safely used on Device 1. We must track which device owns the memory. - -> **[Post-Review v0.6: Critical Type Correction]** -> **`view(CuVector, 1:n)` returns `CuVector`, NOT `SubArray`!** -> -> GPUArrays.jl handles contiguous views via `derive()` which returns a new GPU array -> sharing the same memory buffer (see `~/.julia/packages/GPUArrays/.../src/host/base.jl:302`). -> This is fundamentally different from CPU where `view()` returns `SubArray`. -> -> **Implications for pool design**: -> 1. We **cannot cache views separately** from backing vectors on GPU -> 2. Instead, we store `CuVector{T}` directly and return slices via `view()` on each call -> 3. View creation is cheap (no allocation, just metadata), so no caching benefit -> 4. This simplifies the GPU pool: no `views` or `view_lengths` fields needed - -```julia -# ext/AdaptiveArrayPoolsCUDAExt/types.jl - -using CUDA - -# IMPORTANT: Unlike CPU, GPU views are derived CuArrays, not SubArrays. -# view(::CuVector{T}, ::UnitRange) -> CuVector{T} (shared memory, different offset/length) -# This means: -# 1. "views" vector would just hold more CuVectors (no savings) -# 2. We skip view caching entirely - just return view(vec, 1:n) each time -# 3. View creation is O(1) metadata operation, no GPU memory allocation - -mutable struct CuTypedPool{T} <: AbstractTypedPool{T, CuVector{T}} - # Storage (GPU vectors) - vectors::Vector{CuVector{T}} - - # View length cache (for resize decision, but no view object cache) - # The actual view is created fresh each time since it's just metadata - view_lengths::Vector{Int} - - # N-D cache (same structure as CPU) - nd_arrays::Vector{Any} - nd_dims::Vector{Any} - nd_ptrs::Vector{UInt} - nd_next_way::Vector{Int} - - # State (identical to CPU) - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} -end - -# Constructor with sentinel pattern -function CuTypedPool{T}() where T - CuTypedPool{T}( - CuVector{T}[], Int[], # No views vector! - Any[], Any[], UInt[], Int[], - 0, [0], [0] - ) -end - -# GPU-optimized fixed slots (different from CPU!) -const GPU_FIXED_SLOT_FIELDS = ( - :float32, # Primary (GPU-optimized) - :float64, # Precision when needed - :float16, # ML inference (added per AI review) - :int32, # Indexing (GPU-preferred) - :int64, # Large indices - :complexf32, # FFT, signal processing - :complexf64, # High-precision complex - :bool, # Masks -) - -mutable struct CuAdaptiveArrayPool <: AbstractArrayPool - # Fixed slots (GPU-optimized order: Float32 first) - float32::CuTypedPool{Float32} - float64::CuTypedPool{Float64} - float16::CuTypedPool{Float16} # Added per AI review - int32::CuTypedPool{Int32} - int64::CuTypedPool{Int64} - complexf32::CuTypedPool{ComplexF32} - complexf64::CuTypedPool{ComplexF64} - bool::CuTypedPool{Bool} - - others::IdDict{DataType,Any} - _current_depth::Int - _untracked_flags::Vector{Bool} - - # Safety: Track which device this pool belongs to (use public API!) - device_id::Int -end - -function CuAdaptiveArrayPool() - dev = CUDA.device() - CuAdaptiveArrayPool( - CuTypedPool{Float32}(), CuTypedPool{Float64}(), CuTypedPool{Float16}(), - CuTypedPool{Int32}(), CuTypedPool{Int64}(), - CuTypedPool{ComplexF32}(), CuTypedPool{ComplexF64}(), - CuTypedPool{Bool}(), - IdDict{DataType,Any}(), 1, [false], - CUDA.deviceid(dev) # Use public API, not internal .handle - ) -end -``` - ---- - -## 4. Generic Functions with Minimal Dispatch - -### 4.1 Allocation Dispatch Point - -The **only** type-specific function needed: - -```julia -# src/acquire.jl - CPU default -""" - allocate_vector(tp::AbstractTypedPool{T}, n::Int) -> V - -Allocate a new vector of type V with n elements. -This is the single dispatch point for storage-specific allocation. -""" -@inline allocate_vector(::AbstractTypedPool{T,Vector{T}}, n::Int) where T = - Vector{T}(undef, n) - -# ext/ - GPU override (THE ONLY METHOD EXTENSION NEEDS TO ADD!) -@inline allocate_vector(::AbstractTypedPool{T,CuVector{T}}, n::Int) where T = - CuVector{T}(undef, n) -``` - -> **[AI Review: Interaction with CUDA.jl Allocator]** -> It is important to note that `CuVector{T}(undef, n)` uses `CUDA.jl`'s own internal memory pool. -> **Why do we need another pool?** -> 1. **Overhead Reduction**: Even cached CUDA allocations have Julia-side overhead (struct creation, finalizer registration). `AdaptiveArrayPools` reuses the *Julia objects* (`CuArray` structs) and views, reducing GC pressure and allocation latency further. -> 2. **Logical Grouping**: It allows "rewinding" a whole block of temporary allocations in one go, which `CUDA.jl`'s allocator doesn't support (it's `malloc`/`free` style). - -### 4.2 get_view! Implementation - -> **[Post-Review v0.6: CPU vs GPU Differences]** -> Due to type differences (`view(Vector, 1:n) → SubArray` vs `view(CuVector, 1:n) → CuVector`), -> the CPU and GPU implementations differ slightly. CPU caches view objects; GPU creates them fresh. - -#### CPU Version (existing, unchanged) - -```julia -# src/acquire.jl - CPU implementation (caches SubArray views) - -function get_view!(tp::AbstractTypedPool{T,Vector{T}}, n::Int) where {T} - tp.n_active += 1 - idx = tp.n_active - - # 1. Expand pool if needed - if idx > length(tp.vectors) - push!(tp.vectors, allocate_vector(tp, n)) - new_view = view(tp.vectors[idx], 1:n) - push!(tp.views, new_view) # Cache the SubArray - push!(tp.view_lengths, n) - # ... growth warning ... - return new_view - end - - # 2. Cache hit (return cached SubArray - ZERO ALLOC) - @inbounds cached_len = tp.view_lengths[idx] - if cached_len == n - return @inbounds tp.views[idx] - end - - # 3. Cache miss - resize and update cached view - @inbounds vec = tp.vectors[idx] - if length(vec) < n - resize!(vec, n) - end - new_view = view(vec, 1:n) - @inbounds tp.views[idx] = new_view - @inbounds tp.view_lengths[idx] = n - return new_view -end -``` - -#### GPU Version (extension) - -> **[Post-Review v0.6: resize! Cost Warning]** -> `resize!(::CuVector, n)` with capacity increase triggers: -> 1. New GPU buffer allocation -> 2. Async copy of existing elements (even if we don't need them!) -> -> For pools, we typically don't need old data. Consider using `CUDA.unsafe_free!` + fresh -> allocation instead, or just allocating oversized initially. This is a **performance -> optimization opportunity** for v1.1+. - -```julia -# ext/AdaptiveArrayPoolsCUDAExt/acquire.jl - -# GPU version: no view caching (view() returns CuVector, not SubArray) -function AdaptiveArrayPools.get_view!(tp::CuTypedPool{T}, n::Int) where {T} - tp.n_active += 1 - idx = tp.n_active - - # 1. Expand pool if needed - if idx > length(tp.vectors) - push!(tp.vectors, allocate_vector(tp, n)) - push!(tp.view_lengths, n) - # Return fresh view (no caching - view creates CuVector metadata) - return view(tp.vectors[idx], 1:n) - end - - # 2. Check if resize needed - @inbounds cached_len = tp.view_lengths[idx] - @inbounds vec = tp.vectors[idx] - - if length(vec) < n - # WARNING: resize! on CuVector copies old data (wasteful for pools) - # TODO v1.1: Consider CUDA.unsafe_free! + fresh alloc instead - resize!(vec, n) - end - - @inbounds tp.view_lengths[idx] = n - - # Always create fresh view (O(1) metadata, no GPU allocation) - return view(vec, 1:n) -end -``` - -### 4.3 get_nd_view! Implementation - -> **[Post-Review v0.6: reshape Behavior on GPU]** -> `reshape(::CuVector, dims)` also uses GPUArrays' `derive()` mechanism, returning a -> `CuArray{T,N}` (not `ReshapedArray`). This is actually simpler - we get a proper -> GPU array that CUDA kernels can use directly. - -```julia -# src/acquire.jl - Works for both, but return types differ: -# - CPU: ReshapedArray{T,N,SubArray{...}} -# - GPU: CuArray{T,N} (via derive) - -@inline function get_nd_view!(tp::AbstractTypedPool{T}, dims::NTuple{N,Int}) where {T,N} - total_len = safe_prod(dims) - flat_view = get_view!(tp, total_len) - return reshape(flat_view, dims) # CPU: ReshapedArray, GPU: CuArray -end -``` - -### 4.4 Generic get_nd_array! (minimal dispatch) - -```julia -# src/acquire.jl - -# CPU version uses unsafe_wrap -@inline function wrap_array(::AbstractTypedPool{T,Vector{T}}, - flat_view, dims::NTuple{N,Int}) where {T,N} - unsafe_wrap(Array{T,N}, pointer(flat_view), dims) -end - -# ext/ - GPU version -@inline function wrap_array(::AbstractTypedPool{T,CuVector{T}}, - flat_view, dims::NTuple{N,Int}) where {T,N} - # Use reshape - returns CuArray{T,N} via GPUArrays derive() - reshape(flat_view, dims) -end - -# Generic implementation -@inline function get_nd_array!(tp::AbstractTypedPool{T}, dims::NTuple{N,Int}) where {T,N} - total_len = safe_prod(dims) - flat_view = get_view!(tp, total_len) - slot = tp.n_active - - # ... cache lookup logic (identical) ... - - # DISPATCH POINT for array wrapping - arr = wrap_array(tp, flat_view, dims) - - # ... cache update logic (identical) ... - - return arr -end -``` - -> **[Post-Review v0.6: GPU reshape Clarification]** -> `reshape(::CuArray, dims)` returns a `CuArray{T,N}` (via GPUArrays `derive()`), **NOT** -> `ReshapedArray`. This is actually better for GPU kernels - they work directly with -> `CuArray` without any wrapper overhead. The `derive()` mechanism shares the underlying -> GPU memory buffer with different offset/strides metadata. - ---- - -## 5. State Management (100% Reusable) - -### 5.1 Generic State Functions - -All state functions operate only on `n_active` and checkpoint vectors - pure CPU operations. - -```julia -# src/state.jl - These work for ANY AbstractTypedPool! - -@inline function _checkpoint_typed_pool!(tp::AbstractTypedPool, depth::Int) - push!(tp._checkpoint_n_active, tp.n_active) - push!(tp._checkpoint_depths, depth) - nothing -end - -@inline function _rewind_typed_pool!(tp::AbstractTypedPool, current_depth::Int) - # Orphan cleanup - while @inbounds tp._checkpoint_depths[end] > current_depth - pop!(tp._checkpoint_depths) - pop!(tp._checkpoint_n_active) - end - - # Restore - if @inbounds tp._checkpoint_depths[end] == current_depth - pop!(tp._checkpoint_depths) - tp.n_active = pop!(tp._checkpoint_n_active) - else - tp.n_active = @inbounds tp._checkpoint_n_active[end] - end - nothing -end - -function _reset_typed_pool!(tp::AbstractTypedPool) - tp.n_active = 0 - empty!(tp._checkpoint_n_active) - push!(tp._checkpoint_n_active, 0) - empty!(tp._checkpoint_depths) - push!(tp._checkpoint_depths, 0) - tp -end - -# Concrete dispatches (trivial wrappers) -reset!(tp::TypedPool) = _reset_typed_pool!(tp) -reset!(tp::CuTypedPool) = _reset_typed_pool!(tp) # ext/ adds this -``` - -### 5.2 empty! (Type-Specific) - -`empty!` needs to clear storage, but the logic is identical: - -```julia -# src/state.jl - Generic implementation - -function Base.empty!(tp::AbstractTypedPool) - empty!(tp.vectors) - empty!(tp.views) # CPU only (GPU CuTypedPool has no views field) - empty!(tp.view_lengths) - empty!(tp.nd_arrays) - empty!(tp.nd_dims) - empty!(tp.nd_ptrs) - empty!(tp.nd_next_way) - _reset_typed_pool!(tp) - tp -end - -# GPU-specific version (no views field) -function Base.empty!(tp::CuTypedPool) - empty!(tp.vectors) - empty!(tp.view_lengths) - empty!(tp.nd_arrays) - empty!(tp.nd_dims) - empty!(tp.nd_ptrs) - empty!(tp.nd_next_way) - _reset_typed_pool!(tp) - tp -end -``` - -> **[Post-Review v0.6: GPU Memory Release Clarification]** -> `empty!(tp.vectors)` **removes Julia references** to `CuVector` objects. This does NOT -> guarantee immediate VRAM release! The actual GPU memory lifecycle is: -> -> 1. **Reference removed** → CuArray becomes GC-eligible -> 2. **GC runs** → CuArray finalizer queued -> 3. **Finalizer runs** → Returns memory to CUDA.jl's internal pool -> 4. **CUDA.jl pool decision** → May or may not release to driver -> -> For **immediate VRAM release**, use `CUDA.reclaim()` after `empty!()`: -> ```julia -> empty!(get_task_local_cuda_pool()) -> GC.gc() # Force finalizers to run -> CUDA.reclaim() # Request CUDA.jl to release cached memory -> ``` - ---- - -## 6. Task-Local Pool Design - -> **[AI Review: Multi-Device Safety]** -> The original design for `get_task_local_cuda_pool` was unsafe for multi-GPU workflows. If a task switches devices (e.g., `CUDA.device!(1)`), it must not use the pool created for Device 0. -> **Revised Design**: We use a `Dict{Int, CuAdaptiveArrayPool}` in task local storage to manage one pool per device per task. - -### 6.1 Separate Keys & Device Awareness - -```julia -# src/task_local_pool.jl -const _POOL_KEY = :ADAPTIVE_ARRAY_POOL - -@inline function get_task_local_pool() - pool = get(task_local_storage(), _POOL_KEY, nothing) - if pool === nothing - pool = AdaptiveArrayPool() - task_local_storage(_POOL_KEY, pool) - end - return pool::AdaptiveArrayPool -end - -# ext/AdaptiveArrayPoolsCUDAExt/task_local_pool.jl -const _CU_POOL_KEY = :ADAPTIVE_ARRAY_POOL_CUDA - -@inline function get_task_local_cuda_pool() - # Get the dictionary of pools (one per device) - pools = get(task_local_storage(), _CU_POOL_KEY, nothing) - if pools === nothing - pools = Dict{Int, CuAdaptiveArrayPool}() - task_local_storage(_CU_POOL_KEY, pools) - end - - # Get current device ID using public API - dev_id = CUDA.deviceid(CUDA.device()) - - # Get or create pool for this device - if !haskey(pools, dev_id) - pools[dev_id] = CuAdaptiveArrayPool() # Constructor captures device_id - end - - return pools[dev_id] -end -``` - -> **[Post-Review v0.6: Public API for Device ID]** -> Always use `CUDA.deviceid(dev)` instead of `dev.handle`. The `.handle` field is internal -> and may change between CUDA.jl versions. `deviceid()` is the stable public API. - -### 6.2 Rationale for Separation - -| Scenario | Benefit | -|----------|---------| -| Mixed CPU/GPU workflow | Use both pools independently | -| GPU memory pressure | `empty!(cuda_pool)` without affecting CPU | -| Different lifecycles | CPU warm, GPU cleared per batch | -| **Multi-GPU** | **Safety**: Prevents cross-device access errors | -| Debugging | Clear distinction in profiling | - ---- - -## 7. Macro Design - -### 7.1 Recommended: Unified Macro with Backend Symbol - -```julia -# Unified API - single macro with optional backend symbol -@with_pool pool begin ... end # CPU (default, :cpu implied) -@with_pool :cuda pool begin ... end # GPU via CUDA -@with_pool :metal pool begin ... end # GPU via Metal (future) -@with_pool :cpu pool begin ... end # Explicit CPU - -# Without pool name (auto-generated) -@with_pool begin ... end # CPU default -@with_pool :cuda begin ... end # GPU -``` - -**Advantages:** -- Single macro to learn -- Easy backend switching (`:cuda` → `:metal`) -- Future-proof (just add new symbols in extensions) -- Clean, consistent API - -### 7.2 Implementation - -> **[Post-Review v0.6: Zero-Overhead Backend Selection]** -> The original `Dict{Symbol, Function}` registry has a critical flaw: runtime dictionary -> lookup weakens type inference, preventing the compiler from inlining the pool getter. -> This conflicts with our "zero overhead for CPU path" goal. -> -> **Solution**: Use `Val{:backend}` dispatch instead. Extensions add methods at load time, -> and the compiler can fully inline the call chain. - -```julia -# src/macros.jl - Val-based dispatch for zero overhead - -""" - _get_pool_for_backend(::Val{:cpu}) -> AdaptiveArrayPool - -Get task-local pool for the specified backend. Extensions add methods for their backends. -Using Val{Symbol} enables compile-time dispatch and full inlining. -""" -@inline _get_pool_for_backend(::Val{:cpu}) = get_task_local_pool() - -# Fallback with helpful error message -@noinline function _get_pool_for_backend(::Val{B}) where B - error("Pool backend :$B not found. Did you forget to load the extension (e.g., `using CUDA`)?") -end - -# Macro signatures -macro with_pool(backend::QuoteNode, pool_name, expr) - _generate_pool_code_with_backend(backend.value, pool_name, expr) -end - -macro with_pool(backend::QuoteNode, expr) - # Backend symbol without pool name - pool_name = gensym(:pool) - _generate_pool_code_with_backend(backend.value, pool_name, expr) -end - -macro with_pool(pool_name, expr) - # No backend = CPU default - _generate_pool_code_with_backend(:cpu, pool_name, expr) -end - -macro with_pool(expr) - # No backend, no pool name - pool_name = gensym(:pool) - _generate_pool_code_with_backend(:cpu, pool_name, expr) -end - -function _generate_pool_code_with_backend(backend::Symbol, pool_name, expr) - transformed_expr = _transform_acquire_calls(expr, pool_name) - - # Use Val{backend} for compile-time dispatch - fully inlinable! - quote - local $(esc(pool_name)) = $_get_pool_for_backend($(Val{backend}())) - checkpoint!($(esc(pool_name))) - try - $(esc(transformed_expr)) - finally - rewind!($(esc(pool_name))) - end - end -end -``` - -> **Why Val{:backend} instead of Dict?** -> -> | Approach | Lookup Cost | Type Inference | Inlining | -> |----------|-------------|----------------|----------| -> | `Dict{Symbol,Function}` | O(1) hash | ❌ Returns `Function` | ❌ Dynamic call | -> | `Val{:cpu}` dispatch | O(0) compiled | ✅ Concrete type | ✅ Full inlining | -> -> With Val dispatch, `@with_pool :cpu` compiles to exactly the same code as the -> original non-backend version—zero overhead. - -### 7.3 Extension Registration - -```julia -# ext/AdaptiveArrayPoolsCUDAExt/macros.jl - -# Add method for :cuda backend via Val dispatch (no __init__ needed!) -@inline AdaptiveArrayPools._get_pool_for_backend(::Val{:cuda}) = get_task_local_cuda_pool() - -# Optional: Explicit macro alias for users who prefer it -macro with_cuda_pool(pool_name, expr) - esc(:(@with_pool :cuda $pool_name $expr)) -end - -macro with_cuda_pool(expr) - esc(:(@with_pool :cuda $expr)) -end - -export @with_cuda_pool # Optional explicit alias -``` - -> **Note**: With Val dispatch, no `__init__` registration is needed. The method is added -> when the extension module loads, and Julia's method dispatch handles the rest. - -### 7.4 Design Trade-offs - -| Approach | Pros | Cons | -|----------|------|------| -| **Unified** (`@with_pool :cuda`) | Single API, easy switching, extensible | Symbol must be literal | -| **Explicit** (`@with_cuda_pool`) | Clear intent, better autocomplete | Multiple macros to learn | -| **Hybrid** (both available) | User choice | Slight API redundancy | - -**Recommendation: Hybrid approach** - unified macro as primary API, explicit aliases optional. - ---- - -## 8. Package Extension Structure - -### 8.1 Project.toml Changes - -```toml -[weakdeps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" - -[extensions] -AdaptiveArrayPoolsCUDAExt = "CUDA" -``` - -### 8.2 File Structure - -``` -AdaptiveArrayPools/ -├── src/ -│ ├── AdaptiveArrayPools.jl -│ ├── types.jl # + AbstractTypedPool{T,V}, AbstractArrayPool -│ ├── acquire.jl # + allocate_vector, wrap_array dispatch points -│ ├── state.jl # Generic _checkpoint/_rewind/_reset (unchanged logic) -│ ├── task_local_pool.jl # (unchanged) -│ ├── macros.jl # + _get_pool_for_backend(::Val{:cpu}) dispatch -│ └── utils.jl # (unchanged) -└── ext/ - └── AdaptiveArrayPoolsCUDAExt/ - ├── AdaptiveArrayPoolsCUDAExt.jl # ~25 lines - ├── types.jl # ~50 lines (no views field!) - ├── acquire.jl # ~30 lines (GPU-specific get_view!) - ├── dispatch.jl # ~35 lines (+ checkpoint correction) - ├── task_local_pool.jl # ~25 lines (multi-device, public API) - └── macros.jl # ~15 lines (@with_cuda_pool) -``` - -**Total extension: ~180 lines** (slightly more due to GPU-specific get_view!) - -### 8.3 Extension Entry Point - -```julia -# ext/AdaptiveArrayPoolsCUDAExt/AdaptiveArrayPoolsCUDAExt.jl - -module AdaptiveArrayPoolsCUDAExt - -using AdaptiveArrayPools -using AdaptiveArrayPools: AbstractTypedPool, AbstractArrayPool, - allocate_vector, wrap_array, get_view!, - _checkpoint_typed_pool!, _rewind_typed_pool!, - _reset_typed_pool!, _get_pool_for_backend, - CACHE_WAYS, checkpoint!, rewind!, reset! -using CUDA - -include("types.jl") -include("acquire.jl") # GPU-specific get_view! -include("dispatch.jl") -include("task_local_pool.jl") -include("macros.jl") - -# Exports -export CuAdaptiveArrayPool, CuTypedPool -export get_task_local_cuda_pool -export @with_cuda_pool - -end # module -``` - -### 8.4 dispatch.jl - -```julia -# ext/AdaptiveArrayPoolsCUDAExt/dispatch.jl - -# THE KEY DISPATCH METHODS - -@inline AdaptiveArrayPools.allocate_vector( - ::AbstractTypedPool{T,CuVector{T}}, n::Int -) where T = CuVector{T}(undef, n) - -@inline AdaptiveArrayPools.wrap_array( - ::AbstractTypedPool{T,CuVector{T}}, flat_view, dims::NTuple{N,Int} -) where {T,N} = reshape(flat_view, dims) - -# get_typed_pool! dispatches for CuAdaptiveArrayPool -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Float32}) = p.float32 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Float64}) = p.float64 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Float16}) = p.float16 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Int32}) = p.int32 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Int64}) = p.int64 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{ComplexF32}) = p.complexf32 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{ComplexF64}) = p.complexf64 -@inline AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{Bool}) = p.bool - -# Fallback for other types (with checkpoint correction!) -@inline function AdaptiveArrayPools.get_typed_pool!(p::CuAdaptiveArrayPool, ::Type{T}) where T - 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) - end - tp - end::CuTypedPool{T} -end -``` - -> **[Post-Review v0.6: Checkpoint Correction for Dynamic Pools]** -> When a new `CuTypedPool{T}` is created inside a `@with_pool` scope (i.e., when -> `_current_depth > 1`), we must initialize its checkpoint state to match the current -> depth. Otherwise, `rewind!` would pop from an incorrect checkpoint stack state. -> -> This mirrors the CPU implementation in `src/types.jl:230-238`. - ---- - -## 9. Memory Layout Clarification - -### 9.1 Why `Vector{CuVector{T}}` (not `CuVector{CuVector{T}}`) - -``` -✅ Correct: Vector{CuVector{T}} - - CPU RAM GPU VRAM - ┌─────────────────┐ ┌─────────────────┐ - │ Vector │ │ │ - │ ├─ CuVec meta1 ─┼──────────┼─► data1 [...] │ - │ ├─ CuVec meta2 ─┼──────────┼─► data2 [...] │ - │ └─ CuVec meta3 ─┼──────────┼─► data3 [...] │ - └─────────────────┘ └─────────────────┘ - - Pool management: CPU Computation: GPU -``` - -### 9.2 What Lives Where - -| Component | Location | Reason | -|-----------|----------|--------| -| Pool struct | CPU | Julia runtime | -| `vectors::Vector{...}` | CPU | Pool indexing | -| CuVector metadata | CPU | Julia object wrapper | -| CuVector data | **GPU** | Actual computation | -| n_active, checkpoints | CPU | State management | - ---- - -## 10. Migration Path - -### 10.1 Phase 1: Abstract Types (Non-Breaking) - -**Changes to src/:** -```julia -# types.jl -+ abstract type AbstractTypedPool{T, V<:AbstractVector{T}} end -+ abstract type AbstractArrayPool end -- mutable struct TypedPool{T} -+ mutable struct TypedPool{T} <: AbstractTypedPool{T, Vector{T}} -- mutable struct AdaptiveArrayPool -+ mutable struct AdaptiveArrayPool <: AbstractArrayPool - -# acquire.jl -+ allocate_vector(::AbstractTypedPool{T,Vector{T}}, n) where T = Vector{T}(undef, n) -+ wrap_array(::AbstractTypedPool{T,Vector{T}}, view, dims) where {T,N} = unsafe_wrap(...) -# Change get_view!, get_nd_array! signatures to use AbstractTypedPool - -# state.jl -# Change _checkpoint_typed_pool!, _rewind_typed_pool! to use AbstractTypedPool -``` - -**Breaking potential**: None (only adding supertypes and using more general signatures) - -### 10.2 Phase 2: CUDA Extension - -**New files in ext/:** -- Minimal implementation as described above - -**Breaking potential**: None (purely additive) - -### 10.3 Phase 3: Macro Enhancement (Optional) - -- Consider Option B unified macro -- Add `@with_cuda_pool` first, evaluate need for unification - ---- - -## 11. Example Usage (Target API) - -### 11.1 Basic Usage - Unified Macro - -```julia -using AdaptiveArrayPools -using CUDA # Triggers extension loading, registers :cuda backend - -# CPU workflow (default, unchanged) -function cpu_compute(data) - @with_pool pool begin - tmp = acquire!(pool, Float64, length(data)) - tmp .= data - sum(tmp) - end -end - -# GPU workflow - using :cuda backend symbol -function gpu_compute(data::CuVector) - @with_pool :cuda pool begin - A = acquire!(pool, Float32, 1000, 1000) # Returns CuMatrix{Float32} - B = acquire!(pool, Float32, 1000, 1000) - - A .= CUDA.rand(1000, 1000) - B .= A .* 2 - - sum(B) - end -end - -# Explicit CPU backend (equivalent to default) -function explicit_cpu_compute(data) - @with_pool :cpu pool begin - tmp = acquire!(pool, Float64, length(data)) - tmp .= data - sum(tmp) - end -end -``` - -### 11.2 Mixed CPU/GPU Workflow - -```julia -function mixed_compute(host_data::Vector{Float32}) - # CPU pool for staging - @with_pool cpu_pool begin - staging = acquire!(cpu_pool, Float32, length(host_data)) - staging .= host_data - - # Nested GPU pool - @with_pool :cuda gpu_pool begin - device_data = acquire!(gpu_pool, Float32, length(staging)) - copyto!(device_data, staging) # CPU → GPU - device_data .= device_data .^ 2 - copyto!(staging, device_data) # GPU → CPU - end # GPU pool rewinds here - - sum(staging) - end # CPU pool rewinds here -end -``` - -### 11.3 Without Pool Name (Auto-generated) - -```julia -# When you don't need to reference the pool directly -function simple_gpu_compute() - @with_pool :cuda begin - # pool name auto-generated, use get_task_local_cuda_pool() if needed - A = acquire!(get_task_local_cuda_pool(), Float32, 100, 100) - sum(A) - end -end - -# Or use the explicit getter within the block -function gpu_with_getter() - @with_pool :cuda begin - pool = get_task_local_cuda_pool() - A = acquire!(pool, Float32, 100, 100) - B = acquire!(pool, Float32, 100, 100) - A .+ B - end -end -``` - -### 11.4 Backend Switching (Same Code, Different Backend) - -```julia -# Parameterized backend - useful for testing/benchmarking -function compute_on_backend(data, backend::Symbol) - if backend == :cpu - @with_pool pool begin - tmp = acquire!(pool, Float32, length(data)) - tmp .= data - sum(tmp) - end - elseif backend == :cuda - @with_pool :cuda pool begin - tmp = acquire!(pool, Float32, length(data)) - tmp .= data - sum(tmp) - end - end -end - -# Note: Backend symbol must be literal in macro (compile-time) -# For runtime dispatch, use explicit pool getters: -function runtime_backend_dispatch(data, use_gpu::Bool) - pool = use_gpu ? get_task_local_cuda_pool() : get_task_local_pool() - checkpoint!(pool) - try - tmp = acquire!(pool, Float32, length(data)) - tmp .= data - sum(tmp) - finally - rewind!(pool) - end -end -``` - -### 11.5 Explicit Pool Management (Advanced) - -```julia -# Manual checkpoint/rewind for fine-grained control -function explicit_pool_management() - cpu = get_task_local_pool() - gpu = get_task_local_cuda_pool() - - # Checkpoint both pools - checkpoint!(cpu) - checkpoint!(gpu) - try - cpu_buf = acquire!(cpu, Float64, 1000) - gpu_buf = acquire!(gpu, Float32, 1000) - - # ... computation ... - - finally - # Rewind in reverse order (LIFO) - rewind!(gpu) - rewind!(cpu) - end -end - -# Clear GPU memory when under pressure -function memory_sensitive_workflow() - @with_pool :cuda pool begin - # Heavy GPU computation - A = acquire!(pool, Float32, 10000, 10000) - # ... - end - - # Explicitly free GPU memory if needed - empty!(get_task_local_cuda_pool()) - - # Continue with CPU work - @with_pool pool begin - # CPU pool unaffected - end -end -``` - -### 11.6 Future: Multiple GPU Backends - -```julia -# When Metal.jl extension is added (future) -using Metal # Registers :metal backend - -function apple_silicon_compute() - @with_pool :metal pool begin - A = acquire!(pool, Float32, 1000, 1000) # MtlMatrix{Float32} - # Metal-specific computation - end -end -``` - -> **[Post-Review v0.6: Backend Symbol Must Be Literal]** -> The macro `@with_pool :backend` requires a **literal symbol** (`:cuda`, `:metal`), -> not a variable containing a symbol. This is a Julia macro limitation—the backend -> is resolved at macro expansion time (compile time), not runtime. -> -> **This does NOT work:** -> ```julia -> const GPU_BACKEND = Sys.isapple() ? :metal : :cuda -> @with_pool GPU_BACKEND pool begin ... end # ERROR: GPU_BACKEND is not a QuoteNode -> ``` -> -> **For runtime backend selection, use explicit pool getters:** -> ```julia -> function portable_gpu_compute(use_metal::Bool) -> pool = use_metal ? get_task_local_metal_pool() : get_task_local_cuda_pool() -> checkpoint!(pool) -> try -> A = acquire!(pool, Float32, 1000, 1000) -> # ... computation ... -> finally -> rewind!(pool) -> end -> end -> ``` -> -> **Or use `@static` for compile-time platform selection:** -> ```julia -> function portable_gpu_compute() -> @static if Sys.isapple() -> @with_pool :metal pool begin -> # Metal path -> end -> else -> @with_pool :cuda pool begin -> # CUDA path -> end -> end -> end -> ``` - ---- - -## 12. Open Questions - -### 12.1 Resolved - -1. **Code duplication in extension** → Solved with parametric abstract types -2. **Macro approach** → Hybrid: unified `@with_pool :cuda` + optional `@with_cuda_pool` -3. **Memory layout** → `Vector{CuVector{T}}` is correct -4. **Float16 support** → **Added** to GPU fixed slots (per AI review) -5. **Multi-Device Safety** → **Solved** with `Dict{Int, Pool}` in task local storage (per AI review) -6. **unsafe_wrap for GPU** → Use `reshape` instead (per AI review) -7. **[v0.6] GPU view type** → `view(CuVector, 1:n)` returns `CuVector`, not `SubArray`. Pool design simplified. -8. **[v0.6] Zero-overhead backend selection** → `Val{:backend}` dispatch instead of Dict registry -9. **[v0.6] GPU checkpoint correction** → Added to `get_typed_pool!` fallback for `others` dict -10. **[v0.6] Device ID API** → Use `CUDA.deviceid(dev)` instead of internal `.handle` -11. **[v0.6] Backend symbol literal requirement** → Documented; `@static if` for platform selection - -### 12.2 Stream Synchronization (Critical Safety Documentation) - -> **[Post-Review v0.6: Expanded Safety Documentation]** - -**The Problem**: `rewind!` logically "frees" pooled memory. If a GPU kernel is still -running asynchronously using that memory, and the pool re-issues it for a new allocation, -**data corruption** or **use-after-free** occurs. - -**When It's Safe** (no synchronization needed): -- Single Task, default stream: Julia tasks typically use CUDA's default stream, which - serializes operations. `rewind!` happens after all prior operations complete. -- `CUDA.@sync` inside the block: Explicit synchronization before rewind. - -**When It's DANGEROUS** (must synchronize): - -1. **Passing arrays to other Tasks**: - ```julia - @with_pool :cuda pool begin - A = acquire!(pool, Float32, 1000) - @spawn begin - # DANGER: This task may still be using A after rewind! - expensive_computation!(A) - end - end # rewind! happens here - A is now invalid! - ``` - **Fix**: Wait for spawned task before exiting scope. - -2. **Explicit async streams**: - ```julia - @with_pool :cuda pool begin - A = acquire!(pool, Float32, 1000) - stream = CUDA.stream() - CUDA.@sync stream begin - # Kernel launched on non-default stream - my_kernel!(A; stream) - end - # If no @sync: kernel may still be running when rewind! executes - end - ``` - **Fix**: `CUDA.synchronize(stream)` or use `CUDA.@sync` before scope ends. - -3. **Kernel launch then immediate exit**: - ```julia - @with_pool :cuda pool begin - A = acquire!(pool, Float32, 1000) - @cuda threads=1024 my_kernel!(A) - # Kernel is async! May still be running... - end # rewind! immediately follows! - ``` - **Fix**: `CUDA.synchronize()` or `CUDA.@sync @cuda ...` - -**Recommendation for Documentation**: -```julia -# GPU POOLING SAFETY RULES -# -# 1. DO NOT pass pooled arrays to other Tasks without synchronization -# 2. DO synchronize before @with_pool block ends if using async streams -# 3. PREFER `CUDA.@sync` around kernel launches in pooled scopes -# 4. WHEN IN DOUBT: `CUDA.synchronize()` before the block ends -``` - -### 12.3 Still Open - -1. **Typed checkpoint for GPU**: Reuse existing macro logic? - - Should work with minimal changes - - Need to export `_transform_acquire_calls` etc. - -2. **resize! optimization for GPU** (v1.1+): - - Current: `resize!(CuVector, n)` copies old data (wasteful for pools) - - Consider: `CUDA.unsafe_free!` + fresh allocation, or pre-allocate oversized - -3. **Multi-backend single macro**: Support multiple pools in one call? - - Tuple syntax: `@with_pool (:cpu, cpu_pool) (:cuda, cuda_pool) begin ... end` - - Pro: Cleaner for mixed workflows, guaranteed proper rewind order - - Con: More complex macro implementation, less common use case - - Alternative: Nested `@with_pool` blocks (current approach) - - > **[AI Review]**: The tuple syntax is elegant but maybe over-engineering for V1. - ---- - -## 13. Summary: What Changes Where - -### src/ Changes (Phase 1) - -| File | Changes | -|------|---------| -| types.jl | Add abstract types, inherit from them | -| acquire.jl | Add `allocate_vector`, `wrap_array` dispatch points; generalize signatures | -| state.jl | Generalize to `AbstractTypedPool` | -| macros.jl | Add `_get_pool_for_backend(::Val{:cpu})` dispatch (NOT Dict registry) | -| Others | No changes | - -### ext/ New Files (Phase 2) - -| File | Lines | Content | -|------|-------|---------| -| AdaptiveArrayPoolsCUDAExt.jl | ~20 | Module, imports, exports | -| types.jl | ~50 | CuTypedPool (no views field!), CuAdaptiveArrayPool (+ Float16, device_id) | -| acquire.jl | ~30 | GPU-specific `get_view!` (no view caching) | -| dispatch.jl | ~35 | allocate_vector, wrap_array, get_typed_pool! (with checkpoint correction) | -| task_local_pool.jl | ~25 | get_task_local_cuda_pool (multi-device aware, public API) | -| macros.jl | ~25 | @with_cuda_pool | -| **Total** | **~155** | | - ---- - -## Changelog - -| Version | Date | Changes | -|---------|------|---------| -| 0.1 | 2024-12-10 | Initial draft | -| 0.2 | 2024-12-10 | Redesigned with parametric abstract types for maximum code reuse | -| 0.3 | 2024-12-10 | Unified macro design (`@with_pool :cuda`), comprehensive usage examples | -| 0.3.1 | 2024-12-10 | Added open questions: macro style preference, multi-backend single macro | -| 0.4 | 2024-12-10 | AI Review: Added Float16, device_id, multi-device pool getter, stream sync warning | -| 0.5 | 2024-12-10 | Merged AI feedback with restored full documentation | -| 0.6 | 2024-12-14 | **Post-Review Revision**: (1) Fixed GPU view type—`view(CuVector,1:n)` returns `CuVector` via GPUArrays `derive()`, not `SubArray`; simplified pool design by removing view caching. (2) Replaced Dict registry with `Val{:backend}` dispatch for zero-overhead backend selection. (3) Added checkpoint correction to GPU `get_typed_pool!` fallback. (4) Fixed `device_id` to use public API `CUDA.deviceid()`. (5) Clarified `empty!` semantics (reference removal ≠ VRAM release). (6) Documented `resize!` cost on GPU. (7) Expanded stream synchronization safety documentation. (8) Fixed backend symbol literal requirement (removed invalid `GPU_BACKEND` variable example). | diff --git a/docs/design/fixed_slots_codegen_design.md b/docs/design/fixed_slots_codegen_design.md deleted file mode 100644 index 7b57ec92..00000000 --- a/docs/design/fixed_slots_codegen_design.md +++ /dev/null @@ -1,364 +0,0 @@ -# Fixed Slots Iteration Automation Design Document - -## 1. Problem Definition - -### Current State -Iteration over fixed slot types was **manually repeated** across multiple functions: - -```julia -# Inside checkpoint! -_checkpoint_typed_pool!(pool.float64, depth) -_checkpoint_typed_pool!(pool.float32, depth) -_checkpoint_typed_pool!(pool.int64, depth) -_checkpoint_typed_pool!(pool.int32, depth) -_checkpoint_typed_pool!(pool.complexf64, depth) -_checkpoint_typed_pool!(pool.complexf32, depth) -_checkpoint_typed_pool!(pool.bool, depth) - -# Inside rewind! - same pattern repeated -_rewind_typed_pool!(pool.float64, depth) -_rewind_typed_pool!(pool.float32, depth) -... - -# Inside empty! - repeated again -empty!(pool.float64) -empty!(pool.float32) -... -``` - -### Improvement Goals -- **Centralized iteration logic**: Define once, use everywhere -- **Zero allocation**: No runtime overhead -- **IDE support preserved**: Keep struct definitions explicit - ---- - -## 2. Design Decision - -### Option B Adopted: const tuple + @generated (automate iteration only) - -**Core principle**: Keep struct definition manual, automate only iteration - -```julia -# 1. Keep struct explicitly defined (full IDE support) -mutable struct AdaptiveArrayPool - float64::TypedPool{Float64} - float32::TypedPool{Float32} - int64::TypedPool{Int64} - int32::TypedPool{Int32} - complexf64::TypedPool{ComplexF64} - complexf32::TypedPool{ComplexF32} - bool::TypedPool{Bool} - others::IdDict{DataType, Any} - _current_depth::Int - _untracked_flags::Vector{Bool} -end - -# 2. Define field names as const tuple -const FIXED_SLOT_FIELDS = (:float64, :float32, :int64, :int32, :complexf64, :complexf32, :bool) - -# 3. Use @generated for zero-allocation iteration -@generated function foreach_fixed_slot(f::F, pool::AdaptiveArrayPool) where {F} - exprs = [:(f(getfield(pool, $(QuoteNode(field))))) for field in FIXED_SLOT_FIELDS] - quote - $(exprs...) - nothing - end -end -``` - ---- - -## 3. Detailed Implementation - -### 3.1 types.jl Changes - -```julia -# ============================================================================== -# Fixed Slot Configuration -# ============================================================================== - -""" - FIXED_SLOT_FIELDS - -Fixed slot field names for iteration. Used by `foreach_fixed_slot`. - -Note: When adding/removing fixed slots, update BOTH: -1. This tuple -2. The AdaptiveArrayPool struct definition below -""" -const FIXED_SLOT_FIELDS = (:float64, :float32, :int64, :int32, :complexf64, :complexf32, :bool) - -# ============================================================================== -# AdaptiveArrayPool (explicit definition - full IDE support) -# ============================================================================== - -mutable struct AdaptiveArrayPool - # Fixed Slots: common types with zero lookup overhead - # NOTE: Keep in sync with FIXED_SLOT_FIELDS above - float64::TypedPool{Float64} - float32::TypedPool{Float32} - int64::TypedPool{Int64} - int32::TypedPool{Int32} - complexf64::TypedPool{ComplexF64} - complexf32::TypedPool{ComplexF32} - bool::TypedPool{Bool} - - # Fallback: rare types - others::IdDict{DataType, Any} - - # Untracked acquire detection - _current_depth::Int - _untracked_flags::Vector{Bool} -end - -# ... constructor, get_typed_pool! etc. remain unchanged ... - -# ============================================================================== -# Zero-Allocation Iteration -# ============================================================================== - -""" - foreach_fixed_slot(f, pool::AdaptiveArrayPool) - -Apply function `f` to each fixed slot TypedPool. -Zero allocation via compile-time unrolling. - -## Example -```julia -foreach_fixed_slot(pool) do tp - _checkpoint_typed_pool!(tp, depth) -end -``` -""" -@generated function foreach_fixed_slot(f::F, pool::AdaptiveArrayPool) where {F} - exprs = [:(f(getfield(pool, $(QuoteNode(field))))) for field in FIXED_SLOT_FIELDS] - quote - $(exprs...) - nothing - end -end -``` - -### 3.2 state.jl Changes - -```julia -function checkpoint!(pool::AdaptiveArrayPool) - pool._current_depth += 1 - push!(pool._untracked_flags, false) - depth = pool._current_depth - - # Fixed slots - zero allocation via @generated - foreach_fixed_slot(pool) do tp - _checkpoint_typed_pool!(tp, depth) - end - - # Others - fallback types - for p in values(pool.others) - _checkpoint_typed_pool!(p, depth) - end - nothing -end - -function rewind!(pool::AdaptiveArrayPool) - depth = pool._current_depth - - # Fixed slots - zero allocation - foreach_fixed_slot(pool) do tp - _rewind_typed_pool!(tp, depth) - end - - # Others - for tp in values(pool.others) - _rewind_typed_pool!(tp, depth) - end - - pop!(pool._untracked_flags) - pool._current_depth -= 1 - nothing -end - -function Base.empty!(pool::AdaptiveArrayPool) - # Fixed slots - foreach_fixed_slot(empty!, pool) - - # Others - for tp in values(pool.others) - empty!(tp) - end - empty!(pool.others) - - # Reset untracked detection state (1-based sentinel pattern) - pool._current_depth = 1 # 1 = global scope (sentinel) - empty!(pool._untracked_flags) - push!(pool._untracked_flags, false) # Sentinel: global scope starts with false - pool -end -``` - ---- - -## 4. Type Add/Remove Procedure - -### Adding UInt8 - -**Locations requiring manual update (2 places)**: - -```julia -# 1. Update FIXED_SLOT_FIELDS -const FIXED_SLOT_FIELDS = (:float64, :float32, :int64, :int32, :complexf64, :complexf32, :bool, :uint8) - -# 2. Update AdaptiveArrayPool struct -mutable struct AdaptiveArrayPool - float64::TypedPool{Float64} - float32::TypedPool{Float32} - int64::TypedPool{Int64} - int32::TypedPool{Int32} - complexf64::TypedPool{ComplexF64} - complexf32::TypedPool{ComplexF32} - bool::TypedPool{Bool} - uint8::TypedPool{UInt8} # ← Added - ... -end - -# 3. Update constructor -function AdaptiveArrayPool() - AdaptiveArrayPool( - TypedPool{Float64}(), - ... - TypedPool{UInt8}(), # ← Added - ... - ) -end - -# 4. Add get_typed_pool! dispatch -@inline get_typed_pool!(p::AdaptiveArrayPool, ::Type{UInt8}) = p.uint8 -``` - -**Automatically updated**: -- `checkpoint!` internal iteration -- `rewind!` internal iteration -- `empty!` internal iteration -- All code using `foreach_fixed_slot` - ---- - -## 5. Testing Strategy - -```julia -@testset "Fixed Slot Iteration" begin - pool = AdaptiveArrayPool() - - # Verify FIXED_SLOT_FIELDS and struct synchronization - for field in FIXED_SLOT_FIELDS - @test hasfield(AdaptiveArrayPool, field) - @test getfield(pool, field) isa TypedPool - end - - # Verify foreach_fixed_slot visits all slots - count = Ref(0) - foreach_fixed_slot(pool) do tp - count[] += 1 - end - @test count[] == length(FIXED_SLOT_FIELDS) - - # Zero allocation verification - pool2 = AdaptiveArrayPool() - foreach_fixed_slot(identity, pool2) # warmup - allocs = @allocated foreach_fixed_slot(identity, pool2) - @test allocs == 0 -end -``` - ---- - -## 6. Benefits - -### 6.1 Full IDE Support -- Explicit struct definition → autocomplete, Go to Definition work correctly -- Perfect type inference -- LSP/Language Server compatible - -### 6.2 Simple Implementation -- Single `@generated` function automates iteration -- No `@eval` needed → no precompilation concerns -- Most existing code preserved - -### 6.3 Easy Debugging -- Clear struct definition allows field inspection in debugger -- Compatible with tools like `@infiltrate` - -### 6.4 Zero Runtime Overhead -```julia -# @generated unrolls at compile time: -# foreach_fixed_slot(f, pool) is equivalent to: -f(pool.float64) -f(pool.float32) -f(pool.int64) -f(pool.int32) -f(pool.complexf64) -f(pool.complexf32) -f(pool.bool) -``` - ---- - -## 7. Drawbacks and Considerations - -### 7.1 Synchronization Required (2 places) -```julia -# These two locations must always be in sync: -const FIXED_SLOT_FIELDS = (:float64, :float32, ...) # 1 -mutable struct AdaptiveArrayPool # 2 - float64::TypedPool{Float64} - ... -end -``` - -**Mitigation**: Explicit comments + test verification - -### 7.2 @generated First-Call Cost -```julia -# Recompiles for different closures -foreach_fixed_slot(x -> checkpoint!(x, 1), pool) # Compiles -foreach_fixed_slot(x -> rewind!(x, 1), pool) # Compiles again -``` - -**Impact**: Slight effect on TTFX (Time To First X) -**Mitigation**: Warmup (precompile) at package load - -### 7.3 Metaprogramming Knowledge Required -```julia -@generated function foreach_fixed_slot(f::F, pool) where {F} - # Understanding this code requires @generated knowledge - exprs = [:(f(getfield(pool, $(QuoteNode(field))))) for field in FIXED_SLOT_FIELDS] - ... -end -``` - -**Mitigation**: Thorough comments and docstrings - ---- - -## 8. Option Comparison Summary - -| Aspect | Current (Manual) | Option B (Adopted) | Option C (@eval) | -|--------|------------------|-------------------|------------------| -| Modification locations | 6+ places | 2 places + α | 1 place | -| IDE support | Perfect | Perfect | Partial | -| Complexity | Low | Low | High | -| Debugging | Easy | Easy | Difficult | -| Type addition safety | May miss | Test-verified | Automatic | - ---- - -## 9. Conclusion - -**Reasons for adopting Option B**: - -1. **Practical balance**: Removes repetitive code without the complexity of full automation (Option C) -2. **IDE support preserved**: Maintains the most important developer experience -3. **Low risk**: Uses only `@generated` without `@eval`, ensuring precompilation stability -4. **Incremental improvement**: Improves only iteration while preserving most existing code - -Since type changes are rare (1-2 times during package lifetime), and struct definition synchronization across 2 locations can be sufficiently verified by tests, Option B is the optimal choice. diff --git a/docs/design/hybrid_api_design.md b/docs/design/hybrid_api_design.md deleted file mode 100644 index b05d0f17..00000000 --- a/docs/design/hybrid_api_design.md +++ /dev/null @@ -1,527 +0,0 @@ -# Hybrid API Design: acquire! vs unsafe_acquire! - -> **Update (v0.2.2, feat/new_array_nd)**: The `unsafe_acquire!` path no longer uses -> `unsafe_wrap` + N-way cache on Julia 1.11+ CPU. Instead, it uses `setfield!`-based -> wrapper reuse — **0-alloc for any number of dimension patterns** (no eviction limit). -> The N-way cache (`CACHE_WAYS`) is now only used by the **CUDA** backend and the -> **Julia 1.10 legacy** fallback. The `acquire!` → `ReshapedArray` path is unchanged. -> `TypedPool` fields changed: `nd_arrays`/`nd_dims`/`nd_ptrs`/`nd_next_way` → -> `nd_wrappers::Vector{Union{Nothing, Vector{Any}}}`. - -## Executive Summary - -Redesigning `AdaptiveArrayPools.jl`'s N-D array acquisition API with a **Two Tools Strategy**: - -| API | Return Type | Use Case | Allocation Characteristics | -|-----|-------------|----------|---------------------------| -| `acquire!` | `ReshapedArray` (fixed) | General use, Static dispatch | No cache needed, relies on compiler optimization | -| `unsafe_acquire!` | `Array` (fixed) | Dynamic dispatch, FFI | Cache hit: 0, miss: 112 bytes | - -**Core Principle**: Return type does not change based on state (Type Stability guaranteed) - -> **Note**: ReshapedArray's "0 allocation" depends on compiler's SROA (Scalar Replacement of Aggregates) and -> escape analysis. Not always guaranteed - heap allocation may occur if the value escapes from the function. - ---- - -## Problem Statement - -### Current State (v0.2.0) - -``` -acquire!(pool, T, dims...) - └─> get_nd_view!() - └─> get_nd_array!() ← uses unsafe_wrap - └─> 112 bytes on cache miss! -``` - -- Both `acquire!` and `unsafe_acquire!` internally use `unsafe_wrap` -- Always 112 bytes allocation on cache miss -- Tried to reduce miss rate with N-way cache, but 100% miss on cyclic patterns - -### v0.1.2 Approach - -``` -acquire!(pool, T, dims...) - └─> get_view!(tp, total_len) ← 1D view (cached) - └─> reshape(view, dims) ← 0 bytes always! -``` - -- `reshape(view, dims)` creates a wrapper object, but heap allocation can be avoided via compiler optimization (SROA/escape analysis) -- Simple and predictable - ---- - -## Why Not Mixed Return Types? - -### Proposed (but rejected) Approach - -```julia -# ❌ BAD: Array on cache hit, View on miss -function acquire!(pool, T, dims...) - if cache_hit - return cached_array::Array{T,N} - else - return reshape(view, dims)::ReshapedArray{...} - end -end -``` - -### Problem: Type Instability - -| Aspect | Impact | -|--------|--------| -| **Compiler inference** | `Union{Array, ReshapedArray}` → Union splitting or dynamic dispatch | -| **Performance** | Execution slowdown while trying to achieve zero-alloc | -| **API semantics** | Same function returning different types → confusion | -| **Module boundaries** | Inference widens when storing result or passing to other modules | - -**AI Feedback Quote**: -> "State-dependent returns become Union{Array, ReshapedArray} from external view, breaking API-level type stability." - ---- - -## Recommended Design: Two Tools Strategy - -### Principles - -1. **Fixed return type**: Each API always returns the same type -2. **Purpose separation**: Users choose API based on situation -3. **Simple implementation**: Minimize complex cache logic - -### API Design - -#### 1. `acquire!` → ReshapedArray (regression to v0.1.2 style) - -```julia -@inline function acquire!(pool::AdaptiveArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N} - tp = get_typed_pool!(pool, T) - total_len = safe_prod(dims) - flat_view = get_view!(tp, total_len) # 1D view (cached, 0 alloc) - return reshape(flat_view, dims) # ReshapedArray (0 alloc always!) -end -``` - -**Characteristics**: -- Always returns `ReshapedArray{T, N, SubArray{T, 1, Vector{T}, ...}, ...}` -- No `unsafe_wrap` call → no Array header creation cost (112B) even on cache miss -- N-way cache unnecessary (simple 1D view cache sufficient) - -**Use Cases**: -- General `Flux` layers (`mul!`, `broadcast`) -- Code where static dispatch is guaranteed -- Most use cases - -**Constraints**: -- Escape optimization may fail in type-unspecified call paths, causing wrapper allocation -- Incompatible with APIs requiring strict `Array` type (rare) - -#### 2. `unsafe_acquire!` → Array (maintains v0.2.0 + N-way Cache) - -```julia -@inline function unsafe_acquire!(pool::AdaptiveArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N} - tp = get_typed_pool!(pool, T) - return get_nd_array!(tp, dims) # Array with slot-based + N-way cache -end -``` - -**Characteristics**: -- Always returns `Array{T, N}` -- Cache hit: 0 bytes, Cache miss: 112 bytes -- Maintains existing N-way cache (4-way) - -**Use Cases**: -- Type-unspecified call paths (e.g., `TGLFNNmodel._pooled_chain` - no concrete type parameters) -- FFI / ccall -- Special APIs that strictly require `StridedArray` -- Places requiring strict `Array` type - -**Benefits**: -- Array is already a heap-allocated object → no additional allocation when reusing cached instance -- Avoids wrapper object optimization issues in type-unspecified paths - -### API Aliases - -For clarity, explicit aliases are provided: - -```julia -# Main APIs -export acquire!, unsafe_acquire! - -# Explicit Aliases -export acquire_view!, acquire_array! - -"""Alias for [`acquire!`](@ref). Returns a ReshapedArray (View).""" -const acquire_view! = acquire! - -"""Alias for [`unsafe_acquire!`](@ref). Returns an Array (via unsafe_wrap).""" -const acquire_array! = unsafe_acquire! -``` - ---- - -## Comparison Matrix - -| Strategy | Return Type | Cache Miss Cost | Type-Unspecified Path* | Type Stable | -|----------|-------------|-----------------|------------------------|-------------| -| **acquire! (new)** | `ReshapedArray` | **0 bytes** (no unsafe_wrap) | May allocate wrapper | **✓** | -| **unsafe_acquire!** | `Array` | 112 bytes | **0 bytes** (on cache hit) | **✓** | -| ~~Mixed (rejected)~~ | `Union{...}` | 0 bytes | Unspecified | **✗** | - -*Type-unspecified path: Calls through abstract fields without concrete type parameters, etc. Compiler cannot apply escape optimization, causing wrapper object heap allocation. - -### Recommended API by Situation - -| Situation | acquire! | unsafe_acquire! | Recommendation | -|-----------|----------|-----------------|----------------| -| Type-specified call path | Optimizable | 0 bytes (hit) | `acquire!` | -| Variable dims (cyclic pattern) | Optimizable | cache miss occurs | `acquire!` | -| Type-unspecified path | Wrapper alloc | **0 bytes** (hit) | **`unsafe_acquire!`** | -| FFI / raw pointer | N/A | 0 bytes | `unsafe_acquire!` | - ---- - -## Implementation Specification - -### Data Structure Changes (types.jl) - -#### Constants -```julia -const CACHE_WAYS = 4 -``` - -#### TypedPool Struct Layout - -```julia -mutable struct TypedPool{T} - # --- Backing Storage --- - vectors::Vector{Vector{T}} - - # --- 1D Cache (Simple 1-way or Direct) --- - views::Vector{SubArray{T, 1, Vector{T}, Tuple{UnitRange{Int}}, true}} - view_lengths::Vector{Int} - - # --- N-D Array Cache (N-way Set Associative) --- - # Layout: Flat Vector. Index = (slot_idx - 1) * CACHE_WAYS + way_idx - nd_arrays::Vector{Any} # Stores Array{T, N} - nd_dims::Vector{Any} # Stores NTuple{N, Int} - nd_ptrs::Vector{UInt} # Stores objectid/pointer for validation - nd_next_way::Vector{Int} # Round-Robin counter per slot (1 per slot) - - # --- State --- - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} -end -``` - -**Implemented Changes** (compared to initial design): -- **Removed**: `nd_views` (No longer needed as `acquire!` returns `ReshapedArray`) ✓ -- **Updated**: `nd_arrays`, `nd_dims`, `nd_ptrs` store `CACHE_WAYS` items per active slot ✓ -- **Added**: `nd_next_way::Vector{Int}` for Round-Robin replacement index per slot ✓ - -### Logic Implementation (core.jl) - -#### `acquire!` (The Fast Path) - -**Goal**: Always return `ReshapedArray`. No N-D cache lookup. - -```julia -@inline function get_nd_view!(tp::TypedPool{T}, dims::NTuple{N, Int}) where {T, N} - len = safe_prod(dims) - flat_view = get_view!(tp, len) - return reshape(flat_view, dims) -end -``` - -#### `unsafe_acquire!` (The N-way Cache Path) - -**Goal**: Return `Array`. Use N-way cache with Linear Search + Round-Robin Replacement. - -**Algorithm**: -1. Get 1D view: `flat_view = get_view!(tp, prod(dims))` -2. Get current pointer: `current_ptr = UInt(pointer(flat_view))` -3. Calculate Base Index: `base = (tp.n_active - 1) * CACHE_WAYS` -4. **Search (Hit Check)**: - - Loop `k` from `1` to `CACHE_WAYS` - - Check if `nd_dims[base + k] == dims` **AND** `nd_ptrs[base + k] == current_ptr` - - If match: Return `nd_arrays[base + k]` -5. **Miss (Replacement)**: - - Get victim way from `nd_next_way[tp.n_active]` - - Target Index: `target = base + victim_way + 1` - - Create Array: `arr = unsafe_wrap(Array{T, N}, pointer(flat_view), dims)` - - **Update Cache**: - - `nd_arrays[target] = arr` - - `nd_dims[target] = dims` - - `nd_ptrs[target] = current_ptr` - - **Update Round-Robin**: Increment `nd_next_way` (modulo `CACHE_WAYS`) - - Return `arr` - ---- - -## Implementation Plan - -### Phase 1: Simplify `acquire!` N-D Path - -**File**: `src/core.jl` - -**Before**: -```julia -@inline function get_nd_view!(tp::TypedPool{T}, dims::NTuple{N, Int}) where {T, N} - arr = get_nd_array!(tp, dims) # uses unsafe_wrap - idx = tp.n_active - # ... complex caching logic - new_view = view(arr, ntuple(_ -> Colon(), Val(N))...) - return new_view # SubArray{T,N,Array{T,N}} -end -``` - -**After**: -```julia -@inline function get_nd_view!(tp::TypedPool{T}, dims::NTuple{N, Int}) where {T, N} - total_len = safe_prod(dims) - flat_view = get_view!(tp, total_len) # 1D view (cached) - return reshape(flat_view, dims) # ReshapedArray (0 alloc!) -end -``` - -### Phase 2: Maintain `unsafe_acquire!` Cache - -**No changes** - maintain current implementation: -- `get_nd_array!` → `unsafe_wrap` + slot-based cache -- Maintain N-way cache (4-way) -- 112 bytes allocation on cache miss - -### Phase 3: TypedPool Field Updates - -Update struct as specified in Implementation Specification above. - -### Phase 4: Test Updates - -**Files**: `test/test_nway_cache.jl`, `test/test_zero_allocation.jl` - -```julia -@testset "acquire! returns ReshapedArray" begin - pool = AdaptiveArrayPool() - @with_pool pool begin - m = acquire!(pool, Float64, 10, 10) - @test m isa Base.ReshapedArray - @test size(m) == (10, 10) - end -end - -@testset "acquire! is always zero-allocation" begin - pool = AdaptiveArrayPool() - - # 5-way cycling (exceeds any cache) - still 0 alloc! - function test_5way!(p) - dims_list = ((5, 10), (10, 5), (7, 7), (3, 16), (4, 12)) - for dims in dims_list - checkpoint!(p) - acquire!(p, Float64, dims...) # ReshapedArray - rewind!(p) - end - end - - test_5way!(pool); test_5way!(pool) - allocs = @allocated test_5way!(pool) - @test allocs == 0 # Always zero, regardless of pattern! -end - -@testset "unsafe_acquire! returns Array" begin - pool = AdaptiveArrayPool() - @with_pool pool begin - m = unsafe_acquire!(pool, Float64, 10, 10) - @test m isa Array - @test size(m) == (10, 10) - end -end -``` - -### Phase 5: Documentation - -**CHANGELOG.md** (not a breaking change, behavior improvement): -```markdown -## [Unreleased] -### Changed -- `acquire!` N-D path now returns `ReshapedArray` instead of `SubArray{Array}` - - Always zero-allocation, regardless of cache hit/miss - - Simpler implementation, no N-D cache dependency -- `unsafe_acquire!` continues to return `Array` with N-way cache - - Use this when dynamic dispatch or raw Array is needed -``` - -**Docstring Updates**: -```julia -""" - acquire!(pool, Type{T}, dims...) -> ReshapedArray{T,N,...} - -Acquire a view with dimensions `dims` from the pool. - -Returns a `ReshapedArray` backed by pool memory. **Zero creation cost** - no -`unsafe_wrap` call needed. Compiler may optimize away heap allocation via -SROA/escape analysis in type-specified paths. - -For type-unspecified paths (struct fields without concrete type parameters), -use [`unsafe_acquire!`](@ref) instead - cached Array instances can be reused. - -## Example -```julia -@with_pool pool begin - m = acquire!(pool, Float64, 64, 100) # ReshapedArray - m .= 1.0 - result = sum(m) -end -``` -""" -``` - -```julia -""" - unsafe_acquire!(pool, Type{T}, dims...) -> Array{T,N} - -Acquire a raw `Array` backed by pool memory. - -Returns an `Array` object. Since Array is already heap-allocated, the cached -instance can be reused without wrapper allocation overhead. - -## When to use -- Type-unspecified paths (e.g., struct fields without concrete type parameters) -- FFI / ccall requiring raw pointers -- APIs that strictly require `Array` type - -## Allocation behavior -- Cache hit: 0 bytes (cached Array instance reused) -- Cache miss: 112 bytes (Array header creation) - -## Example -```julia -@with_pool pool begin - m = unsafe_acquire!(pool, Float64, 64, 100) # Matrix{Float64} - # Safe for type-unspecified paths - some_abstract_field.process(m) # 0 bytes - cached instance reused -end -``` -""" -``` - ---- - -## Verification Checklist - -1. **Type Check**: `acquire!` must return `ReshapedArray`. `unsafe_acquire!` must return `Array`. -2. **Allocation Check**: - - `acquire!`: 0 allocations always - - `unsafe_acquire!`: 0 allocations on cache hit - - `unsafe_acquire!`: 0 allocations on interleaved access (e.g., alternating 10x10 and 20x20) thanks to N-way cache -3. **Safety**: Ensure `unsafe_acquire!` validates pointers (re-wraps if the backing vector was resized) - ---- - -## TurbulentTransport Integration - -### Changed File: `src/tglf_nn.jl` - -**Already Applied** (line 277): -```julia -@with_pool pool function flux_array!(out_y::AbstractMatrix{T}, fluxmodel::TGLFNNmodel, x::AbstractMatrix{T}; ...) where {T<:Real} - # ... - # NOTE: Use unsafe_acquire! (returns Array) instead of acquire! (returns ReshapedArray) - # because _pooled_chain field lacks concrete type parameters, causing - # escape optimization failure. Array (cached instance) avoids wrapper allocation. - xx = unsafe_acquire!(pool, T, size(x)) - # ... - fluxmodel._pooled_chain(out_y, xx) # 0 bytes - cached Array instance reused -end -``` - -### No Change Needed: `src/pooled_layers.jl` - -`PooledDense`, `PooledActivation` are in **static dispatch** environment: -- Types are known at compile time -- Maintain use of `acquire!` (ReshapedArray) -- ReshapedArray is also 0 bytes in static dispatch - -```julia -@inline function _pooled_dense_forward!(pd::PooledDense, x::AbstractVecOrMat) - pool = get_task_local_pool() - # acquire! usage OK - static dispatch environment - out = acquire!(pool, Float64, size(d.weight, 1), size(xT, 2)) - mul!(out, d.weight, xT) # ReshapedArray is StridedArray ✓ - return Flux.NNlib.bias_act!(d.σ, out, d.bias) -end -``` - ---- - -## Summary - -### Before (v0.2.0) - -``` -┌─────────────────────────────────────────────────────────────┐ -│ acquire!() ──┬──> get_nd_view!() ──> get_nd_array!() │ -│ │ │ │ │ -│ │ │ unsafe_wrap (112B miss) │ -│ │ │ ↓ │ -│ │ └──────> SubArray{Array} ←──────────┘ -│ │ │ -│ unsafe_acquire!() ──> get_nd_array!() ──> Array │ -│ │ │ -│ unsafe_wrap (112B miss) │ -└─────────────────────────────────────────────────────────────┘ -``` - -### After (Hybrid) - -``` -┌─────────────────────────────────────────────────────────────┐ -│ acquire!() ──> get_view!() ──> reshape() ──> ReshapedArray │ -│ │ │ │ -│ 1D cache 0 bytes always! │ -│ (0 alloc) │ -│ │ -│ unsafe_acquire!() ──> get_nd_array!() ──> Array │ -│ │ │ -│ unsafe_wrap + N-way cache │ -│ (0B hit, 112B miss) │ -└─────────────────────────────────────────────────────────────┘ -``` - -### Decision Matrix for Users - -``` -┌─────────────────────────────────────────────────────────────┐ -│ Which API to use? │ -├─────────────────────────────────────────────────────────────┤ -│ │ -│ Is the code path type-unspecified? │ -│ (abstract fields without concrete type params, │ -│ runtime-determined function calls) │ -│ │ -│ YES ──────────────> unsafe_acquire!() │ -│ │ │ │ -│ │ Returns Array │ -│ │ (cached instance reused) │ -│ │ │ -│ NO ───────────────> acquire!() │ -│ │ │ -│ Returns ReshapedArray │ -│ (0 bytes creation) │ -│ │ -└─────────────────────────────────────────────────────────────┘ -``` - ---- - -## Resolved Questions - -1. **N-way cache retention level**: Configurable via `CACHE_WAYS` preference (default: 4-way). ✅ -2. **nd_views field removal**: Removed. `acquire!` now returns `ReshapedArray` via `reshape()`. ✅ -3. **Backward compatibility**: `acquire!` returns `ReshapedArray` (a type of `AbstractArray`), maintaining API compatibility. ✅ - ---- - -## References - -- [nd_array_approach_comparison.md](./nd_array_approach_comparison.md) - Benchmark results and boxing analysis diff --git a/docs/design/macro-linenumbernode-improvement.md b/docs/design/macro-linenumbernode-improvement.md deleted file mode 100644 index 26eab676..00000000 --- a/docs/design/macro-linenumbernode-improvement.md +++ /dev/null @@ -1,369 +0,0 @@ -# @with_pool Macro LineNumberNode Improvement Plan - -## Goal -Utilize `__source__` and `LineNumberNode` to improve coverage, stack trace, and debugging - -## Target File -- `/Users/yoo/.julia/dev/AdaptiveArrayPools/src/macros.jl` - ---- - -## Phase 1: Add Helper Functions - -### 1.1 LineNumberNode Insertion Helper (New) - -**Location**: Add before `_generate_pool_code` function - -```julia -""" - _maybe_add_source_location!(expr, source) - -Insert source location LineNumberNode at the beginning of an Expr block. -No-op if source is nothing or expr is not an Expr(:block, ...). -""" -function _maybe_add_source_location!(expr::Expr, source::Union{LineNumberNode,Nothing}) - if source !== nothing && expr.head === :block - pushfirst!(expr.args, LineNumberNode(source.line, source.file)) - end - return expr -end -_maybe_add_source_location!(expr, ::Nothing) = expr -``` - -**Benefits**: Common application across all return paths, reduces risk of omission/drift - -### 1.2 Function Body LineNumberNode Correction Helper (New) - -**Goal**: Correct with `__source__` **only when no LNN exists at the body top level** - -> **Background**: `body` is the **user code AST** obtained from `func_def.args[2]`. -> Existing LNNs point to **user file lines** and must be preserved. -> The problem is **short function forms** like `f(x) = ...` that have no LNN. - -```julia -""" - _has_toplevel_lnn(body) -> Bool - -Check if body has a LineNumberNode at the top level (within first few args). -More robust than checking only args[1], handles Expr(:meta) etc. -""" -function _has_toplevel_lnn(body) - body isa Expr && body.head === :block || return false - # Check first 3 args for LNN (handles :meta, :line annotations, etc.) - for i in 1:min(3, length(body.args)) - body.args[i] isa LineNumberNode && return true - end - return false -end - -""" - _ensure_body_has_toplevel_lnn(body, source) - -Ensure body has a LineNumberNode at the top level. -- If body already has a top-level LNN, preserve it (user file line info) -- If not, prepend source LNN (macro call location as fallback) - -Returns a new Expr to avoid mutating the original AST. -""" -function _ensure_body_has_toplevel_lnn(body, source::Union{LineNumberNode,Nothing}) - source === nothing && return body - - # Check if top-level LNN already exists (robust check) - if _has_toplevel_lnn(body) - return body # Preserve existing user file LNN - end - - # No top-level LNN → add source as fallback (no mutation) - lnn = LineNumberNode(source.line, source.file) - if body isa Expr && body.head === :block - return Expr(:block, lnn, body.args...) - else - return Expr(:block, lnn, body) - end -end -``` - -**Benefits**: -- **User body LNN preserved**: If existing top-level LNN exists, keep it (accurate body line) -- **Short function form handling**: If no LNN, correct with `__source__` -- **Mutation prevention**: Returns new Expr to protect original AST - ---- - -## Phase 2: Modify Helper Function Signatures - -### 2.1 Keyword Argument Approach (Recommended) - -To avoid the risk of fixed `:cpu` default, add `source` as keyword argument: - -| Function | Search Pattern | Change | -|----------|----------------|--------| -| `_generate_pool_code` | `function _generate_pool_code(pool_name, expr, force_enable)` | `(...; source::Union{LineNumberNode,Nothing}=nothing)` | -| `_generate_pool_code_with_backend` | `function _generate_pool_code_with_backend(backend::Symbol, pool_name, expr, force_enable::Bool)` | `(...; source::Union{LineNumberNode,Nothing}=nothing)` | -| `_generate_function_pool_code` | `function _generate_function_pool_code(pool_name, func_def, force_enable, disable_pooling, backend::Symbol=:cpu)` | `(...; source::Union{LineNumberNode,Nothing}=nothing)` | -| `_generate_function_pool_code_with_backend` | `function _generate_function_pool_code_with_backend(backend::Symbol, pool_name, func_def, disable_pooling::Bool)` | `(...; source::Union{LineNumberNode,Nothing}=nothing)` | - -**Benefits**: Minimal changes to existing call sites, solves `backend` default value issue - ---- - -## Phase 3: Pass source to Internal Function Calls - -### Inside `_generate_pool_code` (Search: `_generate_function_pool_code(pool_name`) - -```julia -# Before -return _generate_function_pool_code(pool_name, expr, force_enable, true, :cpu) -return _generate_function_pool_code(pool_name, expr, force_enable, false) - -# After (pass as keyword argument) -return _generate_function_pool_code(pool_name, expr, force_enable, true, :cpu; source) -return _generate_function_pool_code(pool_name, expr, force_enable, false; source) -``` - -### Inside `_generate_pool_code_with_backend` (Search: `_generate_function_pool_code_with_backend(backend`) - -```julia -# After -_generate_function_pool_code_with_backend(backend, pool_name, expr, ...; source) -``` - ---- - -## Phase 4: LineNumberNode Insertion - -Call helper before returning each `quote ... end` block: - -```julia -result = quote - # ... generated code ... -end -_maybe_add_source_location!(result, source) -return result -``` - -### Insertion Locations (Based on Search Patterns) - -**`_generate_pool_code`** (Search: `function _generate_pool_code`): -- `return quote ... end` in `!USE_POOLING` branch -- `return quote ... end` in `force_enable` branch -- `return quote ... end` in `else` branch - -**`_generate_pool_code_with_backend`** (Search: `function _generate_pool_code_with_backend`): -- All `return quote ... end` with same pattern - -**`_generate_function_pool_code`** (Search: `function _generate_function_pool_code`): -- After `transformed_body` creation: `transformed_body = _ensure_body_has_toplevel_lnn(transformed_body, source)` -- Also correct `body` in `disable_pooling` path: `body = _ensure_body_has_toplevel_lnn(body, source)` -- Then `new_body = quote ... end` -- **(Optional) To make wrapper appear as call-site**: Add `_maybe_add_source_location!(new_body, source)` - -**`_generate_function_pool_code_with_backend`** (Search: `function _generate_function_pool_code_with_backend`): -- Apply `_ensure_body_has_toplevel_lnn(..., source)` to `transformed_body` and `body` before constructing `new_body` -- **(Optional)** Can also apply `_maybe_add_source_location!` to `new_body` - -> **Core Principle**: Preserve if top-level LNN exists (user line), correct with `__source__` if not (short function form) -> **Note**: Wrapper code (checkpoint/try/finally) lines may still point to macros.jl. Inserting LNN to `new_body` improves this but is not required. - ---- - -## Phase 5: Macro Definition Modifications - -**@with_pool** (Search: `macro with_pool`): -```julia -macro with_pool(pool_name, expr) - _generate_pool_code(pool_name, expr, true; source=__source__) -end -# Same pattern for remaining 3 -``` - -**@maybe_with_pool** (Search: `macro maybe_with_pool`): -```julia -macro maybe_with_pool(pool_name, expr) - _generate_pool_code(pool_name, expr, false; source=__source__) -end -# Same pattern for remaining 3 -``` - ---- - -## Phase 6: Testing - -### 6.1 Robust Search Helpers (test/test_macro_expansion.jl) - -```julia -""" - find_linenumbernode_with_line(expr, target_line) -> Union{LineNumberNode, Nothing} - -Recursively search for a LineNumberNode matching target_line. -More robust than checking only the first LNN (handles block forms where -_maybe_add_source_location! may insert LNN before user code LNN). -""" -function find_linenumbernode_with_line(expr, target_line::Int) - if expr isa LineNumberNode && expr.line == target_line - return expr - elseif expr isa Expr - for arg in expr.args - result = find_linenumbernode_with_line(arg, target_line) - result !== nothing && return result - end - end - return nothing -end - -""" - has_valid_linenumbernode(expr) -> Bool - -Check if expr contains any LineNumberNode with valid line info. -""" -function has_valid_linenumbernode(expr) - if expr isa LineNumberNode - return expr.line > 0 && expr.file !== :none - elseif expr isa Expr - for arg in expr.args - has_valid_linenumbernode(arg) && return true - end - end - return false -end - -""" - get_function_body(expr) -> Union{Expr, Nothing} - -Extract function body from a function definition expression. -Handles both `function f() ... end` and `f() = ...` forms. -""" -function get_function_body(expr) - if expr isa Expr - if expr.head === :function && length(expr.args) >= 2 - return expr.args[2] - elseif expr.head === :(=) && expr.args[1] isa Expr && expr.args[1].head === :call - return expr.args[2] - end - # Recurse for wrapped expressions - for arg in expr.args - result = get_function_body(arg) - result !== nothing && return result - end - end - return nothing -end -``` - -### 6.2 Test Cases (Full Coverage) - -> **Test Strategy**: Verify "existence of LNN matching expected line" rather than "first LNN". -> In block forms, `_maybe_add_source_location!` may insert additional LNNs, -> so checking for existence of LNN with specific line is more robust. - -```julia -@testset "Source location preservation" begin - # Test 1: @with_pool block form - @testset "@with_pool block" begin - expected_line = @__LINE__ + 2 - expr = @macroexpand @with_pool pool begin - v = acquire!(pool, Float64, 10) - end - # Check if LNN matching expected line exists - lnn = find_linenumbernode_with_line(expr, expected_line) - @test lnn !== nothing - @test lnn.file !== :none - # At minimum, valid LNN must exist - @test has_valid_linenumbernode(expr) - end - - # Test 2: @with_pool function form - @testset "@with_pool function" begin - expected_line = @__LINE__ + 2 - func_expr = @macroexpand @with_pool pool function test_func(n) - acquire!(pool, Float64, n) - end - body = get_function_body(func_expr) - @test body !== nothing - lnn = find_linenumbernode_with_line(body, expected_line) - @test lnn !== nothing - end - - # Test 3: @maybe_with_pool - @testset "@maybe_with_pool" begin - expected_line = @__LINE__ + 2 - expr = @macroexpand @maybe_with_pool pool begin - v = acquire!(pool, Float64, 10) - end - lnn = find_linenumbernode_with_line(expr, expected_line) - @test lnn !== nothing - end - - # Test 4: Backend variant (@with_pool :cpu) - @testset "@with_pool :cpu backend" begin - expected_line = @__LINE__ + 2 - expr = @macroexpand @with_pool :cpu pool begin - v = acquire!(pool, Float64, 10) - end - lnn = find_linenumbernode_with_line(expr, expected_line) - @test lnn !== nothing - end - - # Test 5: Without pool name (implicit gensym) - @testset "@with_pool without pool name" begin - expected_line = @__LINE__ + 2 - expr = @macroexpand @with_pool begin - inner_function() - end - lnn = find_linenumbernode_with_line(expr, expected_line) - @test lnn !== nothing - end - - # Test 6: Short-form function (f(x) = ...) - Case without LNN, corrected with __source__ - @testset "@with_pool short function" begin - expected_line = @__LINE__ + 1 - func_expr = @macroexpand @with_pool pool test_func(x) = acquire!(pool, Float64, x) - body = get_function_body(func_expr) - @test body !== nothing - # Short function originally has no LNN, so corrected with __source__ - lnn = find_linenumbernode_with_line(body, expected_line) - @test lnn !== nothing - end -end -``` - -### 6.3 Verification Command -```bash -julia --project -e 'using Pkg; Pkg.test()' -``` - ---- - -## Expected Results - -| Item | Before Improvement | After Improvement | -|------|-------------------|-------------------| -| Coverage | signature uncovered | Properly mapped | -| Stack trace | macros.jl:XXX | Original source:line | -| Breakpoint | Inside macros.jl | Improved to inside body | - ---- - -## Considerations - -1. **Use Keyword Arguments**: Add `source` as keyword arg to minimize impact on existing call sites -2. **Use Helper Functions**: Use `_maybe_add_source_location!` for consistent insertion across all paths -3. **Body Line Correction**: Use `_ensure_body_has_toplevel_lnn` to preserve top-level LNN, correct with `__source__` if not present -4. **Robust Tests**: Search-based verification resistant to AST structure changes + line number accuracy verification -5. **esc() Interaction**: `LineNumberNode` is unrelated to hygiene → insert at quote block top -6. **try-finally**: Lines inside wrapper still point to macros.jl (acceptable) -7. **CUDA Extension**: Only registers backend dispatch, no macro definitions → no changes needed - ---- - -## Change Summary - -| Phase | Work | Estimated Change | -|-------|------|-----------------| -| 1 | Add helper functions (3: `_maybe_add_source_location!`, `_has_toplevel_lnn`, `_ensure_body_has_toplevel_lnn`) | +35 lines | -| 2 | Modify signatures (4 functions) | 4 lines modified | -| 3 | Modify internal calls | ~5 lines modified | -| 4 | LineNumberNode insertion | ~10 lines added | -| 5 | Modify macro definitions (8) | 8 lines modified | -| 6 | Add tests (3 helpers + 6 tests) | +80 lines | -| **Total** | | ~140 lines | diff --git a/docs/design/nd_array_approach_comparison.md b/docs/design/nd_array_approach_comparison.md deleted file mode 100644 index 04653bff..00000000 --- a/docs/design/nd_array_approach_comparison.md +++ /dev/null @@ -1,440 +0,0 @@ -# N-D Array Approach Comparison: unsafe_wrap vs ReshapedArray - -> **Update (v0.2.2, feat/new_array_nd)**: The N-way set-associative cache described in this -> document has been **superseded on Julia 1.11+ CPU** by `setfield!`-based wrapper reuse -> (`nd_wrappers` indexed by dimensionality N). This achieves **0-alloc for unlimited dimension -> patterns** — no eviction, no `CACHE_WAYS` limit. The N-way cache remains in use for -> **CUDA** and the **Julia 1.10 legacy** path. The `acquire!` → `ReshapedArray` path is -> unchanged. See `src/acquire.jl` and `src/types.jl` for the current implementation. - -## Summary - -This document analyzes two approaches for returning N-dimensional arrays from AdaptiveArrayPools: - -1. **Current (v1.1.x)**: `unsafe_wrap(Array, pointer, dims)` with N-way cache -2. **Proposed (v1.0.2 style)**: `reshape(view(backing, 1:n), dims)` without cache - -**Recommendation**: Switch back to ReshapedArray approach for `acquire!` N-D path. - ---- - -## Key Finding: SubArray Wrapper Allocation - -### The 48-byte Problem - -When using `acquire!` vs `unsafe_acquire!` in real code: - -```julia -# In flux_array! (TurbulentTransport) -xx = unsafe_acquire!(pool, T, size(x)) # 0 bytes - returns Array directly -xx = acquire!(pool, T, size(x)) # 48 bytes - SubArray wrapper! -``` - -**Root Cause**: `acquire!` returns `SubArray`, which allocates its wrapper struct (48 bytes) when it escapes the function scope. - -### Allocation Breakdown - -| API | Return Type | Allocation | -|-----|-------------|------------| -| `unsafe_acquire!(pool, T, m, n)` | `Matrix{T}` | 0 bytes (cache hit), 112 bytes (miss) | -| `acquire!(pool, T, m, n)` | `SubArray{..., Matrix}` | **48 bytes** (wrapper) + 112 bytes (miss) | - -### Why SubArray Allocates - -**Fundamental difference:** - -``` -unsafe_acquire! returns: -┌─────────────────────────────────────────────────────┐ -│ Pool backing Vector │ -│ [████████████████████████████████████] │ -│ ↓ │ -│ Array header lives in pool cache (reused) │ -│ → Returns pointer to EXISTING object (0 alloc) │ -└─────────────────────────────────────────────────────┘ - -acquire! returns: -┌─────────────────────────────────────────────────────┐ -│ Pool backing Vector │ -│ [████████████████████████████████████] │ -│ ↓ │ -│ Array in cache │ -│ ↓ │ -│ NEW SubArray struct (parent, indices, stride...) │ -│ → Creates NEW wrapper object (48 bytes!) │ -└─────────────────────────────────────────────────────┘ -``` - -SubArray is stack-allocated **only when**: -1. Used entirely within a single function -2. Compiler can prove it doesn't escape - -In `flux_array!`, `xx` escapes because: -- Passed to `_pooled_chain(out_y, xx)` -- Used across multiple loop iterations -- Compiler can't optimize away the wrapper - -### Root Cause: Type-Unspecified Path → Escape Optimization Failure - -**Critical Finding**: The core reason for allocation is that **compiler escape optimization fails in type-unspecified call paths**. - -> **Correction**: SubArray is a mutable struct (not immutable/isbits). -> The explanation that "dynamic dispatch causes boxing" is inaccurate. -> Precisely: wrapper objects are created at construction time, and the key factor is -> whether the compiler can optimize to stack allocation through escape analysis. - -#### Why it happens in `flux_array!`: - -```julia -# TGLFNNmodel struct (tglf_nn.jl) -struct TGLFNNmodel <: TGLFmodel - fluxmodel::Flux.Chain - # ... - _pooled_chain::PooledChain # ← No concrete type parameter! -end -``` - -The `_pooled_chain` field is declared without concrete type parameters, -so the call `fluxmodel._pooled_chain(out_y, xx)` is **not recompiled** and wrapper object optimization is not applied. - -#### Escape Optimization Failure: - -| Condition | Compiler Behavior | Result | -|-----------|-------------------|--------| -| Type specified + no escape | SROA/escape analysis applied | Stack allocation or elimination possible | -| Type specified + escape | Partial optimization possible | Depends on situation | -| **Type unspecified** | Optimization not applicable | **Wrapper object heap allocation** | - -``` -Type-specified path: -┌─────────────────────────────────────────────────────────────────┐ -│ Compiler knows the type │ -│ → SROA/escape analysis can be applied │ -│ → Wrapper object can be stack-allocated or completely removed │ -└─────────────────────────────────────────────────────────────────┘ - -Type-unspecified path (e.g., call through abstract field): -┌─────────────────────────────────────────────────────────────────┐ -│ Compiler doesn't know concrete type │ -│ → Escape analysis cannot be applied │ -│ → Wrapper object is heap-allocated │ -└─────────────────────────────────────────────────────────────────┘ -``` - -#### Array vs View Types in Type-Unspecified Paths: - -| Type | Characteristic | In Type-Unspecified Path | -|------|----------------|--------------------------| -| `Array` | Already heap-allocated object | Cached instance reuse → **No additional allocation** | -| `SubArray` | Requires wrapper object | Escape optimization failure → **Wrapper allocation** | -| `ReshapedArray` | Requires wrapper object | Escape optimization failure → **Wrapper allocation** | - -**Key insight**: Array is an object that already exists on the heap, so returning the same instance from cache incurs no additional allocation. -In contrast, SubArray/ReshapedArray create new wrapper objects each time, and optimization is not applied when type is unspecified. - -### Solutions for the Wrapper Allocation Problem - -#### 1. Use `unsafe_acquire!` (Recommended for this case) - -`unsafe_acquire!` returns `Array`, which is already a heap-allocated object (cached instance can be reused): - -```julia -# flux_array! in tglf_nn.jl -xx = unsafe_acquire!(pool, T, size(x)) # Returns Matrix{T} → cache hit = 0 alloc -``` - -✅ Zero allocation on cache hit (cached Array instance reused) -✅ No code changes to TGLFNNmodel needed -✅ Safe since `xx` is only used as scratch memory - -#### 2. Parameterize TGLFNNmodel (Fundamental fix) - -```julia -struct TGLFNNmodel{M<:Flux.Chain, P<:PooledChain} <: TGLFmodel - fluxmodel::M - _pooled_chain::P # Now compiler knows exact type -end -``` - -✅ Enables compiler escape optimization -✅ SubArray/ReshapedArray may become zero-alloc (depends on SROA/escape analysis) -❌ Requires significant code changes -❌ Changes serialization behavior - -#### 3. Function Barrier - -```julia -# Force type specialization -@inline function _call_pooled_chain(chain::PooledChain{M}, out, x) where M - chain(out, x) -end -``` - -⚠️ May not help if extracting `_pooled_chain` from struct is also dynamic - -### Implications for API Design - -**Key Insight: Wrapper Types Depend on Compiler Optimization** - -> **Note**: Whether wrapper types (SubArray, ReshapedArray) allocate depends on compiler SROA/escape analysis. -> The comparison below assumes **type-specified paths**. In type-unspecified paths, all wrapper types allocate. - -| Approach | Return Type | Creation Cost | Type-Specified Path | Type-Unspecified Path | -|----------|-------------|---------------|---------------------|----------------------| -| `unsafe_acquire!` | `Array` | 112 bytes (miss) | cache hit: 0 | cache hit: 0 | -| `acquire!` (current) | `SubArray{Array}` | 112 bytes (miss) | Optimizable | Wrapper allocation | -| **`acquire!` (reshape)** | `ReshapedArray{View}` | **0 bytes** | **Optimizable** | Wrapper allocation | - -**Advantages of ReshapedArray approach**: - -1. ✅ No creation cost (no unsafe_wrap call) -2. ✅ Compiler optimization possible in type-specified paths -3. ✅ BLAS compatible (StridedArray) -4. ✅ Same operation performance (mul!, broadcast) - -**`unsafe_acquire!` is better in type-unspecified paths**: -- Array already exists on heap → no additional allocation when reusing cached instance - ---- - -## Benchmark Results - -### Test Environment -- Benchmark file: `benchmark/nd_approach_comparison.jl` -- Tests 8-10 specifically compare N-way cache behavior - -### Allocation Comparison - -| Scenario | unsafe_wrap | reshape | Savings | -|----------|-------------|---------|---------| -| Single call (cache miss) | 112 bytes | **0 bytes** | 100% | -| 3-way cycling × 100 | 33,600 bytes | **0 bytes** | 100% | -| 5-way cycling × 100 | 56,000 bytes | **0 bytes** | 100% | -| With 4-way cache (3-way pattern) | 0 bytes | **0 bytes** | - | -| With 4-way cache (5-way pattern) | 56,000 bytes | **0 bytes** | 100% | - -### Performance Comparison - -| Operation | unsafe_wrap | reshape | Winner | -|-----------|-------------|---------|--------| -| mul! (BLAS) | 9.92 μs | 9.96 μs | Tie | -| Broadcast σ.(x) | 24.2 μs | 23.5 μs | Tie | -| Dense layer | 35.3 μs | 33.7 μs | Tie | -| 3-way cycling | 5.87 μs | **0.38 μs** | reshape (15x) | -| 5-way cycling | 10.38 μs | **0.59 μs** | reshape (18x) | - -### Type Information - -All three types are `StridedArray` (BLAS compatible): - -```julia -# unsafe_wrap -Matrix{Float64} # isa StridedArray ✓ - -# SubArray of unsafe_wrap -SubArray{Float64, 2, Matrix{Float64}, ...} # isa StridedArray ✓ - -# ReshapedArray of 1D view -Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, ...}, ...} # isa StridedArray ✓ -``` - -### Memory Layout - -All have identical column-major layout: -``` -Strides: (1, 4) # Same for all three -``` - ---- - -## Current Implementation (unsafe_wrap) - -### Call Path -``` -acquire!(pool, Float64, 64, 100) - └─> get_nd_view!(tp, (64, 100)) - └─> get_nd_array!(tp, (64, 100)) - ├─> get_view!(tp, 6400) # 1D view (0 alloc, cached) - └─> unsafe_wrap(...) # 112 bytes on cache miss! -``` - -### Code Location -`src/core.jl:129`: -```julia -arr = unsafe_wrap(Array{T, N}, pointer(flat_view), dims) -``` - -### N-way Cache Structure (Current Implementation) -```julia -# In TypedPool (src/types.jl) - used by unsafe_acquire! only -# Note: nd_views was removed since acquire! now uses reshape() -nd_arrays::Vector{Any} # Cached Array objects (from unsafe_wrap) -nd_dims::Vector{Any} # Cached dimension tuples -nd_ptrs::Vector{UInt} # Cached pointers for invalidation -nd_next_way::Vector{Int} # Round-robin counter per slot -``` - ---- - -## Proposed Implementation (ReshapedArray) - -### Call Path -``` -acquire!(pool, Float64, 64, 100) - └─> get_nd_view!(tp, (64, 100)) - └─> get_view!(tp, 6400) # 1D view (0 alloc, cached) - └─> reshape(view, dims) # 0 alloc always! -``` - -### Proposed Code Change -```julia -# Replace get_nd_view! in src/core.jl -@inline function get_nd_view!(tp::TypedPool{T}, dims::NTuple{N, Int}) where {T, N} - total_len = safe_prod(dims) - flat_view = get_view!(tp, total_len) # 1D view (cached) - return reshape(flat_view, dims) # Zero-alloc ReshapedArray -end -``` - ---- - -## Pros and Cons - -### ReshapedArray Approach (Proposed) - -#### Pros -1. **Zero allocation always** - No 112-byte allocation regardless of cache hit/miss -2. **No N-D cache needed** - Simpler code, less memory overhead -3. **Faster** - No cache lookup overhead (0.38μs vs 5.87μs for cycling patterns) -4. **BLAS compatible** - ReshapedArray is StridedArray -5. **Same performance** - Identical mul!/broadcast speed -6. **No Bélády's Anomaly** - Works with any access pattern (5-way, 10-way, etc.) -7. **Simpler TypedPool** - Can remove nd_arrays, nd_dims, nd_ptrs fields - -#### Cons -1. **Return type changes** - `SubArray{..., Array{...}}` → `ReshapedArray{..., SubArray{...}}` -2. **Some libraries might check `isa Array`** - Rare, but possible (not BLAS though) -3. **Slightly different printing** - Display shows as ReshapedArray - -### unsafe_wrap Approach (Current) - -#### Pros -1. **Returns actual Array** - Some code might expect `Matrix{Float64}` -2. **Cache hits are zero-alloc** - When pattern fits in N-way cache - -#### Cons -1. **112 bytes per cache miss** - Adds up with varying batch sizes -2. **N-way cache complexity** - Extra fields, cache lookup logic -3. **Bélády's Anomaly** - 5+ patterns = 100% miss with 4-way cache -4. **Slower cycling** - Cache lookup overhead even on hits - ---- - -## Impact Analysis - -### TurbulentTransport Usage - -In `src/pooled_layers.jl`: - -```julia -# PooledDense (line 86) -out = acquire!(pool, Float64, size(d.weight, 1), size(xT, 2)) -mul!(out, d.weight, xT) # Works with ReshapedArray ✓ - -# PooledActivation (line 54) -out = acquire!(pool, Float64, size(x)) -out .= pa.σ.(x) # Works with ReshapedArray ✓ -``` - -Both use cases are compatible with ReshapedArray: -- `mul!` accepts any `StridedMatrix` -- Broadcasting works on any `AbstractArray` - -### Flux.NNlib.bias_act! - -```julia -Flux.NNlib.bias_act!(d.σ, out, d.bias) -``` - -This function accepts `AbstractArray` - ReshapedArray is compatible. - -### unsafe_acquire! Unchanged - -For code that explicitly needs raw `Array` (FFI, specific BLAS paths): -```julia -unsafe_acquire!(pool, Float64, 64, 100) # Still returns Matrix{Float64} -``` - -This API remains unchanged and still uses `unsafe_wrap` with caching. - ---- - -## Migration Path - -### Phase 1: Modify acquire! N-D path -```julia -# src/core.jl - Replace get_nd_view! -@inline function get_nd_view!(tp::TypedPool{T}, dims::NTuple{N, Int}) where {T, N} - total_len = safe_prod(dims) - flat_view = get_view!(tp, total_len) - return reshape(flat_view, dims) -end -``` - -### Phase 2: Simplify TypedPool (optional) -Remove N-D cache fields if `unsafe_acquire!` usage is rare: -- `nd_views`, `nd_arrays`, `nd_dims`, `nd_ptrs` - -### Phase 3: Update documentation -- Note return type change in CHANGELOG -- Update docstrings for `acquire!` - ---- - -## Conclusion - -### Two Separate Problems, Two Solutions - -This investigation revealed **two distinct allocation issues**: - -#### Problem 1: N-D Array Creation (unsafe_wrap vs reshape) - -| Metric | unsafe_wrap + cache | reshape | -|--------|---------------------|---------| -| Allocation (miss) | 112 bytes | **0 bytes** | -| Allocation (hit) | 0 bytes | **0 bytes** | -| Speed (cycling) | 5-10 μs | **0.3-0.6 μs** | -| BLAS compat | ✓ | ✓ | -| Code complexity | High (cache) | **Low** | -| Works with any pattern | ✗ (≤4 ways) | **✓ (any)** | - -**Solution**: Switch `acquire!` N-D path to use `reshape(view, dims)` instead of `unsafe_wrap`. - -#### Problem 2: Type-Unspecified Path Wrapper Allocation - -> **Correction**: SubArray and ReshapedArray are both **mutable structs**. -> Whether allocation occurs depends on whether compiler escape optimization can be applied. - -| Type | Type-Specified Path | Type-Unspecified Path | -|------|---------------------|----------------------| -| `Array` | cache hit: 0 bytes | cache hit: **0 bytes** ✓ | -| `SubArray` | Optimizable | **Wrapper allocation** ✗ | -| `ReshapedArray` | Optimizable | **Wrapper allocation** ✗ | - -**Key difference**: -- `Array`: Object already exists on heap → no additional allocation when reusing cached instance -- Wrapper types: New object created each time → heap allocation when optimization fails in type-unspecified paths - -**Solution for TurbulentTransport**: Use `unsafe_acquire!` which returns `Array` (cached instance reusable). - -### Summary - -| Context | Recommended API | Reason | -|---------|-----------------|--------| -| Type-specified path (general use) | `acquire!` → ReshapedArray | 0 bytes creation, compiler optimization possible | -| Type-unspecified path | `unsafe_acquire!` → Array | Cached instance can be reused | -| FFI / raw pointer needs | `unsafe_acquire!` → Array | Direct memory access | - -**The N-way cache was solving the wrong problem** - caching `Array` objects when `reshape` is already zero-cost for creation. - -For type-unspecified paths (like `TGLFNNmodel._pooled_chain` without concrete type parameters), `unsafe_acquire!` returning `Array` is the correct choice because cached Array instances can be reused without additional allocation. diff --git a/docs/design/selective_memory_reclamation.md b/docs/design/selective_memory_reclamation.md deleted file mode 100644 index 4a4eba59..00000000 --- a/docs/design/selective_memory_reclamation.md +++ /dev/null @@ -1,417 +0,0 @@ -# Inactive Slot Trimming for AdaptiveArrayPools - -- **Date:** 2026-06-28 -- **Status:** Implemented (CPU + Metal validated on hardware; CUDA precompile-validated, GPU execution via CI). -- **Scope:** New manual `trim!` API. CPU + Metal + CUDA (CUDA mirrors Metal). -- **Author:** design session - ---- - -## 1. Motivation - -AdaptiveArrayPools keeps backing buffers after scopes exit so later `acquire!` -calls can reuse memory and stay allocation-free on the hot path. This is the -right default for normal workloads. - -The downside is memory retention after an unusual workload. A rare large input -can grow one or more slots far beyond the usual working set. After the scope -rewinds, those slots are no longer active, but the pool still keeps strong -references to their backing buffers and cached wrappers. Julia's GC cannot -collect those buffers while the pool still references them. - -We want a manual operation that says: - -> Keep the arrays currently in use. Drop the inactive pooled buffers that are -> only being retained for possible future reuse. - -This is not an immediate OS/VRAM release guarantee. The goal is to release the -pool-held references so the old buffers become GC-eligible once no user code -still references them. - -## 2. Core Concept: Active vs Inactive Slots - -Each typed pool has: - -- **active slots:** `1:tp.n_active` - - Currently claimed by the live call stack. - - Must remain valid. -- **inactive retained slots:** `(tp.n_active + 1):length(tp.vectors)` - - Already released by `rewind!` or previous scope exit. - - Kept only as a reuse cache. - - Safe for a manual cleanup operation to detach from the pool. - -`trim!` operates only on inactive retained slots. It never touches active slots. - -This makes the API easier to reason about than "deep" or "level" terminology: -the persistent thing is the typed pool's slot array, not a durable level object. - -## 3. Proposed API - -```julia -trim!(pool::AdaptiveArrayPool; force_gc::Bool = false) -trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) -trim!(; force_gc::Bool = false) -``` - -Metal extension: - -```julia -trim!(pool::MetalAdaptiveArrayPool; force_gc::Bool = false) -trim!(pool::MetalAdaptiveArrayPool, ::Type{T}; force_gc::Bool = false) -``` - -Disabled pools: - -```julia -trim!(::DisabledPool; force_gc::Bool = false) -trim!(::DisabledPool, ::Type{T}; force_gc::Bool = false) -``` - -Return a small summary: - -```julia -(; - slots_released::Int, - wrappers_released::Int, - estimated_bytes_released::Int, - gc_triggered::Bool, -) -``` - -`force_gc=false` is the default. When `true`, `trim!` calls `GC.gc()` after -detaching inactive references. This still does not promise immediate OS or GPU -memory return; it only asks Julia to run collection after the pool stops holding -the buffers. - -## 4. Semantics - -For each affected typed pool: - -1. Let `keep = tp.n_active`. -2. Estimate storage held by inactive slots for the return summary. -3. Drop backing-buffer references by truncating `tp.vectors` to `keep`. -4. Drop cached wrapper references for inactive slots by truncating each - wrapper-cache vector to `keep`. -5. Leave checkpoint stacks, depth state, typed-pool objects, and live slots - unchanged. - -Illustrative core: - -```julia -function _trim_inactive_typed_pool!(tp) - keep = tp.n_active - old_len = length(tp.vectors) - - released = old_len - keep - bytes = _inactive_storage_bytes(tp, keep + 1, old_len) - - resize!(tp.vectors, keep) - _trim_inactive_wrappers!(tp, keep) - - return (; slots_released = released, - wrappers_released = ..., - estimated_bytes_released = bytes) -end -``` - -No replacement vectors are allocated. If a later workload needs more slots, the -existing `acquire!` cold path will allocate the needed backing buffers at the -actual requested size. - -This avoids speculative "right-sizing" or reserve heuristics. A cleanup API -should not guess future sizes or allocate warm reserves on the user's behalf. - -## 5. Relationship to Existing Operations - -| Function | `n_active` | Backing buffers | Wrapper caches | Depth/checkpoint state | -|---|---:|---|---|---| -| `rewind!` | restored by one scope | keep all | keep all | pop one scope | -| `reset!` | set to 0 | keep all | keep all | reset to sentinel | -| `empty!` | set to 0 | drop all | drop all | reset to sentinel | -| `trim!` | unchanged | drop inactive only | drop inactive only | unchanged | - -Important distinction: - -- `trim!` is not `reset!`. If slots are still active, `trim!` preserves them. -- If user code has allowed `n_active` to grow at top level, call `reset!(pool)` - first if the intent is to mark those slots inactive, then call `trim!(pool)`. -- At normal top level after `@with_pool` scopes have rewound, `n_active` is - usually `0`, so `trim!(pool)` releases all retained backing buffers while - keeping the pool object and lightweight structural state. - -## 6. Why Not Allocate Smaller Replacement Buffers? - -An earlier idea was to replace a huge inactive buffer with a smaller new buffer -such as `1.5x` or `2x` the expected normal size. - -That is not part of this design. - -Reasons: - -- It performs a new allocation during a memory-reclamation call. -- The "normal future size" is only a guess without hot-path usage tracking. -- If the user actually needs that slot again, the existing `acquire!` path will - allocate the right size when it has the real request. -- The simplest and most honest behavior is: detach inactive references now, - allocate later only if needed. - -## 7. Implementation Notes - -### 7.1 Backing Buffers - -CPU typed pools can detach inactive backing buffers with: - -```julia -resize!(tp.vectors, tp.n_active) -``` - -This removes references to the inactive `Vector`s. The old vectors become -GC-eligible once no other references exist. - -For `BitTypedPool`, the same outer-vector truncation works; the dropped -`BitVector`s own their chunk arrays. - -For Metal, the same structural operation removes Julia references to inactive -`MtlArray`s. Actual VRAM return still depends on Julia GC and Metal.jl's memory -pool behavior, matching the caveat already documented for `empty!`. - -### 7.2 Wrapper Caches - -Dropping `tp.vectors` is not enough when cached wrappers still reference the old -backing storage. - -For each cached wrapper vector in `tp.arr_wrappers`, truncate it to `tp.n_active` -or clear inactive entries before truncation: - -```julia -for wrappers in tp.arr_wrappers - wrappers === nothing && continue - old_len = length(wrappers) - released += max(0, old_len - tp.n_active) - resize!(wrappers, min(length(wrappers), tp.n_active)) -end -``` - -The wrapper cache remains structurally reusable. Future `acquire!` calls past -the current active region will take the existing cold path and create fresh -wrappers if needed. - -### 7.3 Byte Accounting - -`estimated_bytes_released` is an estimate of storage that the pool stopped -retaining, not proof that RSS/VRAM immediately dropped. - -Suggested helpers: - -- CPU `Vector`: `Base.summarysize(v)` because `resize!(v, 0)` can leave capacity - retained even when `length(v) == 0`. -- `BitVector`: `Base.summarysize(v)` for the same capacity reason. -- Metal `MtlArray`: use GPU allocation size (`sizeof(v)` or `maxsize`-based - helper consistent with existing Metal stats), plus CPU wrapper size only if - useful for diagnostics. - -Avoid `length(v) * sizeof(T)` for released slots in safety mode: those vectors -may have logical length `0` while still retaining large capacity. - -## 8. Safety Contract - -`trim!` preserves active slots and their wrappers. Arrays currently in use by -the live call stack remain valid. - -Inactive retained slots are different: they have already been released by pool -scope management. User code should not rely on arrays/views from those slots. -After `trim!`, any pool-owned references to those inactive buffers are gone. - -In runtime-check mode (`S >= 1`), released slots have already been invalidated -by `rewind!`/`reset!` where applicable. `trim!` does not need to add new hot-path -safety metadata. - -`trim!` should be idempotent: - -- Calling it twice in a row releases nothing on the second call. -- Calling it inside a nested computation keeps currently active slots. -- Calling it at top level after normal rewind releases retained buffers. - -## 9. Backend Structure - -Generic core in `src/state.jl`: - -- `_trim_inactive_typed_pool!(tp::AbstractTypedPool)` -- `_trim_inactive_wrappers!(tp, keep::Int)` -- `_inactive_storage_bytes(tp, first::Int, last::Int)` - -Per-pool summaries are accumulated inline in the `trim!(pool)` loops over the -fixed slots and `others` (no separate merge helper is needed). - -CPU pool entry points: - -- `trim!(pool::AdaptiveArrayPool; force_gc=false)` -- `trim!(pool::AdaptiveArrayPool, ::Type{T}; force_gc=false)` -- `trim!(; force_gc=false)` for `get_task_local_pool()` - -Metal extension entry points: - -- Thin methods in `ext/AdaptiveArrayPoolsMetalExt/state.jl` -- Reuse generic typed-pool helpers where field layout matches. -- Add Metal-specific storage-byte helper if needed. - -CUDA: - -- Implemented as a structural mirror of Metal (`CuTypedPool` has the same field - layout; `_inactive_storage_bytes` overridden to use `sizeof` for GPU bytes). -- Validated locally by extension precompilation (method dispatch resolves); - GPU execution covered by `test/cuda/test_trim.jl`, which runs on CI with a - functional CUDA device (skipped on non-CUDA machines). - -Legacy (Julia < 1.12): - -- The legacy pool (`src/legacy/`) uses an N-way set-associative cache with a - different field layout (`nd_arrays`/`nd_dims`/… instead of `arr_wrappers`), so - real reclamation is not implemented there. -- `trim!` is still **defined and exported** on the legacy path as a NO-OP that - returns a zero summary and warns once (`@warn ... maxlog=1`). This keeps the - public API callable across the full `[compat] julia = "1.10"` range, so - dependent packages can `import`/call `trim!` on any supported Julia version - without a load-time or runtime `UndefVarError`. Real reclamation requires - Julia 1.12+. (Adding a real legacy implementation later would be non-breaking.) - -## 10. Phase Plan (TDD) - -Execute phases one at a time. - -### Phase 1 — CPU `trim!` - -**Goal:** Implement `trim!` for CPU pools and typed pools. - -**Test strategy:** - -- Active slots are preserved. -- Inactive `tp.vectors` entries are removed. -- Inactive wrapper-cache entries are removed. -- `reset!(pool); trim!(pool)` drops all retained CPU buffers. -- `trim!(pool, T)` affects only that type. -- `DisabledPool` returns a zero summary. -- `force_gc=true` sets `gc_triggered=true`. - -**Tasks:** - -- [ ] RED: Add CPU tests in `test/test_state.jl` or a new targeted test file. -- [ ] GREEN: Add generic typed-pool trimming helpers and CPU API methods. -- [ ] REFACTOR: Share summary aggregation helpers and docstrings. - -**Quality gate:** - -- [ ] Targeted CPU state tests pass. -- [ ] Existing allocation/zero-allocation tests for normal acquire paths still pass. -- [ ] `trim!` does not add code to hot acquire or automatic rewind paths. - -**Rollback:** Revert `src/state.jl`, export changes, and new tests. - -### Phase 2 — Safety Mode + Wrapper Reference Verification - -**Goal:** Prove trimming actually detaches pool-held references, including -cached wrappers. - -**Test strategy:** - -- Create cached `Array`/`BitArray` wrappers for inactive slots, call `trim!`, - and assert wrapper cache length is truncated. -- In `S=1`, released slots that were invalidated to length zero are still - counted using capacity-aware byte accounting. -- Subsequent `acquire!` after trim self-heals by allocating fresh backing storage - and wrapper entries. - -**Tasks:** - -- [ ] RED: Add tests for wrapper-cache truncation and S=1 byte estimates. -- [ ] GREEN: Fix byte accounting and wrapper-cache cleanup details. -- [ ] REFACTOR: Keep helper names and docstrings aligned with "inactive slots". - -**Quality gate:** - -- [ ] Targeted CPU safety tests pass. -- [ ] No stale wrapper retains old backing storage through the pool cache. - -**Rollback:** Revert helper changes and tests from this phase. - -### Phase 3 — Metal Parity - -**Goal:** Add the same inactive-slot trimming API to Metal pools. - -**Test strategy:** - -- Active Metal slots are preserved. -- Inactive Metal backing `MtlArray` references are removed. -- Inactive Metal wrapper-cache entries are removed. -- `force_gc=true` calls Julia GC after references are detached. -- Return summary uses Metal-appropriate storage estimates. - -**Tasks:** - -- [ ] RED: Add Metal trim tests under `test/metal/`. -- [ ] GREEN: Add thin Metal pool entry points and byte helper overrides. -- [ ] REFACTOR: Keep CPU/Metal implementation split consistent with `reset!` and `empty!`. - -**Quality gate:** - -- [ ] Targeted Metal tests pass on the M1 machine. -- [ ] CPU tests from Phases 1-2 remain green. - -**Rollback:** Revert Metal extension changes and Metal tests. - -### Phase 4 — Public Docs - -**Goal:** Document `trim!` as the manual inactive-slot cleanup operation. - -**Test strategy:** - -- Doctest or docs build if available. -- API reference includes `trim!`. -- User docs clearly distinguish `trim!`, `reset!`, and `empty!`. - -**Tasks:** - -- [ ] RED: Add docs references that fail until API exists. -- [ ] GREEN: Update API docs, examples, and exported function list. -- [ ] REFACTOR: Tighten wording around GC eligibility vs immediate memory return. - -**Quality gate:** - -- [ ] Docs build succeeds. -- [ ] No public docs imply immediate OS/VRAM release. - -**Rollback:** Revert docs and export/reference updates. - -## 11. Deferred Ideas - -These are intentionally out of scope for the first implementation: - -- Outlier-only trimming (`policy=:largest`, `min_bytes`, `ratio`). -- Dry-run reporting. -- Per-type or per-backend default thresholds. -- Hysteresis or usage tracking. -- Automatic periodic trimming. -- Replacement/reserve buffers for "right-sized" warm reuse. - -All of these add policy. The first version should be a simple primitive: - -> Drop inactive retained references; preserve active working set. - -## 12. Risks - -| Risk | Impact | Likelihood | Mitigation | Detection | -|---|---:|---:|---|---| -| Cached wrappers keep old buffers alive | High | Medium | Truncate wrapper caches with `tp.vectors` | Wrapper-cache reference tests | -| Byte summary reports zero for resized inactive buffers | Medium | High in `S=1` | Use capacity-aware estimates (`summarysize`, Metal allocation size) | S=1 byte-accounting tests | -| Users expect immediate RSS/VRAM drop | Medium | Medium | Document GC eligibility and `force_gc` limits | Docs review | -| `trim!` confused with `reset!` | Medium | Medium | Keep `n_active` unchanged and document `reset!; trim!` pattern | API docs/tests | -| Metal memory pool behavior differs from CPU | Medium | Medium | Match `empty!` caveats; test reference removal, not immediate VRAM | Metal tests | - -## 13. Open Items - -- Final summary field names: `slots_released` vs `inactive_slots_released`, - `estimated_bytes_released` vs `bytes_detached`. -- Whether zero-argument `trim!()` should stay CPU-only or eventually accept a - backend selector. -- Whether `force_gc=true` should call `GC.gc()` once or use a stronger collection - pattern. Default remains `false`. diff --git a/docs/design/untracked_acquire_design.md b/docs/design/untracked_acquire_design.md deleted file mode 100644 index c1568fbd..00000000 --- a/docs/design/untracked_acquire_design.md +++ /dev/null @@ -1,614 +0,0 @@ -# AdaptiveArrayPools: Untracked Acquire Handling Design - -## 1. Background - -### 1.1 Current System -Julia array pool library. The `@with_pool` macro manages temporary array allocation through checkpoint/rewind. - -```julia -@with_pool p begin - v = acquire!(p, Float64, 100) # Acquire array from pool - # ... use ... -end # Automatically returned -``` - -### 1.2 Optimization: Typed Checkpoint/Rewind -The macro extracts types from acquire! calls at compile time, checkpointing/rewinding only those types (77% performance improvement). - -```julia -# Code generated by macro -checkpoint!(pool, Float64) # Only checkpoint Float64 stack -try - v = acquire!(pool, Float64, 100) -finally - rewind!(pool, Float64) # Only rewind Float64 stack -end -``` - -### 1.3 Current Data Structures -```julia -mutable struct TypedPool{T} - vectors::Vector{Vector{T}} # Reusable arrays - n_active::Int # Number of currently active arrays - _checkpoint_n_active::Vector{Int} # Checkpointed n_active values - _checkpoint_depths::Vector{Int} # Depth at which each checkpoint was saved -end - -mutable struct AdaptiveArrayPool - float64::TypedPool{Float64} - float32::TypedPool{Float32} - int64::TypedPool{Int64} - int32::TypedPool{Int32} - complexf64::TypedPool{ComplexF64} - complexf32::TypedPool{ComplexF32} - bool::TypedPool{Bool} - others::IdDict{DataType, Any} - _current_depth::Int - _untracked_flags::Vector{Bool} -end -``` - ---- - -## 2. Problem: Untracked Acquire - -### 2.1 Scenario -```julia -function inner_helper(pool) - acquire!(pool, Float32, 50) # Macro can't see this! -end - -@with_pool p begin - v = acquire!(p, Float64, 100) # Macro extracts: Float64 - inner_helper(p) # Float32 acquire is "untracked" -end -``` - -**Problem:** Macro only knows about Float64, generating `checkpoint!(pool, Float64)`, `rewind!(pool, Float64)`. -Float32's n_active is not restored → memory leak or corruption. - -### 2.2 Nested Case -```julia -@with_pool p begin # L1: checkpoint Float64 - acquire!(p, Float64, 10) - - @with_pool p begin # L2: checkpoint Int64 - acquire!(p, Int64, 5) - inner_untracked(p) # Float32 acquire (untracked) - end # L2 rewind - -end # L1 rewind -``` - ---- - -## 3. Proposed Solution - -### 3.1 Core Idea: Shared Implementation Function + Marking Wrapper - -**Previous Problem:** Runtime detection in acquire! → overhead on every call - -**New Approach:** -- `_acquire_impl!` → actual implementation (existing acquire! logic as-is) -- `acquire!` → marking + `_acquire_impl!` call (1-line wrapper) -- Macro transforms `acquire!` calls to `_acquire_impl!` - -```julia -# User code -@with_pool p begin - v = acquire!(p, Float64, 100) -end - -# After macro transformation -checkpoint!(p, Float64) -try - v = _acquire_impl!(p, Float64, 100) # ← Transformed! (no marking) -finally - # ... -end -``` - -### 3.2 Benefits -1. **Zero overhead:** Inside @with_pool, impl is called directly -2. **Automatic detection:** acquire! outside @with_pool passes through marking wrapper → untracked -3. **No code duplication:** Implementation logic in `_acquire_impl!` only once - -### 3.3 Data Structure Changes - -```julia -mutable struct TypedPool{T} - # ... existing fields ... - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} # NEW: Tracks which depth each checkpoint was saved at -end - -mutable struct AdaptiveArrayPool - # ... existing TypedPools ... - - # NEW: Depth tracking (pool level) - _current_depth::Int # Current checkpoint depth - _untracked_flags::Vector{Bool} # Whether untracked occurred per depth (typed vs full rewind decision) -end -``` - -**Key Changes:** -- `_checkpoint_depths` added to `TypedPool` → tracks which depth each checkpoint occurred at -- `_untracked_flags` maintained → for typed rewind vs full rewind decision -- `_full_rewind_with_types!` removed → `_checkpoint_depths` makes regular `rewind!(pool)` sufficient - ---- - -## 4. Implementation Details - -### 4.1 Function Structure - -```julia -# Implementation function (existing acquire! logic moved as-is) -@inline function _acquire_impl!(pool::AdaptiveArrayPool, ::Type{T}, n::Int) where T - tp = get_typed_pool!(pool, T) - return get_view!(tp, n) -end - -@inline function _acquire_impl!(pool::AdaptiveArrayPool, ::Type{T}, dims::Vararg{Int,N}) where {T,N} - tp = get_typed_pool!(pool, T) - return get_nd_view!(tp, dims) -end - -# User-facing: marking wrapper (1 line) -@inline function acquire!(pool::AdaptiveArrayPool, ::Type{T}, n::Int) where T - _mark_untracked!(pool) - _acquire_impl!(pool, T, n) -end - -@inline function acquire!(pool::AdaptiveArrayPool, ::Type{T}, dims::Vararg{Int,N}) where {T,N} - _mark_untracked!(pool) - _acquire_impl!(pool, T, dims...) -end - -# Untracked marking (still needed - for typed rewind vs full rewind decision) -# Note: 1-based sentinel pattern guarantees _current_depth >= 1, so no check needed -@inline function _mark_untracked!(pool::AbstractArrayPool) - # Always mark (_current_depth >= 1 guaranteed by sentinel) - @inbounds pool._untracked_flags[pool._current_depth] = true -end -``` - -**Macro transformation:** `acquire!(p, T, n)` → `_acquire_impl!(p, T, n)` (only function name substitution) - -**Note:** `_mark_untracked!` is maintained - for triggering full rewind when untracked occurs - -### 4.2 unsafe_acquire! and Aliases - -```julia -# Implementation functions -@inline function _unsafe_acquire_impl!(pool::AdaptiveArrayPool, ::Type{T}, n::Int) where T - tp = get_typed_pool!(pool, T) - return get_nd_array!(tp, (n,)) -end - -@inline function _unsafe_acquire_impl!(pool::AdaptiveArrayPool, ::Type{T}, dims::Vararg{Int,N}) where {T,N} - tp = get_typed_pool!(pool, T) - return get_nd_array!(tp, dims) -end - -# User-facing: marking wrapper -@inline function unsafe_acquire!(pool::AdaptiveArrayPool, ::Type{T}, n::Int) where T - _mark_untracked!(pool) - _unsafe_acquire_impl!(pool, T, n) -end - -@inline function unsafe_acquire!(pool::AdaptiveArrayPool, ::Type{T}, dims::Vararg{Int,N}) where {T,N} - _mark_untracked!(pool) - _unsafe_acquire_impl!(pool, T, dims...) -end - -# Aliases (user API) -const acquire_view! = acquire! -const acquire_array! = unsafe_acquire! - -# Aliases (internal impl) -const _acquire_view_impl! = _acquire_impl! -const _acquire_array_impl! = _unsafe_acquire_impl! -``` - -### 4.3 Macro Transformation - -```julia -# Macro transforms acquire! calls in AST to _*_impl! -function _transform_acquire_calls(expr, pool_name) - if expr isa Expr - if expr.head == :call && length(expr.args) >= 2 - fn = expr.args[1] - pool_arg = expr.args[2] - - # Only transform when matching target pool - if pool_arg == pool_name - if fn == :acquire! || fn == :acquire_view! - expr.args[1] = :_acquire_impl! - elseif fn == :unsafe_acquire! || fn == :acquire_array! - expr.args[1] = :_unsafe_acquire_impl! - end - end - end - # Recursively transform - for i in eachindex(expr.args) - expr.args[i] = _transform_acquire_calls(expr.args[i], pool_name) - end - end - return expr -end -``` - -**Transformation Rules:** -| Original | Transformed | -|----------|-------------| -| `acquire!(p, ...)` | `_acquire_impl!(p, ...)` | -| `acquire_view!(p, ...)` | `_acquire_impl!(p, ...)` | -| `unsafe_acquire!(p, ...)` | `_unsafe_acquire_impl!(p, ...)` | -| `acquire_array!(p, ...)` | `_unsafe_acquire_impl!(p, ...)`| - -### 4.4 Checkpoint/Rewind - -```julia -# Typed checkpoint -function checkpoint!(pool::AdaptiveArrayPool, types::Type...) - pool._current_depth += 1 - push!(pool._untracked_flags, false) - depth = pool._current_depth - for T in types - tp = get_typed_pool!(pool, T) - push!(tp._checkpoint_n_active, tp.n_active) - push!(tp._checkpoint_depths, depth) # NEW: record depth - end -end - -# Typed rewind (fast path - when no untracked) -function rewind!(pool::AdaptiveArrayPool, types::Type...) - for T in types - tp = get_typed_pool!(pool, T) - pop!(tp._checkpoint_depths) # NEW: remove depth - tp.n_active = pop!(tp._checkpoint_n_active) - end - pop!(pool._untracked_flags) - pool._current_depth -= 1 -end - -# Full rewind (untracked fallback - simplified with _checkpoint_depths!) -# Note: Uses 1-based sentinel pattern. At global scope (depth=1), delegates to reset!() -function rewind!(pool::AdaptiveArrayPool) - cur_depth = pool._current_depth - - # Safety guard: at global scope (depth=1), no checkpoint to rewind to - # Delegate to reset! which safely clears all n_active counters - if cur_depth == 1 - reset!(pool) - return nothing - end - - # Fixed slots - zero allocation via @generated iteration - foreach_fixed_slot(pool) do tp - _rewind_typed_pool!(tp, cur_depth) - end - - # Process fallback types - for tp in values(pool.others) - _rewind_typed_pool!(tp, cur_depth) - end - - pop!(pool._untracked_flags) - pool._current_depth -= 1 - return nothing -end - -# Internal helper for rewind with orphan cleanup -# Uses 1-based sentinel pattern: no isempty checks needed -@inline function _rewind_typed_pool!(tp::AbstractTypedPool, current_depth::Int) - # 1. Orphaned Checkpoints Cleanup - while @inbounds tp._checkpoint_depths[end] > current_depth - pop!(tp._checkpoint_depths) - pop!(tp._checkpoint_n_active) - end - - # 2. Normal Rewind Logic (Sentinel Pattern) - if @inbounds tp._checkpoint_depths[end] == current_depth - pop!(tp._checkpoint_depths) - tp.n_active = pop!(tp._checkpoint_n_active) - else - # No checkpoint at current depth - restore from parent - tp.n_active = @inbounds tp._checkpoint_n_active[end] - end - nothing -end -``` - -**Key Improvement:** `_checkpoint_depths[end] == depth` comparison enables accurate pop/restore decision → `_full_rewind_with_types!` not needed! - -### 4.5 Zero-Allocation Iteration (Current Implementation) - -> **Note**: The design originally proposed `all_type_stacks()` generator, but was replaced with -> `foreach_fixed_slot()` @generated function for zero allocation via compile-time unrolling. - -```julia -# Current implementation uses @generated for zero allocation -const FIXED_SLOT_FIELDS = (:float64, :float32, :int64, :int32, :complexf64, :complexf32, :bool) - -@generated function foreach_fixed_slot(f::F, pool::AdaptiveArrayPool) where {F} - exprs = [:(f(getfield(pool, $(QuoteNode(field))))) for field in FIXED_SLOT_FIELDS] - quote - Base.@_inline_meta - $(exprs...) - nothing - end -end -``` - -**Benefits over generator approach**: -- Zero allocation via compile-time unrolling -- Full inlining for hot paths -- No runtime iteration overhead - -### 4.6 Macro Generated Code - -```julia -@with_pool p begin - v = acquire!(p, Float64, 100) - inner_function(p) # untracked acquire possible here -end - -# Generated code: -local p = get_task_local_pool() -checkpoint!(p, Float64) -try - local _result = begin - v = _acquire_impl!(p, Float64, 100) # ← Transformed! (no marking) - inner_function(p) # acquire! inside not transformed → passes through marking wrapper → untracked - end - _result -finally - if p._untracked_flags[p._current_depth] - rewind!(p) # Full rewind (handled accurately with _checkpoint_depths) - else - rewind!(p, Float64) # Typed rewind (fast) - end -end -``` - -**Change:** `_full_rewind_with_types!(p, Float64)` → `rewind!(p)` (simplified!) - ---- - -## 5. Behavior Verification - -### 5.1 Simple Case (no untracked) - -```julia -@with_pool p begin - v = acquire!(p, Float64, 100) -end -``` - -**Flow:** -1. `checkpoint!(p, Float64)` → _current_depth=1, _untracked_flags=[false] -2. `_acquire_impl!(p, Float64, 100)` → no marking (macro transformed) -3. `_untracked_flags[1] == false` → `rewind!(p, Float64)` (fast path) - -### 5.2 Untracked Case - -```julia -function helper(p) - acquire!(p, Float32, 50) # Defined outside @with_pool → not transformed! -end - -@with_pool p begin - v = acquire!(p, Float64, 100) - helper(p) -end -``` - -**Flow:** -1. `checkpoint!(p, Float64)` → _current_depth=1, _untracked_flags=[false], Float64._checkpoint_depths=[1] -2. `_acquire_impl!(p, Float64, 100)` → no marking (macro transformed) -3. `acquire!(p, Float32, 50)` → marking wrapper → `_mark_untracked!` → _untracked_flags=[true] -4. `_untracked_flags[1] == true` → `rewind!(p)` (full rewind) - - Float64: _checkpoint_depths[end]=1 == depth=1 → pop - - Float32: _checkpoint_depths=[] → if n_active=0 do nothing, if >0 restore - -### 5.3 Nested Case - -``` -L1: @with_pool (Float64) - L2: @with_pool (Int64) + untracked Float32 - L3: @with_pool (Bool) -``` - -| Timing | depth | Float64 saved/depths | Int64 saved/depths | Float32 saved/depths | Bool saved/depths | -|--------|-------|---------------------|-------------------|---------------------|------------------| -| Initial | 0 | []/[] | []/[] | []/[] | []/[] | -| L1 checkpoint(F64) | 1 | [0]/[1] | []/[] | []/[] | []/[] | -| L2 checkpoint(I64) | 2 | [0]/[1] | [0]/[2] | []/[] | []/[] | -| Float32 untracked | 2 | [0]/[1] | [0]/[2] | []/[] (n=1) | []/[] | -| L3 checkpoint(Bool) | 3 | [0]/[1] | [0]/[2] | []/[] | [0]/[3] | -| L3 rewind(Bool) | 2 | [0]/[1] | [0]/[2] | []/[] | []/[] | -| L2 rewind() full | 1 | [0]/[1] | []/[] | []/[] (n=0) | []/[] | -| L1 rewind(F64) | 0 | []/[] | []/[] | []/[] | []/[] | - -**L2 Full Rewind Detail (depth=2):** -- Float64: depths[end]=1 ≠ 2 → no pop, restore only ✓ -- Int64: depths[end]=2 == 2 → **pop!** ✓ (now works correctly!) -- Float32: depths=[] → n_active to 0 (untracked cleanup) ✓ -- Bool: depths=[] → n_active=0, do nothing ✓ - -**Key:** Thanks to `_checkpoint_depths`, we can accurately know Int64 was checkpointed at depth=2! - ---- - -## 6. API Summary - -### 6.1 User API (unchanged) -```julia -acquire!(pool, T, dims...) # Acquire array (returns view) -unsafe_acquire!(pool, T, dims...) # Acquire array (returns Array) -acquire_view!(pool, T, dims...) # alias for acquire! -acquire_array!(pool, T, dims...) # alias for unsafe_acquire! -``` - -### 6.2 Internal API (newly added) -```julia -# Implementation functions (macro transforms to direct call) -_acquire_impl!(pool, T, dims...) -_unsafe_acquire_impl!(pool, T, dims...) - -# Aliases (const) -_acquire_view_impl! = _acquire_impl! -_acquire_array_impl! = _unsafe_acquire_impl! - -# Untracked marking -_mark_untracked!(pool) - -# TypedPool iteration -all_type_stacks(pool) # Generator -foreach_type_stack(f, pool) # Callback pattern -``` - -### 6.3 Struct Changes -```julia -mutable struct TypedPool{T} - # existing fields... - n_active::Int - _checkpoint_n_active::Vector{Int} - _checkpoint_depths::Vector{Int} # NEW: checkpoint depth tracking -end - -mutable struct AdaptiveArrayPool - # existing fields... - _current_depth::Int # NEW - _untracked_flags::Vector{Bool} # NEW -end -``` - ---- - -## 7. Performance Characteristics - -| Path | Condition | Behavior | Overhead | -|------|-----------|----------|----------| -| Fast path | no untracked | typed rewind | Minimal (same as existing) | -| Fallback | untracked exists | full rewind | Iterate all types (accurate handling with _checkpoint_depths) | - -### 7.1 acquire! Overhead -```julia -# Implementation function (existing logic as-is) -@inline function _acquire_impl!(pool, T, n) - tp = get_typed_pool!(pool, T) - get_view!(tp, n) -end - -# User wrapper (marking added) -@inline function acquire!(pool, T, n) - _mark_untracked!(pool) # Added: if + array access - _acquire_impl!(pool, T, n) -end - -@inline function _mark_untracked!(pool) - if pool._current_depth > 0 # branch (usually false) - @inbounds pool._untracked_flags[pool._current_depth] = true - end -end -``` - -**Overhead Analysis:** -- Outside @with_pool: `_current_depth == 0` → branch only (very light) -- Inside @with_pool: macro calls `_acquire_impl!` directly → **zero overhead** - ---- - -## 8. Edge Case: Parent Scope Corruption Prevention - -### 8.1 Problem Scenario -```julia -# acquire outside @with_pool -v_parent = acquire!(p, Int64, 10) # Int64 n_active = 1 - -@with_pool p begin # checkpoint(Float64) - Int64 not checkpointed - v = acquire!(p, Float64, 100) - - # Untracked Int64 acquire! - v_child = acquire!(p, Int64, 5) # Int64 n_active = 2 -end -# Full rewind occurs -# Int64: _checkpoint_n_active=[], n_active=2 -# If reset n_active=0, v_parent would be destroyed! -``` - -### 8.2 Solution: Clear Error -```julia -elseif tp.n_active > 0 - # _checkpoint_n_active empty but n_active > 0 - # = attempting rewind without checkpoint - T = eltype(tp) - error(""" - [AdaptiveArrayPools] Cannot rewind type $T: no checkpoint exists. - Found $(tp.n_active) active array(s) that were never checkpointed. - - Fix: Wrap the scope where $T was first acquired in @with_pool. - """) -end -``` - -### 8.3 Error Conditions Summary -| _checkpoint_n_active | n_active | Situation | Action | -|----------------------|----------|-----------|--------| -| len == depth | any | Checkpointed at current depth | pop | -| len > 0, < depth | any | Checkpointed at previous depth | restore to [end] | -| empty | 0 | Not used | do nothing | -| empty | > 0 | **Rewind attempted without checkpoint** | **ERROR** | - -### 8.4 User Resolution -When error occurs, two choices: -1. **Wrap parent in @with_pool too:** - ```julia - @with_pool p begin - v_parent = acquire!(p, Int64, 10) - @with_pool p begin - # inner block - end - end - ``` - -2. **Avoid untracked acquire in inner block:** - ```julia - v_parent = acquire!(p, Int64, 10) - @with_pool p begin - # Don't acquire Int64 - # Or use @with_pool in inner helper too - end - ``` - ---- - -## 9. Feedback Requests - -1. **Function separation approach:** Is the `acquire!` wrapper + `_acquire_impl!` implementation separation appropriate? -2. **length vs _current_depth comparison:** Is it correct in all nested cases? -3. **Macro transformation:** Is AST function name substitution safe? (qualified names, macros, etc.) -4. **_untracked_flags Vector{Bool}:** Are there more efficient alternatives? -5. **_full_rewind_with_types!:** Performance of checking types with Set? -6. **Simpler alternative:** Is there a way to simplify the entire design? - ---- - -## 10. Alternative Review - -### 10.1 Option: Always Full Checkpoint/Rewind -- Pros: Simple implementation -- Cons: Loses 77% performance improvement - -### 10.2 Option: Don't Support Untracked -- Documentation: "Use @with_pool in helper functions too" -- Pros: No implementation changes -- Cons: Silent corruption on user error - -### 10.3 Option: Current Proposal (impl function separation) -- Pros: Maintains performance, safe fallback, no code duplication -- Cons: Requires macro transformation diff --git a/docs/design/vector_resize_memory_behavior.md b/docs/design/vector_resize_memory_behavior.md deleted file mode 100644 index e0194a45..00000000 --- a/docs/design/vector_resize_memory_behavior.md +++ /dev/null @@ -1,247 +0,0 @@ -# Vector Resize Memory Behavior: CPU vs GPU - -## Context -AdaptiveArrayPools uses backing vectors that may need to grow when larger arrays are requested. -Current implementation only grows vectors, never shrinks them. - -**Question**: Should we shrink vectors when smaller sizes are requested? What are the memory implications? - ---- - -## CPU Julia Vector Behavior - -```julia -v = Vector{Float64}(undef, 1000) -resize!(v, 100) # Shrink to 100 elements -resize!(v, 500) # Grow back to 500 -``` - -### Key Facts (needs verification): -1. **Capacity vs Length**: Does Julia Vector maintain separate capacity? -2. **Shrink behavior**: Does `resize!(v, smaller)` release memory immediately? -3. **Regrow cost**: If we shrink then grow again, is there reallocation? - -### My Understanding: -- Julia's `Vector` uses a growth strategy (typically 2x) -- `resize!` to smaller size may NOT release memory (keeps capacity) -- Growing back within capacity is O(1), no allocation -- Memory is only released when Vector is GC'd - -**Question for review**: Is this accurate? Does Julia guarantee capacity preservation on shrink? - ---- - -## GPU CuVector Behavior - -```julia -using CUDA -v = CUDA.zeros(Float64, 1000) -resize!(v, 100) # Shrink - what happens to GPU memory? -resize!(v, 500) # Grow back - allocation? -``` - -### VERIFIED: CUDA.jl resize! Implementation (src/array.jl:889) - -**CuVector has capacity tracking via `A.maxsize` field.** - -```julia -# CUDA.jl constants -const RESIZE_THRESHOLD = 100 * 1024^2 # 100 MiB -const RESIZE_INCREMENT = 32 * 1024^2 # 32 MiB - -function Base.resize!(A::CuVector{T}, n::Integer) where T - n == length(A) && return A - - # only resize when the new length exceeds the capacity or is much smaller - cap = A.maxsize ÷ aligned_sizeof(T) - if n > cap || n < cap ÷ 4 # ← SHRINK THRESHOLD: 25% - len = if n < cap - # shrink to fit (allocates EXACT new size, no over-allocation) - n - elseif A.maxsize > RESIZE_THRESHOLD - # large arrays (>100MB): grow by fixed +32 MiB increments - max(n, cap + RESIZE_INCREMENT ÷ aligned_sizeof(T)) - else - # small arrays (<100MB): double in size - max(n, 2 * length(A)) - end - # ... allocates new buffer, copies data ... - end - # If within capacity: just update length, no reallocation -end -``` - -### Key Findings: - -| Aspect | CUDA.jl CuVector | -|--------|------------------| -| **Capacity tracking** | Yes, via `A.maxsize` | -| **Shrink threshold** | `n < cap ÷ 4` (25%) | -| **Shrink behavior** | Reallocates to EXACT new size | -| **Growth (small <100MB)** | 2x doubling | -| **Growth (large ≥100MB)** | +32 MiB increments | - -### CUDA.jl Memory Management: -- CUDA.jl uses a memory pool (stream-ordered or binned allocator) -- Released memory goes back to pool, not immediately to OS/driver -- `CUDA.reclaim()` forces return to driver -- Pool may return same block on regrow (observed in verification tests) - ---- - -## Current Pool Design Trade-offs - -### Current Approach: Never Shrink -```julia -# In get_view!: -if length(vec) < total_len - resize!(vec, total_len) # Only grow, never shrink -end -new_view = view(vec, 1:total_len) # View handles size -``` - -**Pros**: -- Simple implementation -- Avoids any potential reallocation costs -- Views already handle returning correct size - -**Cons**: -- One large allocation permanently increases memory footprint -- GPU memory is precious and limited -- No way to recover memory without `empty!(pool)` - -### Alternative: Shrink When Significantly Smaller -```julia -if length(vec) < total_len - resize!(vec, total_len) -elseif length(vec) > total_len * 4 # Example: 4x threshold - resize!(vec, total_len) # Shrink to save memory -end -``` - -**Pros**: -- Recovers memory from outlier large allocations -- Better memory efficiency over time - -**Cons**: -- May cause reallocations -- Added complexity -- Need to invalidate cached views on shrink too - ---- - -## Specific Questions for Review - -1. **Julia Vector capacity**: - - Does `resize!(v, smaller)` preserve capacity? - - Is this behavior documented/guaranteed? - - Is there a way to query capacity vs length? - -2. **CuVector resize behavior**: - - Does CUDA.jl's CuVector follow same capacity model? - - What happens to GPU memory on shrink? - - Does CUDA memory pool make shrink "free" anyway? - -3. **Design recommendation**: - - Should pools shrink vectors at some threshold? - - What threshold makes sense? (2x? 4x? 10x?) - - Should CPU and GPU have different policies? - -4. **Memory pressure handling**: - - Should pool respond to memory pressure signals? - - Is there a way to detect "memory is tight"? - ---- - -## Test Code to Verify Behavior - -```julia -# CPU Test -function test_cpu_resize_behavior() - v = Vector{Float64}(undef, 10_000_000) # ~80MB - @show Base.summarysize(v) - - resize!(v, 100) - @show Base.summarysize(v) # Does this shrink? - - resize!(v, 5_000_000) - @show Base.summarysize(v) # Reallocation needed? - - # Is there a way to check capacity? -end - -# GPU Test -function test_gpu_resize_behavior() - CUDA.reclaim() # Start clean - - v = CUDA.zeros(Float64, 10_000_000) # ~80MB GPU - @show CUDA.memory_status() - - resize!(v, 100) - @show CUDA.memory_status() # Memory returned to pool? - - resize!(v, 5_000_000) - @show CUDA.memory_status() # New allocation? -end -``` - ---- - -## Related: View Cache Invalidation - -Currently, when `resize!` is called (grow only), we invalidate all cached views: - -```julia -if length(vec) < total_len - resize!(vec, total_len) - # Invalidate all N-way cache entries for this slot - for k in 1:CUDA_CACHE_WAYS - @inbounds tp.views[base + k] = nothing - @inbounds tp.view_dims[base + k] = nothing - end -end -``` - -If we add shrinking, same invalidation would be needed since shrink can also reallocate. - ---- - -## Summary - -### VERIFIED Results - -| Aspect | CPU Vector | GPU CuVector | -|--------|------------|--------------| -| **Capacity tracking** | Yes (implicit) | Yes (`A.maxsize`) | -| **Capacity preservation on shrink** | Yes (pointer unchanged) | No (reallocates at 25%) | -| **Memory returned on shrink** | No (until GC) | To pool (can be reclaimed) | -| **Regrow cost after shrink** | O(1) within capacity | May realloc (pool often returns same block) | -| **CUDA.jl shrink threshold** | N/A | `n < cap ÷ 4` (25%) | - -### Design Recommendation for AdaptiveArrayPools - -**Current "never shrink" is suboptimal for GPU.** CUDA.jl already implements a 25% threshold, meaning: - -1. **Our explicit `resize!(vec, smaller)` calls would trigger CUDA.jl's internal shrink anyway** if below 25% -2. **We're just deferring the inevitable reallocation** when usage drops significantly -3. **GPU memory is precious** - holding 4x+ more than needed is wasteful - -**Recommendation**: Add lazy shrink for GPU at 25% threshold (matching CUDA.jl): - -```julia -# In get_view! for CuTypedPool: -cap = length(vec) -if total_len > cap - resize!(vec, total_len) # Grow - # invalidate cache... -elseif total_len < cap ÷ 4 - resize!(vec, total_len) # Shrink when using <25% capacity - # invalidate cache... -end -``` - -**Why 25%?** -- Matches CUDA.jl's internal threshold -- Consistent behavior - calling resize! directly would shrink at same point -- Allows 4x variation without reallocation (handles typical size fluctuations) -- Recovers memory from outlier large allocations From 0fc0c7cf2b87c43eabd5d8b8fe1e923b73796287 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 14:30:57 -0700 Subject: [PATCH 02/10] (feat): add touched-others stack fields to AdaptiveArrayPool --- src/state.jl | 10 ++++++++++ src/types.jl | 12 ++++++++++++ test/runtests.jl | 1 + test/test_touched_others.jl | 39 +++++++++++++++++++++++++++++++++++++ 4 files changed, 62 insertions(+) create mode 100644 test/test_touched_others.jl diff --git a/src/state.jl b/src/state.jl index 5ddf3866..e5333b08 100644 --- a/src/state.jl +++ b/src/state.jl @@ -555,6 +555,11 @@ function Base.empty!(pool::AdaptiveArrayPool) empty!(pool._others_ptr_bounds_checkpoints) push!(pool._others_ptr_bounds_checkpoints, 0) # Sentinel + # Reset touched-others tracking (transient scope state) + empty!(pool._touched_others) + empty!(pool._touched_others_checkpoints) + push!(pool._touched_others_checkpoints, 0) # Sentinel + # Reset type touch tracking state (1-based sentinel pattern) pool._current_depth = 1 # 1 = global scope (sentinel) empty!(pool._touched_type_masks) @@ -645,6 +650,11 @@ function reset!(pool::AdaptiveArrayPool{S}) where {S} empty!(pool._others_ptr_bounds_checkpoints) push!(pool._others_ptr_bounds_checkpoints, 0) # Sentinel + # Reset touched-others tracking (transient scope state) + empty!(pool._touched_others) + empty!(pool._touched_others_checkpoints) + push!(pool._touched_others_checkpoints, 0) # Sentinel + # Reset type touch tracking state (1-based sentinel pattern) pool._current_depth = 1 # 1 = global scope (sentinel) empty!(pool._touched_type_masks) diff --git a/src/types.jl b/src/types.jl index 9f159aaa..3595f87d 100644 --- a/src/types.jl +++ b/src/types.jl @@ -390,6 +390,14 @@ mutable struct AdaptiveArrayPool{S} <: AbstractArrayPool _others_ptr_bounds::Vector{UInt} # flat [ptr1,end1,ptr2,end2,...] _others_ptr_bounds_checkpoints::Vector{Int} # per-depth: saved length of bounds vector + # Touched-others tracking (per-scope selective fallback rewind). + # Flat stack of fallback typed pools first-touched at each open depth, plus + # per-depth saved lengths — same pattern as _others_ptr_bounds(+_checkpoints). + # Rewind drains only the current depth's segment: O(touched fallback types + # this scope) instead of O(all registered fallback pools). + _touched_others::Vector{Any} + _touched_others_checkpoints::Vector{Int} + # Borrow registry (S = 1 only) _pending_callsite::String # "" = no pending; set by macro before acquire _pending_return_site::String # "" = no pending; set by macro before validate @@ -420,6 +428,8 @@ function AdaptiveArrayPool{S}() where {S} Any[], # _others_values: empty cache UInt[], # _others_ptr_bounds: no bounds Int[0], # _others_ptr_bounds_checkpoints: sentinel + Any[], # _touched_others: no fallback touches yet + Int[0], # _touched_others_checkpoints: sentinel "", # _pending_callsite: no pending "", # _pending_return_site: no pending nothing, # _borrow_log: lazily created at S=1 @@ -474,6 +484,8 @@ _make_pool(runtime_check::Bool, old::AdaptiveArrayPool) = _make_pool(Int(runtime old._others_values, old._others_ptr_bounds, old._others_ptr_bounds_checkpoints, + old._touched_others, + old._touched_others_checkpoints, "", # _pending_callsite: reset "", # _pending_return_site: reset S >= 1 ? IdDict{Any, String}() : nothing, # _borrow_log diff --git a/test/runtests.jl b/test/runtests.jl index 49c602af..7bf101da 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -56,6 +56,7 @@ else include("test_coverage.jl") include("test_allocation.jl") include("test_fallback_reclamation.jl") + include("test_touched_others.jl") include("test_scope_depth_validation.jl") else include("test_aqua.jl") diff --git a/test/test_touched_others.jl b/test/test_touched_others.jl new file mode 100644 index 00000000..7c32d1f3 --- /dev/null +++ b/test/test_touched_others.jl @@ -0,0 +1,39 @@ +# Tests for the touched-others stack: per-scope selective fallback checkpoint/rewind. +# See docs/plans/DESIGN_fallback_touch_tracking.md. + +using Test +using AdaptiveArrayPools +using AdaptiveArrayPools: get_typed_pool!, checkpoint!, rewind!, + _lazy_checkpoint!, _lazy_rewind!, _typed_lazy_checkpoint!, _typed_lazy_rewind!, + _tracked_mask_for_types, _can_use_typed_path + +# Distinct isbits fallback types (not fixed slots) +struct TOFooA + x::Float64 +end +struct TOFooB + x::Float64 +end +struct TOFooC + x::Float64 +end + +@testset "touched-others: fields & lifecycle" begin + pool = AdaptiveArrayPool() + @test pool._touched_others == Any[] + @test pool._touched_others_checkpoints == [0] + + # reset! clears transient scope state, keeps registry + acquire!(pool, TOFooA, 4) + reset!(pool) + @test isempty(pool._touched_others) + @test pool._touched_others_checkpoints == [0] + @test haskey(pool.others, TOFooA) # registry kept + + # empty! clears everything + acquire!(pool, TOFooA, 4) + empty!(pool) + @test isempty(pool._touched_others) + @test pool._touched_others_checkpoints == [0] + @test !haskey(pool.others, TOFooA) +end From d4efd37f5d30ce40070ef604c832c025e4235e3c Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 14:41:14 -0700 Subject: [PATCH 03/10] (feat): plumb touched-others segment push/drain through checkpoint and rewind variants --- src/state.jl | 33 +++++++++++++++++++++++++++++++ test/test_touched_others.jl | 39 +++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/state.jl b/src/state.jl index e5333b08..7fa79df0 100644 --- a/src/state.jl +++ b/src/state.jl @@ -18,6 +18,7 @@ function checkpoint!(pool::AdaptiveArrayPool) pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) push!(pool._touched_has_others, false) + push!(pool._touched_others_checkpoints, length(pool._touched_others)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) depth = pool._current_depth @@ -52,6 +53,7 @@ Also updates _current_depth and bitmask state for type touch tracking. # _typed_lazy_rewind! iterates pool.others even if _acquire_impl! # (which bypasses _record_type_touch!) is the only acquire path. push!(pool._touched_has_others, _fixed_slot_bit(T) == UInt16(0)) + push!(pool._touched_others_checkpoints, length(pool._touched_others)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) _checkpoint_typed_pool!(get_typed_pool!(pool, T), pool._current_depth) return nothing @@ -83,6 +85,7 @@ compile-time unrolling. Increments _current_depth once for all types. pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) push!(pool._touched_has_others, $has_any_fallback) + push!(pool._touched_others_checkpoints, length(pool._touched_others)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) $(checkpoint_exprs...) nothing @@ -124,6 +127,7 @@ Performance: ~2ns vs ~540ns for full `checkpoint!`. # _LAZY_MODE_BIT = lazy mode flag (bits 0–7 are fixed-slot type bits) push!(pool._touched_type_masks, _LAZY_MODE_BIT) push!(pool._touched_has_others, false) + push!(pool._touched_others_checkpoints, length(pool._touched_others)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) depth = pool._current_depth # Eagerly checkpoint any pre-existing others entries. @@ -173,6 +177,7 @@ function rewind!(pool::AdaptiveArrayPool{S}) where {S} for tp in pool._others_values _rewind_typed_pool!(tp, cur_depth, S) end + _truncate_touched_others!(pool) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) @@ -198,6 +203,7 @@ Also updates _current_depth and bitmask state. return nothing end _rewind_typed_pool!(get_typed_pool!(pool, T), pool._current_depth, S) + _drain_touched_others!(pool, pool._current_depth) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) end @@ -233,6 +239,7 @@ Decrements _current_depth once after all types are rewound. return nothing end $(rewind_exprs...) + _drain_touched_others!(pool, pool._current_depth) if $S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) end @@ -372,6 +379,30 @@ end return nothing end +# ============================================================================== +# Touched-Others Stack (per-scope selective fallback checkpoint/rewind) +# ============================================================================== + +# Rewind and remove the fallback typed pools first-touched in the current scope. +# The current depth's segment of `_touched_others` is (base+1):end, where base is +# the saved length pushed by the matching checkpoint variant. O(touched this scope). +@inline function _drain_touched_others!(pool::AdaptiveArrayPool{S}, d::Int) where {S} + base = pop!(pool._touched_others_checkpoints) + stack = pool._touched_others + for i in (base + 1):length(stack) + _rewind_typed_pool!(@inbounds(stack[i]), d, S) + end + length(stack) > base && resize!(stack, base) + 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) + resize!(pool._touched_others, pop!(pool._touched_others_checkpoints)) + return nothing +end + """ _lazy_rewind!(pool::AdaptiveArrayPool) @@ -393,6 +424,7 @@ Called directly from the macro-generated `finally` clause as a single function c _rewind_typed_pool!(tp, d, S) end end + _drain_touched_others!(pool, d) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) end @@ -460,6 +492,7 @@ guaranteed by the `_TYPED_LAZY_BIT` mode set in `_typed_lazy_checkpoint!`. _rewind_typed_pool!(tp, d, S) end end + _drain_touched_others!(pool, d) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) end diff --git a/test/test_touched_others.jl b/test/test_touched_others.jl index 7c32d1f3..8b6c5968 100644 --- a/test/test_touched_others.jl +++ b/test/test_touched_others.jl @@ -37,3 +37,42 @@ end @test pool._touched_others_checkpoints == [0] @test !haskey(pool.others, TOFooA) end + +@testset "touched-others: checkpoint/rewind plumbing balance" begin + pool = AdaptiveArrayPool() + + # lazy pair + _lazy_checkpoint!(pool) + @test pool._touched_others_checkpoints == [0, 0] + _lazy_rewind!(pool) + @test pool._touched_others_checkpoints == [0] + + # typed single pair (fixed-slot type) + checkpoint!(pool, Float64) + @test length(pool._touched_others_checkpoints) == 2 + rewind!(pool, Float64) + @test pool._touched_others_checkpoints == [0] + + # typed multi pair + checkpoint!(pool, Float64, Int64) + rewind!(pool, Float64, Int64) + @test pool._touched_others_checkpoints == [0] + + # full pair + checkpoint!(pool) + rewind!(pool) + @test pool._touched_others_checkpoints == [0] + + # typed-lazy pair + _typed_lazy_checkpoint!(pool, Float64) + _typed_lazy_rewind!(pool, _tracked_mask_for_types(Float64)) + @test pool._touched_others_checkpoints == [0] + + # nesting + _lazy_checkpoint!(pool) + checkpoint!(pool, Float64) + @test length(pool._touched_others_checkpoints) == 3 + rewind!(pool, Float64) + _lazy_rewind!(pool) + @test pool._touched_others_checkpoints == [0] +end From d31de0b3f172ec18046932b5fb6395ccd36505f9 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 14:48:36 -0700 Subject: [PATCH 04/10] (test): add failing acceptance tests for selective fallback touch tracking --- test/test_touched_others.jl | 167 +++++++++++++++++++++++++++++++++++- 1 file changed, 166 insertions(+), 1 deletion(-) diff --git a/test/test_touched_others.jl b/test/test_touched_others.jl index 8b6c5968..2f312b7b 100644 --- a/test/test_touched_others.jl +++ b/test/test_touched_others.jl @@ -5,7 +5,7 @@ using Test using AdaptiveArrayPools using AdaptiveArrayPools: get_typed_pool!, checkpoint!, rewind!, _lazy_checkpoint!, _lazy_rewind!, _typed_lazy_checkpoint!, _typed_lazy_rewind!, - _tracked_mask_for_types, _can_use_typed_path + _tracked_mask_for_types, _can_use_typed_path, get_task_local_pool, @with_pool # Distinct isbits fallback types (not fixed slots) struct TOFooA @@ -76,3 +76,168 @@ end _lazy_rewind!(pool) @test pool._touched_others_checkpoints == [0] end + +@testset "touched-others: no eager checkpoint on lazy entry (pollution regression)" begin + pool = AdaptiveArrayPool() + # Register three fallback types at global scope, then reset counters + acquire!(pool, TOFooA, 4); acquire!(pool, TOFooB, 4); acquire!(pool, TOFooC, 4) + reset!(pool) + tpA = get_typed_pool!(pool, TOFooA) + tpB = get_typed_pool!(pool, TOFooB) + tpC = get_typed_pool!(pool, TOFooC) + + _lazy_checkpoint!(pool) + # THE regression assertion: unrelated registered fallbacks are NOT touched + @test tpB._checkpoint_depths == [0] + @test tpC._checkpoint_depths == [0] + + v = acquire!(pool, TOFooA, 8) + @test tpA._checkpoint_depths[end] == 2 # first-touch checkpoint at depth 2 + @test length(pool._touched_others) == 1 + @test pool._touched_others[end] === tpA + @test tpA.n_active == 1 + + # Re-acquire same type: no duplicate stack entry + acquire!(pool, TOFooA, 8) + @test length(pool._touched_others) == 1 + + _lazy_rewind!(pool) + @test tpA.n_active == 0 + @test isempty(pool._touched_others) + @test tpB._checkpoint_depths == [0] # still never visited + @test tpC._checkpoint_depths == [0] +end + +@testset "touched-others: nested scopes, same fallback type at two depths" begin + pool = AdaptiveArrayPool() + tpA = get_typed_pool!(pool, TOFooA) + + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, TOFooA, 4) + @test tpA.n_active == 1 + + _lazy_checkpoint!(pool) # depth 3 + acquire!(pool, TOFooA, 4) + acquire!(pool, TOFooA, 4) + @test tpA.n_active == 3 + @test length(pool._touched_others) == 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: nested scope NOT touching outer's fallback" begin + pool = AdaptiveArrayPool() + tpA = get_typed_pool!(pool, TOFooA) + + _lazy_checkpoint!(pool) # depth 2 + acquire!(pool, TOFooA, 4) + _lazy_checkpoint!(pool) # depth 3: does not touch TOFooA + zeros!(pool, 16) # Float64 fixed-slot work only + @test length(pool._touched_others) == 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 = AdaptiveArrayPool() + tpB = get_typed_pool!(pool, TOFooB) + + _typed_lazy_checkpoint!(pool, Float64) + zeros!(pool, 8) # tracked fixed-slot work + acquire!(pool, TOFooB, 4) # untracked helper-style fallback touch + @test pool._touched_has_others[end] == true + @test pool._touched_others[end] === tpB + @test !_can_use_typed_path(pool, _tracked_mask_for_types(Float64)) + _typed_lazy_rewind!(pool, _tracked_mask_for_types(Float64)) + @test tpB.n_active == 0 +end + +@testset "touched-others: tracked fallback type via typed checkpoint!" begin + pool = AdaptiveArrayPool() + + checkpoint!(pool, TOFooA) # fallback T tracked by macro + tpA = get_typed_pool!(pool, TOFooA) + @test pool._touched_others[end] === tpA # pushed at checkpoint + acquire!(pool, TOFooA, 4) # public-API acquire: no double push + @test count(tp -> tp === tpA, pool._touched_others) == 1 + # macro exit path for has_others=true is _typed_lazy_rewind! + _typed_lazy_rewind!(pool, _tracked_mask_for_types(TOFooA)) + @test tpA.n_active == 0 + @test pool._touched_others_checkpoints == [0] +end + +@testset "touched-others: new type registered mid-scope" begin + pool = AdaptiveArrayPool() + _lazy_checkpoint!(pool) + acquire!(pool, TOFooC, 4) # first-ever registration, in-scope + tpC = get_typed_pool!(pool, TOFooC) + @test pool._touched_others[end] === tpC + @test count(tp -> tp === tpC, pool._touched_others) == 1 + _lazy_rewind!(pool) + @test tpC.n_active == 0 +end + +@testset "touched-others: full checkpoint!/rewind! pairing unchanged" begin + pool = AdaptiveArrayPool() + acquire!(pool, TOFooA, 4); reset!(pool) + tpA = get_typed_pool!(pool, TOFooA) + + checkpoint!(pool) # eager: checkpoints ALL others + @test tpA._checkpoint_depths[end] == 2 + acquire!(pool, TOFooA, 4) + @test isempty(pool._touched_others) # guard saw existing depth-2 entry + rewind!(pool) + @test tpA.n_active == 0 + @test pool._touched_others_checkpoints == [0] +end + +@testset "touched-others: similar! records fallback touch" begin + pool = AdaptiveArrayPool() + src = [TOFooA(1.0), TOFooA(2.0)] + _lazy_checkpoint!(pool) + similar!(pool, src) + tpA = get_typed_pool!(pool, TOFooA) + @test pool._touched_others[end] === tpA + _lazy_rewind!(pool) + @test tpA.n_active == 0 +end + +@testset "touched-others: @with_pool integration + exception-leak recovery" begin + # Integration through the real macro (task-local pool) + f_leaf(n) = @with_pool p begin + q = acquire!(p, TOFooA, n) + length(q) + end + @test f_leaf(8) == 8 + tl = get_task_local_pool() + @test get_typed_pool!(tl, TOFooA).n_active == 0 + @test isempty(tl._touched_others) + + # Inner scope throws, outer catches: outer exit must clean up leaked state + function f_outer() + @with_pool p begin + acquire!(p, TOFooB, 4) + try + @with_pool p2 begin + acquire!(p2, TOFooC, 4) + error("boom") + end + catch + end + 1 + end + end + @test f_outer() == 1 + @test tl._current_depth == 1 + @test get_typed_pool!(tl, TOFooB).n_active == 0 + @test get_typed_pool!(tl, TOFooC).n_active == 0 + @test isempty(tl._touched_others) + @test tl._touched_others_checkpoints == [0] + empty!(tl) # leave the task-local pool clean for other test files +end From 8a0fbdde416852bb26a59e6922f8969388354ba4 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 15:31:39 -0700 Subject: [PATCH 05/10] (perf): selective fallback checkpoint/rewind via touched-others stack @with_pool scopes now checkpoint/rewind only the fallback typed pools they actually touch (first-touch lazy checkpoint, per-depth stack segment) instead of eagerly checkpointing and rewinding every registered fallback pool. Fixes order-dependent latency where unrelated fallback registrations slowed hot call sites (active=0 pollution). --- src/acquire.jl | 6 ++++ src/state.jl | 92 ++++++++++++++++++++++++++------------------------ src/types.jl | 3 ++ 3 files changed, 57 insertions(+), 44 deletions(-) diff --git a/src/acquire.jl b/src/acquire.jl index 8fd9ecd8..12ed1f7b 100644 --- a/src/acquire.jl +++ b/src/acquire.jl @@ -310,6 +310,12 @@ end b = _fixed_slot_bit(T) if b == UInt16(0) @inbounds pool._touched_has_others[depth] = true + # First-touch lazy checkpoint for fallback types (identity-list analogue of + # the fixed-slot bitmask path below). depth == 1 (global scope) is exempt: + # there is no scope to rewind to, matching get_typed_pool!'s gate. + if depth > 1 + _touch_fallback_pool!(pool, get_typed_pool!(pool, T), depth) + end else current_mask = @inbounds pool._touched_type_masks[depth] # Lazy checkpoint: lazy mode (bit 15) OR typed lazy mode (bit 14), AND first touch. diff --git a/src/state.jl b/src/state.jl index 7fa79df0..3140284e 100644 --- a/src/state.jl +++ b/src/state.jl @@ -55,7 +55,13 @@ Also updates _current_depth and bitmask state for type touch tracking. push!(pool._touched_has_others, _fixed_slot_bit(T) == UInt16(0)) push!(pool._touched_others_checkpoints, length(pool._touched_others)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) - _checkpoint_typed_pool!(get_typed_pool!(pool, T), pool._current_depth) + if _fixed_slot_bit(T) == UInt16(0) + # Fallback type: checkpoint + record in the touched-others stack so the + # selective rewind path visits exactly this pool. + _touch_fallback_pool!(pool, get_typed_pool!(pool, T), pool._current_depth) + else + _checkpoint_typed_pool!(get_typed_pool!(pool, T), pool._current_depth) + end return nothing end @@ -80,7 +86,13 @@ compile-time unrolling. Increments _current_depth once for all types. # If so, push has_others=true so _typed_lazy_rewind! iterates pool.others # even when _acquire_impl! (bypassing _record_type_touch!) is used. has_any_fallback = any(i -> _fixed_slot_bit(types[i].parameters[1]) == UInt16(0), unique_indices) - checkpoint_exprs = [:(_checkpoint_typed_pool!(get_typed_pool!(pool, types[$i]), pool._current_depth)) for i in unique_indices] + checkpoint_exprs = map(unique_indices) do i + if _fixed_slot_bit(types[i].parameters[1]) == UInt16(0) + :(_touch_fallback_pool!(pool, get_typed_pool!(pool, types[$i]), pool._current_depth)) + else + :(_checkpoint_typed_pool!(get_typed_pool!(pool, types[$i]), pool._current_depth)) + end + end return quote pool._current_depth += 1 push!(pool._touched_type_masks, UInt16(0)) @@ -115,9 +127,11 @@ Increments `_current_depth` and pushes bitmask sentinels — but does **not** sa `_touched_type_masks` marks this depth as lazy mode so that `_record_type_touch!` can trigger lazy first-touch checkpoints. -Existing `others` entries are eagerly checkpointed since there is no per-type -tracking for non-fixed-slot pools; Case B in `_rewind_typed_pool!` handles any -new `others` entries created during the scope (n_active starts at 0 = sentinel). +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. Performance: ~2ns vs ~540ns for full `checkpoint!`. """ @@ -129,13 +143,6 @@ Performance: ~2ns vs ~540ns for full `checkpoint!`. push!(pool._touched_has_others, false) push!(pool._touched_others_checkpoints, length(pool._touched_others)) _runtime_check(pool) && push!(pool._others_ptr_bounds_checkpoints, length(pool._others_ptr_bounds)) - depth = pool._current_depth - # Eagerly checkpoint any pre-existing others entries. - # New others types created during the scope start at n_active=0 (sentinel covers them). - for p in pool._others_values - _checkpoint_typed_pool!(p, depth) - @inbounds pool._touched_has_others[depth] = true - end return nothing end @@ -202,7 +209,11 @@ Also updates _current_depth and bitmask state. reset!(get_typed_pool!(pool, T), S) return nothing end - _rewind_typed_pool!(get_typed_pool!(pool, T), pool._current_depth, S) + # Fixed-slot types rewind directly; fallback types were pushed onto the + # touched-others stack by checkpoint!(pool, T) and are covered by the drain. + if _fixed_slot_bit(T) != UInt16(0) + _rewind_typed_pool!(get_typed_pool!(pool, T), pool._current_depth, S) + end _drain_touched_others!(pool, pool._current_depth) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) @@ -230,7 +241,8 @@ Decrements _current_depth once after all types are rewound. push!(unique_indices, i) end end - rewind_exprs = [:(_rewind_typed_pool!(get_typed_pool!(pool, types[$i]), pool._current_depth, $S)) for i in reverse(unique_indices)] + fixed_indices = [i for i in unique_indices if _fixed_slot_bit(types[i].parameters[1]) != UInt16(0)] + rewind_exprs = [:(_rewind_typed_pool!(get_typed_pool!(pool, types[$i]), pool._current_depth, $S)) for i in reverse(fixed_indices)] reset_exprs = [:(reset!(get_typed_pool!(pool, types[$i]), $S)) for i in unique_indices] return quote # Safety guard: at global scope (depth=1), delegate to reset! @@ -403,6 +415,20 @@ end return nothing end +# First-touch producer: checkpoint a fallback pool at `depth` and record it in the +# touched-others stack, exactly once per depth. The `_checkpoint_depths[end] != depth` +# guard makes this idempotent across the three producer paths (typed checkpoint!, +# _record_type_touch!, get_typed_pool! registration) 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) + if @inbounds(tp._checkpoint_depths[end]) != depth + push!(tp._checkpoint_n_active, tp.n_active) + push!(tp._checkpoint_depths, depth) + push!(pool._touched_others, tp) + end + return nothing +end + """ _lazy_rewind!(pool::AdaptiveArrayPool) @@ -419,11 +445,6 @@ Called directly from the macro-generated `finally` clause as a single function c d = pool._current_depth bits = @inbounds(pool._touched_type_masks[d]) & _TYPE_BITS_MASK _selective_rewind_fixed_slots!(pool, bits) # S propagated via pool type - if @inbounds(pool._touched_has_others[d]) - for tp in pool._others_values - _rewind_typed_pool!(tp, d, S) - end - end _drain_touched_others!(pool, d) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) @@ -440,34 +461,22 @@ end Typed checkpoint that enables lazy first-touch checkpointing for extra types touched by helpers (`use_typed=true`, `_can_use_typed_path=false` path). -Calls `checkpoint!(pool, types...)` (checkpoints only the statically-known types), -then sets `_TYPED_LAZY_BIT` (bit 14) in `_touched_type_masks[depth]` to signal typed lazy mode. +Calls `checkpoint!(pool, types...)` (checkpoints only the statically-known types — +fallback types among them are pushed onto the touched-others stack via +`_touch_fallback_pool!`), then sets `_TYPED_LAZY_BIT` (bit 14) in +`_touched_type_masks[depth]` to signal typed lazy mode. `_record_type_touch!` checks `(mask & _MODE_BITS_MASK) != 0` (bit 14 OR bit 15) to trigger a lazy first-touch checkpoint for each extra type on first acquire, ensuring Case A -(not Case B) applies at rewind and parent `n_active` is preserved correctly. +(not Case B) applies at rewind and parent `n_active` is preserved correctly. Fallback +types touched by helpers (beyond the statically-known `types...`) are first-touch +checkpointed and pushed onto the touched-others stack by `_record_type_touch!`'s +`b == UInt16(0)` branch (Step 3), then drained selectively at rewind. """ @inline function _typed_lazy_checkpoint!(pool::AdaptiveArrayPool, types::Type...) checkpoint!(pool, types...) d = pool._current_depth @inbounds pool._touched_type_masks[d] |= _TYPED_LAZY_BIT - - # Eagerly snapshot pre-existing others entries — mirrors _lazy_checkpoint!. - # _record_type_touch! cannot lazy-checkpoint others types (b==0 branch, no per-type bit). - # Without this, a helper that re-acquires an already-active others type triggers Case B - # at rewind and restores the wrong parent n_active value. - # - # Also set has_others=true when pool.others is non-empty, so _typed_lazy_rewind! - # enters the others loop even for tracked non-fixed-slot types (e.g. CPU Float16) that - # used _acquire_impl! (bypassing _record_type_touch!, leaving has_others=false otherwise). - # 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 pool._others_values - if @inbounds(p._checkpoint_depths[end]) != d - _checkpoint_typed_pool!(p, d) - end - @inbounds pool._touched_has_others[d] = true - end return nothing end @@ -487,11 +496,6 @@ guaranteed by the `_TYPED_LAZY_BIT` mode set in `_typed_lazy_checkpoint!`. touched = @inbounds(pool._touched_type_masks[d]) & _TYPE_BITS_MASK combined = tracked_mask | touched _selective_rewind_fixed_slots!(pool, combined) # S propagated via pool type - if @inbounds(pool._touched_has_others[d]) - for tp in pool._others_values - _rewind_typed_pool!(tp, d, S) - end - end _drain_touched_others!(pool, d) if S >= 1 && length(pool._others_ptr_bounds_checkpoints) > 1 resize!(pool._others_ptr_bounds, pop!(pool._others_ptr_bounds_checkpoints)) diff --git a/src/types.jl b/src/types.jl index 3595f87d..7f12c5a6 100644 --- a/src/types.jl +++ b/src/types.jl @@ -533,6 +533,9 @@ const _FIXED_SLOT_TYPES = Union{Float64, Float32, Int64, Int32, ComplexF64, Comp # _record_type_touch!) would leave has_others=false, causing the # rewind to skip pool.others entirely and leak this new type's n_active. @inbounds p._touched_has_others[p._current_depth] = true + # Record in the touched-others stack so the selective rewind visits this + # brand-new pool (its checkpoint was just pushed above). + push!(p._touched_others, new_tp) end return new_tp end From 5c162ff125ed88a7c17d505abfda5f9b7946c829 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 15:31:44 -0700 Subject: [PATCH 06/10] (test): isolate touched-others @with_pool testset from task-local pool state The task-local pool is a process-wide singleton shared by every test file. @with_pool's fast path does not guarantee cleanup for exceptions that escape its own scope (only nested leaks within one invocation are caught by the entry-depth guard), so earlier files that deliberately throw across a @with_pool boundary (e.g. test_macro_expansion.jl's stack-trace tests) can leave _current_depth elevated before this testset runs. Reset the task-local pool at the start of the testset so its absolute-depth assertions hold regardless of test file execution order. Root-caused by direct repro (with and without the touched-others changes: identical leak in both), confirming it is pre-existing and unrelated to the semantic switch. --- test/test_touched_others.jl | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/test/test_touched_others.jl b/test/test_touched_others.jl index 2f312b7b..510dbad8 100644 --- a/test/test_touched_others.jl +++ b/test/test_touched_others.jl @@ -209,6 +209,16 @@ end end @testset "touched-others: @with_pool integration + exception-leak recovery" begin + # The task-local pool is a process-wide singleton shared with every other test + # file. `@with_pool`'s fast (non-try/finally) path does not guarantee cleanup + # for exceptions that escape ITS OWN scope (only nested leaks within one + # invocation are caught by the entry-depth guard — see macros.jl docstring), + # so earlier test files that deliberately throw across a `@with_pool` boundary + # (e.g. stack-trace tests in test_macro_expansion.jl) can leave `_current_depth` + # elevated. Start from a known-clean baseline so the absolute-depth assertions + # below are independent of test file execution order. + empty!(get_task_local_pool()) + # Integration through the real macro (task-local pool) f_leaf(n) = @with_pool p begin q = acquire!(p, TOFooA, n) From 8ebf053138b00a551a6a69ff05f580ebbb0563c2 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 15:48:01 -0700 Subject: [PATCH 07/10] (test): zero-allocation regression for touched-others roundtrips --- test/test_touched_others.jl | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/test/test_touched_others.jl b/test/test_touched_others.jl index 510dbad8..00d729ea 100644 --- a/test/test_touched_others.jl +++ b/test/test_touched_others.jl @@ -251,3 +251,32 @@ end @test tl._touched_others_checkpoints == [0] empty!(tl) # leave the task-local pool clean for other test files end + +# Function barriers for @allocated measurements +function _to_lazy_roundtrip(pool, n) + _lazy_checkpoint!(pool) + v = acquire!(pool, TOFooA, n) + s = length(v) + _lazy_rewind!(pool) + return s +end + +function _to_macro_roundtrip(n) + return @with_pool p begin + v = acquire!(p, TOFooA, n) + length(v) + end +end + +@testset "touched-others: zero allocation after warmup" begin + pool = AdaptiveArrayPool() + # Pollute the registry to also confirm cost independence at the alloc level + acquire!(pool, TOFooB, 4); acquire!(pool, TOFooC, 4); reset!(pool) + + _to_lazy_roundtrip(pool, 32); _to_lazy_roundtrip(pool, 32) # warmup + @test @allocated(_to_lazy_roundtrip(pool, 32)) == 0 + + _to_macro_roundtrip(32); _to_macro_roundtrip(32) # warmup + @test @allocated(_to_macro_roundtrip(32)) == 0 + empty!(get_task_local_pool()) +end From 4b4207837dfc6f0aae56cdffc021d18d91b90a1d Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 16:45:11 -0700 Subject: [PATCH 08/10] (docs): add reset! fallback-registry note --- src/state.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/state.jl b/src/state.jl index 3140284e..8adce9f3 100644 --- a/src/state.jl +++ b/src/state.jl @@ -669,6 +669,12 @@ reset!(pool) # Restore state, keep allocated memory # Now pool.n_active == 0, but vectors are still available for reuse ``` +!!! note "Fallback registry" + `reset!` keeps registered fallback typed pools (`pool.others`) so their + buffers can be reused. Since the touched-others tracking, inactive registered + fallbacks add no per-scope checkpoint/rewind cost; use `empty!` only when the + buffers themselves should be released. + See also: [`empty!`](@ref), [`rewind!`](@ref) """ function reset!(pool::AdaptiveArrayPool{S}) where {S} From 1ec0984a32755cee2dfb60162222d824f260852d Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 16:45:11 -0700 Subject: [PATCH 09/10] (docs): fix dead design-doc links from removed docs/design files --- docs/src/architecture/design-docs.md | 46 ++++++--------------------- docs/src/architecture/how-it-works.md | 11 +++---- test/test_compact.jl | 2 +- test/test_touched_others.jl | 5 ++- 4 files changed, 19 insertions(+), 45 deletions(-) diff --git a/docs/src/architecture/design-docs.md b/docs/src/architecture/design-docs.md index 4efb4e9d..ed55c6a2 100644 --- a/docs/src/architecture/design-docs.md +++ b/docs/src/architecture/design-docs.md @@ -1,41 +1,15 @@ # Design Documents -For in-depth analysis of design decisions, implementation tradeoffs, and architectural choices, see the design documents in the repository: - -## API Design - -- **[hybrid_api_design.md](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/hybrid_api_design.md)** - Historical two-API strategy (`acquire!` vs `acquire_view!`) and type stability analysis - -## Caching & Performance - -- **[nd_array_approach_comparison.md](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/nd_array_approach_comparison.md)** - N-way cache design (now legacy — replaced by `setfield!` reuse on Julia 1.12+ CPU), boxing analysis, and ReshapedArray benchmarks - -- **[fixed_slots_codegen_design.md](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/fixed_slots_codegen_design.md)** - Zero-allocation iteration via `@generated` functions and fixed-slot type dispatch - -## Macro Internals - -- **[untracked_acquire_design.md](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/untracked_acquire_design.md)** - Macro-based untracked acquire detection and 1-based sentinel pattern - -## Backend Extensions - -- **[cuda_extension_design.md](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/cuda_extension_design.md)** - CUDA backend architecture and package extension loading - ---- - -## Document Overview - -| Document | Focus Area | Key Insights | -|----------|------------|--------------| -| hybrid_api_design | API strategy | Array by default, views via `acquire_view!` | -| nd_array_approach_comparison | Caching | N-way cache (legacy); setfield! reuse on Julia 1.12+ CPU | -| fixed_slots_codegen_design | Codegen | @generated functions enable type-stable iteration | -| untracked_acquire_design | Macro safety | Sentinel pattern ensures correct cleanup | -| cuda_extension_design | GPU support | Seamless CPU/CUDA API parity | +The historical design documents that used to live under `docs/design/` (covering +the hybrid API strategy, the N-D array approach comparison, fixed-slot codegen, +untracked-acquire detection, and the CUDA extension) have been removed as +outdated — the decisions they recorded either shipped and became the current +implementation, or were superseded by later approaches (e.g. the N-way cache +described in some of them is legacy on Julia 1.12+, see below). + +Current architecture documentation lives in [How It Works](how-it-works.md), +which explains the mechanisms that enable zero-allocation array reuse as they +exist today. ## See Also diff --git a/docs/src/architecture/how-it-works.md b/docs/src/architecture/how-it-works.md index d1bb283f..10f512a5 100644 --- a/docs/src/architecture/how-it-works.md +++ b/docs/src/architecture/how-it-works.md @@ -178,10 +178,7 @@ This pattern reduces branching in hot paths where every nanosecond counts. ## Further Reading -For detailed design documents: - -- [`hybrid_api_design.md`](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/hybrid_api_design.md) — Two-API strategy (`acquire!` vs `acquire_view!`) and type stability analysis -- [`nd_array_approach_comparison.md`](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/nd_array_approach_comparison.md) — N-way cache design, boxing analysis, and ReshapedArray benchmarks -- [`untracked_acquire_design.md`](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/untracked_acquire_design.md) — Macro-based untracked acquire detection and 1-based sentinel pattern -- [`fixed_slots_codegen_design.md`](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/fixed_slots_codegen_design.md) — Zero-allocation iteration via `@generated` functions -- [`cuda_extension_design.md`](https://github.com/ProjectTorreyPines/AdaptiveArrayPools.jl/blob/master/docs/design/cuda_extension_design.md) — CUDA backend architecture and extension loading +The historical design documents that used to be linked here have been removed +as outdated. See [Design Documents](design-docs.md) for details, or the +[Type Dispatch & Cache](type-dispatch.md) and [`@with_pool` Macro Internals](macro-internals.md) +pages for the current technical deep-dives. diff --git a/test/test_compact.jl b/test/test_compact.jl index 264f0024..b82f3338 100644 --- a/test/test_compact.jl +++ b/test/test_compact.jl @@ -4,7 +4,7 @@ using AdaptiveArrayPools: checkpoint!, rewind! # ============================================================================== # compact! — capacity compaction (shrink over-allocated backing buffers in place) -# Design: docs/design/capacity_compaction.md | Plan: docs/plans/PLAN_compact.md +# Design: see the compact! docstring in src/state.jl | Plan: docs/plans/PLAN_compact.md # # Model: a slot's backing Vector length is the HIGH-WATER mark (largest ever # acquired; _claim_slot! only grows). The CURRENT use is recorded per-slot in diff --git a/test/test_touched_others.jl b/test/test_touched_others.jl index 00d729ea..544e11f0 100644 --- a/test/test_touched_others.jl +++ b/test/test_touched_others.jl @@ -1,5 +1,8 @@ # Tests for the touched-others stack: per-scope selective fallback checkpoint/rewind. -# See docs/plans/DESIGN_fallback_touch_tracking.md. +# Invariant: a fallback pool has a checkpoint entry at depth d ⟺ it is in the depth-d +# segment of pool._touched_others — except under full checkpoint!(pool), whose eager +# sweep pairs with full rewind!(pool)'s sweep (segment stays empty, truncate-only). +# See the docstrings of _touch_fallback_pool!/_drain_touched_others! in src/state.jl. using Test using AdaptiveArrayPools From 3226741dea06e05a0867b8f55e39715b4d47a4f7 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Thu, 9 Jul 2026 17:07:04 -0700 Subject: [PATCH 10/10] (docs): exempt Full API page from Documenter size_threshold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit reference/api.md renders at 200.05 KiB in CI — just over Documenter's 200 KiB hard limit — because GitHub source links (remotes, CI-only) add a few KiB that local builds don't have. Exempt only this page; all others keep the guardrail. --- docs/make.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/make.jl b/docs/make.jl index 5f25be26..63276f79 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -113,6 +113,10 @@ makedocs( canonical = "https://projecttorreypines.github.io/AdaptiveArrayPools.jl", edit_link = :commit, assets = String[], + # The single-page Full API reference renders just over Documenter's 200 KiB + # hard limit in CI (GitHub source links, absent locally, add a few KiB). + # Exempt only this page; every other page keeps the size guardrail. + size_threshold_ignore = ["reference/api.md"], ), pages = [ "Home" => "index.md",