Untracked Acquire Detection & State Management Improvements #6
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…action - Add support for acquire!(pool, x) form which generates eltype(x) for checkpoint/rewind optimization - Add unsafe_acquire! recognition to _extract_acquire_types - Handle eltype(x) expressions in _filter_static_types: - External variables: use eltype(x) for typed checkpoint - Local variables: fall back to full checkpoint (not defined at checkpoint time) - Update docstrings to reflect new functionality
_extract_acquire_types tests: - similar-style acquire!(pool, x) generates eltype(x) - unsafe_acquire! single and multiple types - mixed acquire!/unsafe_acquire!/similar-style - qualified unsafe_acquire! calls - pool isolation (no AST pollution) _filter_static_types tests: - eltype(x) with external variable (usable) - eltype(x) with local variable (fallback) - eltype mixed with static types - complex expressions in eltype Macro expansion tests: - Similar-style expansion with external array - Local array fallback behavior - unsafe_acquire! type extraction - Mixed traditional/unsafe/similar-style
_extract_acquire_types tests: - custom struct type extraction (MyCustomType) - multiple custom types - type parameter T (from where clause) - multiple type parameters (T, S) - mixed: builtin, custom, type parameter Integration tests: - custom type with filtering - type parameter integration - type parameter vs local variable conflict (T = eltype(x)) - mixed types with filtering Macro expansion tests: - custom struct type expansion - type parameter T expansion - mixed builtin/custom/type parameter
…tion Extends _extract_acquire_types to recognize all acquire function aliases: - acquire_view! (alias for acquire!) - acquire_array! (alias for unsafe_acquire!) This ensures @with_pool macro generates typed checkpoint!/rewind! calls for code using the explicit naming aliases, not just the base functions. Includes comprehensive tests for alias type extraction and integration.
Add fields and functions to support detecting and handling acquire calls from inner functions that the macro cannot see at compile time. Changes to TypedPool: - saved_depths: Track which checkpoint depth each saved_stack entry belongs to Changes to AdaptiveArrayPool: - check_depth: Current checkpoint nesting depth counter - had_untracked: Per-depth flag indicating if untracked acquire occurred New functions: - all_type_stacks(): Iterator over all TypedPools (6 fixed + others) - foreach_type_stack(): Callback pattern for TypedPool iteration Refs: design/untracked_acquire_design.md sections 3.3, 4.5
Implement the core functionality for detecting acquire calls from inner functions that the macro cannot see at compile time (untracked acquires). New internal implementation functions: - _acquire_impl!: Core acquire logic (no marking overhead) - _unsafe_acquire_impl!: Core unsafe_acquire logic (no marking overhead) - _mark_untracked!: Mark current depth as having untracked acquire Updated user-facing functions: - acquire!: Now wraps _acquire_impl! with _mark_untracked! call - unsafe_acquire!: Now wraps _unsafe_acquire_impl! with _mark_untracked! call Updated state management: - checkpoint!: Increments check_depth, initializes had_untracked flag - rewind!: Uses saved_depths for accurate pop/restore decision - _checkpoint_typed_pool!: Helper with depth tracking - _rewind_typed_pool!: Helper with saved_depths logic - empty!: Resets check_depth and had_untracked New aliases for macro transformation: - _acquire_view_impl! = _acquire_impl! - _acquire_array_impl! = _unsafe_acquire_impl! Design: Macro-transformed code calls _*_impl! directly (zero overhead), while code outside @with_pool calls acquire! wrapper which marks untracked. Refs: design/untracked_acquire_design.md sections 4.1-4.4
Transform acquire! calls to _acquire_impl! in macro-generated code, enabling zero-overhead acquire within @with_pool while detecting untracked acquires from inner functions. New transformation function: - _transform_acquire_calls: AST transformation for acquire calls - acquire!/acquire_view! → _acquire_impl! - unsafe_acquire!/acquire_array! → _unsafe_acquire_impl! - Supports qualified names (AAP.acquire!, A.B.C.acquire!, etc.) - Only transforms calls matching the target pool name GlobalRef constants for module-qualified references: - _ACQUIRE_IMPL_REF: GlobalRef to AdaptiveArrayPools._acquire_impl! - _UNSAFE_ACQUIRE_IMPL_REF: GlobalRef to AdaptiveArrayPools._unsafe_acquire_impl! Updated macro code generation: - _generate_pool_code: Transform body, add had_untracked check - _generate_function_pool_code: Transform body, add had_untracked check - Typed rewind now falls back to full rewind if untracked detected Refs: design/untracked_acquire_design.md sections 4.3, 4.6
Add comprehensive tests for the acquire call transformation function used by macros to enable zero-overhead acquire within @with_pool. Test categories: - Basic transformation: acquire!, unsafe_acquire!, acquire_view!, acquire_array! - Qualified names: AdaptiveArrayPools.acquire!, AAP.acquire!, A.B.C.acquire! - Pool name matching: Only transforms calls with matching pool argument - Recursive transformation: Nested blocks, function calls, mixed pools - Similar-style: acquire!(pool, x) form transformation - GlobalRef verification: Correct module binding for _*_impl! functions Total: 33 new tests for _transform_acquire_calls
…llback - Change check_depth to 1-indexed (Julia idiomatic, 1 = global scope) - Add checkpoint-time check: use full checkpoint if parent had untracked acquires - Silently reset n_active=0 for helper-acquired new types instead of erroring - Auto-checkpoint new TypedPools created inside @with_pool scope - Add 4 tests: parent protection, same type protection, new type reset, typed performance
…ames - Rename TypedPool fields: - saved_stack → _checkpoint_n_active (stores n_active at checkpoints) - saved_depths → _checkpoint_depths (stores depth of each checkpoint) - Rename AdaptiveArrayPool fields: - check_depth → _current_depth (current scope depth, 1=global) - had_untracked → _untracked_flags (per-depth untracked flags) - Initialize with sentinel values to eliminate isempty() checks: - TypedPool: _checkpoint_n_active=[0], _checkpoint_depths=[0] - AdaptiveArrayPool: _current_depth=1, _untracked_flags=[false] - Simplify _rewind_typed_pool! by removing redundant isempty checks - Update empty!() functions to restore sentinel values This provides 17-27% speedup in checkpoint/rewind hot paths.
When a nested scope does full checkpoint but typed rewind, checkpoints from the nested scope remain on the stack (orphaned). This caused incorrect state restoration when outer scopes performed full rewind. Changes: - Add orphaned checkpoint cleanup loop in _rewind_typed_pool! - Pop stale checkpoints (depth > current) before normal rewind logic - Add complex nested test cases for L1 @with_pool → L2 untracked → L3 @with_pool
- Add L1→L2→L3 nested tests using unsafe_acquire! instead of acquire! - Add acquire_array! alias tests in nested scenarios (including 2D) - Add mixed acquire!/unsafe_acquire!/acquire_array! test in single scenario - Add acquire_view! alias nested test with ReshapedArray - Add N-D unsafe_acquire! nested tests (2D, 3D, tuple dimensions) Verifies that all acquire variants trigger untracked detection correctly and state management works across nested @with_pool scopes.
Improve code organization by separating concerns: - acquire.jl (~410 lines): safe_prod, get_view!, get_nd_*!, _mark_untracked!, _*_impl!, acquire!, unsafe_acquire!, aliases - state.jl (~240 lines): checkpoint!, rewind!, empty!, @generated typed versions Include order: types.jl → utils.jl → acquire.jl → state.jl → task_local_pool.jl → macros.jl
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6 +/- ##
==========================================
- Coverage 97.27% 94.60% -2.67%
==========================================
Files 6 7 +1
Lines 440 575 +135
==========================================
+ Hits 428 544 +116
- Misses 12 31 +19 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
This PR implements automatic detection of "untracked" acquire! calls (from helper functions or parent scope) to prevent pool state corruption in nested @with_pool scenarios. It also introduces a 1-based sentinel pattern for improved performance and reorganizes code for better maintainability.
Key Changes
- Untracked Acquire Detection: Tracks untracked
acquire!calls via_untracked_flagsper depth level, automatically falling back to full checkpoint/rewind when detected to protect parent scope arrays - 1-Based Sentinel Pattern: Eliminates
isempty()checks in hot paths by maintaining_checkpoint_n_active=[0]and_checkpoint_depths=[0]sentinel values, providing ~2-5% performance improvement - Macro Transformation: Transforms
acquire!calls inside@with_poolto_acquire_impl!, bypassing untracked marking overhead since macro already tracks these calls - Code Organization: Splits
core.jlintoacquire.jl(~410 lines) andstate.jl(~240 lines) for better separation of concerns
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/types.jl | Added _current_depth and _untracked_flags fields to AdaptiveArrayPool; renamed saved_stack to _checkpoint_n_active and added _checkpoint_depths to TypedPool; implemented auto-checkpoint for dynamically created TypedPool instances |
| src/state.jl | New file containing all checkpoint/rewind/empty! implementations with 1-based sentinel pattern and depth-aware rewind logic |
| src/acquire.jl | New file containing acquisition APIs; added _mark_untracked! and internal _acquire_impl!/_unsafe_acquire_impl! functions for macro transformation |
| src/macros.jl | Enhanced type extraction to support unsafe_acquire!, acquire_view!, acquire_array! aliases and eltype() expressions; added _transform_acquire_calls to transform acquire calls to _impl! variants; added untracked flag checks for dynamic fallback |
| src/AdaptiveArrayPools.jl | Updated includes to reference new acquire.jl and state.jl files |
| test/test_state.jl | Added comprehensive nested scope tests covering parent scope protection, helper functions, mixed acquire variants, and N-dimensional arrays; updated allocation tests to reference new field names |
| test/test_macro_internals.jl | Added extensive tests for type extraction with aliases, similar-style forms, custom types, type parameters, and transformation logic |
| test/test_macro_expansion.jl | Added macro expansion tests for new features including similar-style, unsafe_acquire!, and custom types |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Implement automatic detection of "untracked"
acquire!calls (from inner functions or parent scope outside@with_pool), preventing pool state corruption in nested scope scenarios. Also includes code organization improvements.Key Changes
1. Untracked Acquire Detection
acquire!calls (from inner functions or parent scope) could corrupt arrays on rewind when using typed checkpoint/rewind_untracked_flagsper depth level; automatically fallback to full checkpoint/rewind when detected2. 1-Based Sentinel Pattern
isempty()checks in hot paths_checkpoint_n_active=[0],_checkpoint_depths=[0]always have sentinel valuesacquire!calls (branch elimination in_mark_untracked!)3. Macro Transformation
acquire!calls inside@with_poolare transformed to_acquire_impl!acquire!→ properly marked as untracked4. Code Organization
core.jlintoacquire.jl(~410 lines) andstate.jl(~240 lines)5. Other Improvements
TypedPoolinside@with_poolacquire_view!,acquire_array!aliases in type extraction