From ce771412be8b02092e48efeeb2d676ace4734d07 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 26 Jun 2026 10:13:24 -0700 Subject: [PATCH 1/5] (feat): Add rand!/randn! random-array pool constructors Add rand!/randn! convenience constructors as peers of zeros!/ones!, extending Random.rand!/randn! with pool-backed methods. They are re-exported (same binding), so `using AdaptiveArrayPools` exposes them with no conflict alongside `using Random`. - Random added as a hard dependency (stdlib, ~0 load cost) - Forms: rand!(pool, [T,] dims...) and the collection/range sampling form rand!(pool, S, dims...) (e.g. rand!(pool, 1:6, n)); randn! for N(0,1) - Wired into @with_pool at all three sites (type pre-registration, call->_impl! transform, escape detection). The collection form self-records its eltype touch since the macro cannot pre-register eltype(S) - DisabledPool{:cpu} fallbacks; CPU-only (GPU parity deferred) - Seed-equivalent with plain rand/randn: pool arrays are genuine Arrays, so RNG consumption matches exactly Tests (test/test_random.jl): seed-equivalence incl. consecutive-draw and @with_pool RNG-neutrality, zero-allocation, type-stability (@inferred), @maybe_with_pool fallback, escape detection, and DisabledPool. --- Project.toml | 2 + README.md | 4 +- docs/design/rand_randn_pool_constructors.md | 174 ++++++ src/AdaptiveArrayPools.jl | 6 + src/convenience.jl | 158 +++++ src/macros.jl | 21 +- test/Project.toml | 1 + test/runtests.jl | 2 + test/test_random.jl | 621 ++++++++++++++++++++ 9 files changed, 986 insertions(+), 3 deletions(-) create mode 100644 docs/design/rand_randn_pool_constructors.md create mode 100644 test/test_random.jl diff --git a/Project.toml b/Project.toml index 9bc48731..76a76145 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "0.3.6" [deps] Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" [weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" @@ -20,4 +21,5 @@ CUDA = "5" Metal = "1" Preferences = "1" Printf = "1" +Random = "1" julia = "1.10" diff --git a/README.md b/README.md index 2700e840..20d50314 100644 --- a/README.md +++ b/README.md @@ -95,10 +95,12 @@ Common initialization patterns have convenience functions: |----------|---------------| | `zeros!(pool, 10)` | `acquire!` + `fill!(0)` | | `ones!(pool, Float32, 3, 3)` | `acquire!` + `fill!(1)` | +| `rand!(pool, 10)` | `acquire!` + `Random.rand!` (uniform; `rand!(pool, 1:6, n)` samples a range/collection) | +| `randn!(pool, 10)` | `acquire!` + `Random.randn!` (standard normal) | | `similar!(pool, A)` | `acquire!` matching `eltype(A)`, `size(A)` | | `reshape!(pool, A, 3, 4)` | Reshape sharing memory, zero-alloc (1.12+) | -All convenience functions return `Array` types, same as `acquire!`. See [API Reference](https://projecttorreypines.github.io/AdaptiveArrayPools.jl/stable/usage/api#convenience-functions). +All convenience functions return `Array` types, same as `acquire!`. `rand!`/`randn!` extend `Random.rand!`/`randn!`, so they consume the global RNG exactly like `rand`/`randn` (seed with `Random.seed!` for reproducibility). See [API Reference](https://projecttorreypines.github.io/AdaptiveArrayPools.jl/stable/usage/api#convenience-functions). ## Installation diff --git a/docs/design/rand_randn_pool_constructors.md b/docs/design/rand_randn_pool_constructors.md new file mode 100644 index 00000000..79078c0d --- /dev/null +++ b/docs/design/rand_randn_pool_constructors.md @@ -0,0 +1,174 @@ +# Design: `rand!` / `randn!` pool constructors + +**Status:** Approved design, pending implementation +**Branch:** `feature/rand-randn-pool-constructors` +**Author:** Min-Gu Yoo + +## 1. Goal + +Add random-array convenience constructors to the pool API, peers of `zeros!` / +`ones!`: + +```julia +@with_pool pool begin + v = rand!(pool, 100) # Vector{Float64}, U[0,1) + m = rand!(pool, Float32, 8, 8) # Matrix{Float32}, U[0,1) + d = rand!(pool, 1:6, 10) # Vector{Int}, each ∈ 1:6 (dice) + g = randn!(pool, 100) # Vector{Float64}, N(0,1) +end +``` + +These acquire a right-sized array from the pool and fill it in place with random +values, preserving the package's zero-allocation guarantee inside `@with_pool`. + +## 2. Key decisions (settled during brainstorming) + +| Decision | Choice | Rationale | +|---|---|---| +| Dependency | **`Random` as a hard `[deps]`** | `Random` is a stdlib (in the sysimage); `using Random` adds ~0 load cost. The usual reason to use a weakdep/extension (avoid forcing a heavy dep — cf. CUDA/Metal) does not apply. | +| Why not an extension | The convenience functions are **wired into `@with_pool` at 3 sites**; an extension would force a core stub + `GlobalRef` to a conditionally-defined impl and yield cryptic `MethodError`s when `Random` isn't loaded. Hard-dep keeps `_rand_impl!` next to `_zeros_impl!`. | +| Namespace | **`import Random: rand!, randn!` then `export rand!, randn!`** (re-export) | Makes `rand!` a true peer of `zeros!` — works with just `using AdaptiveArrayPools`. Re-exporting the *same binding* means **no conflict warning** if the user also does `using Random`. | +| Distributions (v1) | **`rand!` (uniform) + `randn!` (normal)** | Covers the overwhelming majority of use. `randexp!` deferred (low demand, trivial to add later via same machinery). | +| Collection/range form | **Include `rand!(pool, S, dims)`** (CPU-only) | This is the real "random integers in a range" case (`rand!(pool, 1:6, n)`). More useful than full-range `rand!(pool, Int, n)`. | +| GPU parity | **Deferred to a follow-up PR** | Live GPU pools inherit the core impl via dispatch (see §7), so it is cheap, but Metal's `randn!`/`rand!` MPS coverage needs verification. Kept out of v1 to ship CPU cleanly. | + +## 3. Why this is "wide but shallow" + +The implementation body is ~2 lines per function (`acquire!` → `Random.rand!`). +The work is **parallelism with `zeros!`**: every site that already knows about +`zeros!` must learn `rand!` / `randn!`. This is a consequence of the package +centralizing checkpoint / escape-analysis / type-registration logic in the +macro — new convenience functions are cheap to write but must be *registered* in +the macro. The spec below is therefore primarily a **checklist of mirror-sites**. + +## 4. Public API surface (CPU, v1) + +Mirrors `zeros!`'s four forms, plus the collection form for `rand!`: + +```julia +# Uniform +rand!(pool, dims...) # default_eltype(pool) +rand!(pool, T::Type, dims...) # typed +rand!(pool, dims::NTuple{N,Int}) # tuple form (for macro size(x) handling) +rand!(pool, T::Type, dims::NTuple{N,Int}) +rand!(pool, S, dims...) # sample from collection S; eltype = eltype(S) +rand!(pool, S, dims::NTuple{N,Int}) + +# Normal (T must be AbstractFloat; non-float T inherits Random.randn!'s error) +randn!(pool, dims...) +randn!(pool, T::Type, dims...) +randn!(pool, dims::NTuple{N,Int}) +randn!(pool, T::Type, dims::NTuple{N,Int}) +``` + +**Default eltype** comes from `default_eltype(pool)` (CPU: `Float64`), so +`rand!(pool, 10)` → `U[0,1)` Float64 and `randn!(pool, 10)` → `N(0,1)` Float64. + +**Collection form semantics**: `rand!(pool, S, dims...)` acquires an array of +`eltype(S)` and fills it via `Random.rand!(arr, S)`. Examples: +`rand!(pool, 1:6, 10)` (Int dice), `rand!(pool, ('a','b','c'), 5)` (Char). + +## 5. Components & files to change + +### 5.1 `Project.toml` +- Add `Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"` to `[deps]`. +- Add `Random = "1"` to `[compat]`. + +### 5.2 `src/AdaptiveArrayPools.jl` (main module) +- `import Random: rand!, randn!` (brings names in scope for extending + re-export). +- `import Random` (so `rand`, `randn` allocating fallbacks are reachable for `DisabledPool`). +- Add `rand!, randn!` to the convenience-functions `export` line (next to `zeros!, ones!, ...`). + +### 5.3 `src/convenience.jl` (next to `zeros!` / `_zeros_impl!`) +Public wrappers (each does `_record_type_touch!` + `_set_pending_callsite!` + impl, +identical structure to `zeros!`): +- `rand!(pool::AbstractArrayPool, ...)` — all six forms above. +- `randn!(pool::AbstractArrayPool, ...)` — four forms. + +Internal impls (backend-agnostic — rely on `Random.rand!`/`randn!` dispatch on the array type): +- `_rand_impl!(pool, T, dims...) = (arr = _acquire_impl!(pool, T, dims...); Random.rand!(arr); arr)` +- `_rand_impl!(pool, dims...) = _rand_impl!(pool, default_eltype(pool), dims...)` +- `_rand_impl!(pool, S, dims...) = (arr = _acquire_impl!(pool, eltype(S), dims...); Random.rand!(arr, S); arr)` *(collection)* +- NTuple overloads for each. +- `_randn_impl!(...)` analogous to `_rand_impl!` minus the collection form, using `Random.randn!`. + +**Dispatch disambiguation (collection vs typed vs default):** the collection +argument `S` is constrained to sampleable, non-`Type`, non-`Integer` collections — +`S::Union{AbstractArray, Tuple, AbstractRange, AbstractSet, AbstractString, AbstractDict}` +— so `rand!(pool, 1:6, 10)` picks the collection method while `rand!(pool, 10)` +and `rand!(pool, Int, 10)` pick the dims/typed methods. This mirrors how Base +separates `rand(coll, n)` from `rand(T, n)`. + +### 5.4 `src/convenience.jl` — `DisabledPool{:cpu}` fallbacks +Mirror the `zeros!` disabled-pool methods (`convenience.jl:439-442`): +- `rand!(::DisabledPool{:cpu}, T, dims...) = rand(T, dims...)` +- `rand!(p::DisabledPool{:cpu}, dims...) = rand(default_eltype(p), dims...)` +- `rand!(::DisabledPool{:cpu}, S, dims...) = rand(S, dims...)` *(collection)* +- `randn!(::DisabledPool{:cpu}, T, dims...) = randn(T, dims...)`, default + NTuple forms. + +### 5.5 `src/macros.jl` — the three mirror-sites +1. **Type extraction** `_extract_acquire_types` (~L1143-1189): add `:rand!` / `:randn!`. + - Typed form: read the type token like `zeros!`/`ones!`. + - Default form: use `default_eltype(pool)` like `zeros!`. + - **Collection form (NEW code path):** emit `eltype(S)` as a runtime expression + (constant-folds for literal ranges) rather than reading a literal token. +2. **Call → impl transform** (~L1483-1537, GlobalRefs at L1491-1492, impl-name list at L1584): + - Add `_RAND_IMPL_REF = GlobalRef(@__MODULE__, :_rand_impl!)` and `_RANDN_IMPL_REF`. + - Map `:rand! → _RAND_IMPL_REF`, `:randn! → _RANDN_IMPL_REF` (both the `fn ==` and `QuoteNode` branches). + - Add `:_rand_impl!`, `:_randn_impl!` to the impl-name list (~L1584). + - The transform only swaps the function name; the collection arg passes through unchanged. +3. **Escape-detection symbol lists** (~L1943-1973, and the lists feeding L2032/L2371): + - Add `:rand!`, `:randn!` so `x = rand!(pool, 10); return x` triggers `PoolEscapeError`, + matching `zeros!` behavior. + +## 6. Error handling + +- `randn!(pool, Int, n)` — inherits `Random.randn!`'s error for non-float eltypes + (no special handling; documented). +- Collection form with an eltype the pool can't store — surfaces through the + normal `acquire!` type path (`_record_type_touch!(pool, eltype(S))`). +- Escape of a pool-backed random array — `PoolEscapeError`, via the macro escape + lists (§5.5.3), identical to `zeros!`. + +## 7. GPU note (out of scope for v1, documented for the follow-up) + +`CuAdaptiveArrayPool{S} <: AbstractArrayPool`, and the CUDA extension does **not** +override `_zeros_impl!` for live pools — it inherits the core impl, which works on +GPU because CUDA.jl overloads `fill!` for `CuArray`. By the same mechanism, the +core `_rand_impl! = Random.rand!(arr)` would inherit GPU support **if** +`Random.rand!(::CuArray)` runs on-device. Therefore the GPU follow-up is mostly: +- `DisabledPool{:cuda}` / `{:metal}` fallback methods (`rand!(::DisabledPool{:cuda}, …) = CUDA.rand(…)`), and +- **verification** that `Random.rand!`/`randn!` on `CuArray`/`MtlArray` are GPU-native + (no scalar fallback under `allowscalar(false)`), with **Metal `randn!` coverage explicitly unconfirmed**. + +The **collection form is intentionally CPU-only** — GPU RNGs have no efficient +arbitrary-collection sampler and would scalar-fallback. + +## 8. Testing (`test/test_convenience.jl`, mirroring the `zeros!` testsets) + +- `rand!`: length/size & eltype correct for default, typed (`Float32`, `Int`), and + multi-dim forms; **all values ∈ [0,1)** for float eltypes. +- `rand!(pool, 1:6, n)`: `eltype == Int`, **all values ∈ 1:6**; a Char/tuple collection case. +- `randn!`: length/size & eltype correct; values finite (statistical sanity, not exact equality). +- **Zero-allocation** inside `@with_pool` (reuse the existing alloc-test helper / `test_zero_allocation.jl` pattern). +- **Escape detection**: `x = rand!(pool, 10); return x` triggers `PoolEscapeError` + (mirror the existing `test_compile_escape.jl` style). +- **`DisabledPool` fallback**: returns a plain `Array` of correct type/size. +- **Determinism**: seed via `Random.seed!` (no `rng` arg in v1) and assert reproducible values. +- **Aqua** (`test_aqua.jl`): confirm no new ambiguities, no piracy flag (methods are + on our `AbstractArrayPool` type → not piracy), deps/compat consistent with the new `Random` entry. + +## 9. Documentation + +- Docstrings on `rand!` / `randn!` mirroring the `zeros!` docstring (signatures, + default-eltype note, collection-form example, `See also`). +- README / docs convenience-function section: add `rand!` / `randn!` next to `zeros!` / `ones!`. + +## 10. Explicitly out of scope (YAGNI) + +- GPU parity (separate PR — see §7). +- `randexp!` (trivial later add). +- `rng`-argument forms (`rand!(pool, rng, dims)`) — would complicate macro type extraction. +- `randperm!` / `randcycle!` / `shuffle!` (different semantics — permutation/in-place on a user array). +- `bitrand` / `rand!(pool, Bit, …)` — tied to the deferred `Bit`/`BitArray` handling. +- GPU collection sampling. diff --git a/src/AdaptiveArrayPools.jl b/src/AdaptiveArrayPools.jl index 026cda82..aa89cd86 100644 --- a/src/AdaptiveArrayPools.jl +++ b/src/AdaptiveArrayPools.jl @@ -1,11 +1,17 @@ module AdaptiveArrayPools using Printf +import Random +# Extend (and re-export) Random's rand!/randn! with pool-constructor methods. +# Re-exporting the SAME binding means no conflict warning when a user also +# does `using Random` (see `rand!`/`randn!` in convenience.jl). +import Random: rand!, randn! # Public API export AdaptiveArrayPool, acquire!, pool_stats, get_task_local_pool export acquire_view!, acquire_array! # Explicit naming variants export zeros!, ones!, trues!, falses!, similar!, reshape!, default_eltype # Convenience functions +export rand!, randn! # Random-array convenience constructors (re-exported from Random) export Bit # Sentinel type for BitArray (use with acquire!, trues!, falses!) export @with_pool, @maybe_with_pool, @safe_with_pool, @safe_maybe_with_pool export STATIC_POOLING, MAYBE_POOLING, RUNTIME_CHECK diff --git a/src/convenience.jl b/src/convenience.jl index 0cf7a5e7..d4b1d8d8 100644 --- a/src/convenience.jl +++ b/src/convenience.jl @@ -254,6 +254,151 @@ end end @inline _falses_impl!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N} = _falses_impl!(pool, dims...) +# ============================================================================== +# rand! / randn! - Acquire random-initialized arrays from pool +# ============================================================================== +# These extend `Random.rand!` / `Random.randn!` with pool-constructor methods. +# They are NOT type piracy: the pool type (`AbstractArrayPool`) is ours, and the +# signatures are disjoint from Random's (which dispatch on `AbstractArray` first). +# +# Because `_acquire_impl!` returns a genuine `Array{T,N}` (not a view), filling it +# with `Random.rand!(arr)` consumes the RNG stream identically to `rand(T, dims)` +# — so seeding identically reproduces the same sequence. + +# Sampleable collection types for the `rand!(pool, S, dims...)` form. Disjoint +# from `Int`/`Type` so it never collides with the dims/typed forms. The trailing +# `d1::Int` (≥1 explicit dim) in the vararg method also prevents a bare Int-tuple +# (used as dims, e.g. `rand!(pool, (8,9))`) from being mistaken for a collection. +const _SampleColl = Union{AbstractArray, Tuple, AbstractSet, AbstractDict, AbstractString} + +""" + rand!(pool, dims...) -> Array + rand!(pool, T, dims...) -> Array + rand!(pool, S, dims...) -> Array + rand!(pool, dims::Tuple) / rand!(pool, T, dims::Tuple) / rand!(pool, S, dims::Tuple) + +Acquire a uniform random array from the pool. + +- `rand!(pool, [T,] dims...)`: floats are `U[0,1)`, integers full-range. + Default element type depends on backend (CPU: `Float64`, CUDA: `Float32`). +- `rand!(pool, S, dims...)`: sample from a collection/range `S` + (e.g. `rand!(pool, 1:6, 10)` for dice). Element type is `eltype(S)`. + +Consumes the global RNG identically to `Random.rand!`, so seeding reproduces +the same sequence. Use `Random.seed!` / pass values through `Random` for control. + +## Example +```julia +@with_pool pool begin + v = rand!(pool, 100) # Vector{Float64}, U[0,1) + m = rand!(pool, Float32, 8, 8) # Matrix{Float32} + d = rand!(pool, 1:6, 10) # Vector{Int}, each ∈ 1:6 +end +``` + +See also: [`randn!`](@ref), [`zeros!`](@ref), [`acquire!`](@ref) +""" +@inline function rand!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N} + _record_type_touch!(pool, T) + _set_pending_callsite!(pool, "") + return _rand_impl!(pool, T, dims...) +end +@inline function rand!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N} + _record_type_touch!(pool, default_eltype(pool)) + _set_pending_callsite!(pool, "") + return _rand_impl!(pool, default_eltype(pool), dims...) +end +@inline function rand!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N} + _record_type_touch!(pool, T) + _set_pending_callsite!(pool, "") + return _rand_impl!(pool, T, dims...) +end +@inline function rand!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N} + _record_type_touch!(pool, default_eltype(pool)) + _set_pending_callsite!(pool, "") + return _rand_impl!(pool, default_eltype(pool), dims...) +end +# Collection/range form: type touch is recorded inside `_rand_impl!` (the macro +# cannot pre-register `eltype(S)`, so the impl is self-sufficient — see below). +@inline function rand!(pool::AbstractArrayPool, S::_SampleColl, d1::Int, dims::Vararg{Int, N}) where {N} + _set_pending_callsite!(pool, "") + return _rand_impl!(pool, S, d1, dims...) +end +@inline function rand!(pool::AbstractArrayPool, S::_SampleColl, dims::NTuple{N, Int}) where {N} + _set_pending_callsite!(pool, "") + return _rand_impl!(pool, S, dims) +end + +# Internal implementations (called directly by macro-transformed code). +@inline function _rand_impl!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N} + arr = _acquire_impl!(pool, T, dims...) + Random.rand!(arr) + return arr +end +@inline _rand_impl!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N} = _rand_impl!(pool, default_eltype(pool), dims...) +@inline _rand_impl!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = _rand_impl!(pool, T, dims...) +@inline _rand_impl!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N} = _rand_impl!(pool, default_eltype(pool), dims...) +# Collection form self-records its eltype touch: the macro registers the wrong +# (default) type for `rand!(pool, S, dims)` because `S` is not a syntactic type, +# so this impl records `eltype(S)` itself to keep checkpoint/rewind correct. +@inline function _rand_impl!(pool::AbstractArrayPool, S::_SampleColl, d1::Int, dims::Vararg{Int, N}) where {N} + T = eltype(S) + _record_type_touch!(pool, T) + arr = _acquire_impl!(pool, T, d1, dims...) + Random.rand!(arr, S) + return arr +end +@inline _rand_impl!(pool::AbstractArrayPool, S::_SampleColl, dims::NTuple{N, Int}) where {N} = _rand_impl!(pool, S, dims...) + +""" + randn!(pool, dims...) -> Array + randn!(pool, T, dims...) -> Array + randn!(pool, dims::Tuple) / randn!(pool, T, dims::Tuple) + +Acquire a standard-normal `N(0, 1)` random array from the pool. `T` must be a +floating-point (or complex) type. Default element type depends on backend +(CPU: `Float64`). Consumes the global RNG identically to `Random.randn!`. + +## Example +```julia +@with_pool pool begin + g = randn!(pool, 100) # Vector{Float64}, N(0,1) + m = randn!(pool, Float32, 8, 8) # Matrix{Float32} +end +``` + +See also: [`rand!`](@ref), [`zeros!`](@ref), [`acquire!`](@ref) +""" +@inline function randn!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N} + _record_type_touch!(pool, T) + _set_pending_callsite!(pool, "") + return _randn_impl!(pool, T, dims...) +end +@inline function randn!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N} + _record_type_touch!(pool, default_eltype(pool)) + _set_pending_callsite!(pool, "") + return _randn_impl!(pool, default_eltype(pool), dims...) +end +@inline function randn!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N} + _record_type_touch!(pool, T) + _set_pending_callsite!(pool, "") + return _randn_impl!(pool, T, dims...) +end +@inline function randn!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N} + _record_type_touch!(pool, default_eltype(pool)) + _set_pending_callsite!(pool, "") + return _randn_impl!(pool, default_eltype(pool), dims...) +end + +@inline function _randn_impl!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N} + arr = _acquire_impl!(pool, T, dims...) + Random.randn!(arr) + return arr +end +@inline _randn_impl!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N} = _randn_impl!(pool, default_eltype(pool), dims...) +@inline _randn_impl!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = _randn_impl!(pool, T, dims...) +@inline _randn_impl!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N} = _randn_impl!(pool, default_eltype(pool), dims...) + # ============================================================================== # similar! - Acquire arrays with same type/size as template # ============================================================================== @@ -459,6 +604,19 @@ end @inline falses!(::DisabledPool{:cpu}, dims::Vararg{Int, N}) where {N} = falses(dims...) @inline falses!(::DisabledPool{:cpu}, dims::NTuple{N, Int}) where {N} = falses(dims...) +# --- rand!/randn! for DisabledPool{:cpu} (allocating, plain Array) --- +@inline rand!(::DisabledPool{:cpu}, ::Type{T}, dims::Vararg{Int, N}) where {T, N} = rand(T, dims...) +@inline rand!(p::DisabledPool{:cpu}, dims::Vararg{Int, N}) where {N} = rand(default_eltype(p), dims...) +@inline rand!(::DisabledPool{:cpu}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = rand(T, dims...) +@inline rand!(p::DisabledPool{:cpu}, dims::NTuple{N, Int}) where {N} = rand(default_eltype(p), dims...) +@inline rand!(::DisabledPool{:cpu}, S::_SampleColl, d1::Int, dims::Vararg{Int, N}) where {N} = rand(S, d1, dims...) +@inline rand!(::DisabledPool{:cpu}, S::_SampleColl, dims::NTuple{N, Int}) where {N} = rand(S, dims) + +@inline randn!(::DisabledPool{:cpu}, ::Type{T}, dims::Vararg{Int, N}) where {T, N} = Random.randn(T, dims...) +@inline randn!(p::DisabledPool{:cpu}, dims::Vararg{Int, N}) where {N} = Random.randn(default_eltype(p), dims...) +@inline randn!(::DisabledPool{:cpu}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = Random.randn(T, dims...) +@inline randn!(p::DisabledPool{:cpu}, dims::NTuple{N, Int}) where {N} = Random.randn(default_eltype(p), dims...) + # --- similar! for DisabledPool{:cpu} --- @inline similar!(::DisabledPool{:cpu}, x::AbstractArray) = similar(x) @inline similar!(::DisabledPool{:cpu}, x::AbstractArray, ::Type{T}) where {T} = similar(x, T) diff --git a/src/macros.jl b/src/macros.jl index 99e310a9..11eed70d 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -1185,8 +1185,12 @@ function _extract_acquire_types(expr, target_pool, types = Set{Any}()) # trues!/falses! (always uses Bit type) elseif fn in (:trues!, :falses!) || fn_name in (:trues!, :falses!) push!(types, :Bit) - # zeros!/ones! - elseif fn in (:zeros!, :ones!) || fn_name in (:zeros!, :ones!) + # zeros!/ones!/rand!/randn! + # NOTE: for rand!(pool, S, dims) (collection/range form), the + # third arg S is not a syntactic type, so the default_eltype + # branch is taken here (a harmless over-registration). The real + # eltype(S) is recorded at runtime inside `_rand_impl!`. + elseif fn in (:zeros!, :ones!, :rand!, :randn!) || fn_name in (:zeros!, :ones!, :rand!, :randn!) if nargs >= 3 third_arg = expr.args[3] # Check if third arg looks like a type (Symbol starting with uppercase or curly) @@ -1494,6 +1498,8 @@ const _TRUES_IMPL_REF = GlobalRef(@__MODULE__, :_trues_impl!) const _FALSES_IMPL_REF = GlobalRef(@__MODULE__, :_falses_impl!) const _SIMILAR_IMPL_REF = GlobalRef(@__MODULE__, :_similar_impl!) const _RESHAPE_IMPL_REF = GlobalRef(@__MODULE__, :_reshape_impl!) +const _RAND_IMPL_REF = GlobalRef(@__MODULE__, :_rand_impl!) +const _RANDN_IMPL_REF = GlobalRef(@__MODULE__, :_randn_impl!) function _transform_acquire_calls(expr, pool_name) if expr isa Expr @@ -1521,6 +1527,10 @@ function _transform_acquire_calls(expr, pool_name) expr = Expr(:call, _SIMILAR_IMPL_REF, expr.args[2:end]...) elseif fn == :reshape! expr = Expr(:call, _RESHAPE_IMPL_REF, expr.args[2:end]...) + elseif fn == :rand! + expr = Expr(:call, _RAND_IMPL_REF, expr.args[2:end]...) + elseif fn == :randn! + expr = Expr(:call, _RANDN_IMPL_REF, expr.args[2:end]...) elseif fn isa Expr && fn.head == :. && length(fn.args) >= 2 # Qualified name: AdaptiveArrayPools.acquire! etc. qn = fn.args[end] @@ -1540,6 +1550,10 @@ function _transform_acquire_calls(expr, pool_name) expr = Expr(:call, _SIMILAR_IMPL_REF, expr.args[2:end]...) elseif qn == QuoteNode(:reshape!) expr = Expr(:call, _RESHAPE_IMPL_REF, expr.args[2:end]...) + elseif qn == QuoteNode(:rand!) + expr = Expr(:call, _RAND_IMPL_REF, expr.args[2:end]...) + elseif qn == QuoteNode(:randn!) + expr = Expr(:call, _RANDN_IMPL_REF, expr.args[2:end]...) end end end @@ -1583,6 +1597,7 @@ const _IMPL_FUNC_NAMES = Set{Symbol}( :_acquire_impl!, :_acquire_view_impl!, :_zeros_impl!, :_ones_impl!, :_trues_impl!, :_falses_impl!, :_similar_impl!, :_reshape_impl!, + :_rand_impl!, :_randn_impl!, ] ) @@ -1942,6 +1957,7 @@ const _ALL_ACQUIRE_NAMES = Set{Symbol}( :acquire!, :acquire_view!, :acquire_array!, :zeros!, :ones!, :similar!, :reshape!, :trues!, :falses!, + :rand!, :randn!, ] ) @@ -1957,6 +1973,7 @@ const _ARRAY_ACQUIRE_NAMES = Set{Symbol}( [ :acquire!, :acquire_array!, :zeros!, :ones!, :similar!, :reshape!, + :rand!, :randn!, ] ) diff --git a/test/Project.toml b/test/Project.toml index fcdde041..3f8edd72 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,4 +2,5 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/runtests.jl b/test/runtests.jl index 74042cce..8c418c96 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,6 +41,7 @@ else include("test_fixed_slots.jl") include("test_backend_macro_expansion.jl") include("test_convenience.jl") + include("test_random.jl") include("test_reshape.jl") include("test_bitarray.jl") include("test_coverage.jl") @@ -68,6 +69,7 @@ else include("test_fixed_slots.jl") include("test_backend_macro_expansion.jl") include("test_convenience.jl") + include("test_random.jl") include("test_reshape.jl") include("test_bitarray.jl") include("test_coverage.jl") diff --git a/test/test_random.jl b/test/test_random.jl new file mode 100644 index 00000000..85dc8c12 --- /dev/null +++ b/test/test_random.jl @@ -0,0 +1,621 @@ +# ============================================================================== +# Tests for rand! / randn! pool constructors (random-array convenience functions) +# ============================================================================== +# Covers: +# - basic random properties (shape, eltype, value range / membership) +# - EXACT seed-equivalence with Base/Random (our rand!/randn! must consume the +# RNG stream identically to rand()/randn(), incl. consecutive-draw and +# @with_pool RNG-neutrality) +# - @with_pool integration (macro transform, size(x) form, rewind reuse) +# - zero-allocation for the 0-alloc forms (Float64 rand!/randn!, range rand!) +# - escape detection, DisabledPool fallback, error handling, namespace re-export +# +# NOTE: seed-equivalence checks compare INSIDE the @with_pool scope and return a +# Bool — the pool array never crosses the scope boundary (avoids tangling with +# escape detection / slot-reuse aliasing). + +using Random +using AdaptiveArrayPools: _extract_acquired_vars + +@testset "Random Constructors (rand!/randn!)" begin + + # ========================================================================= + # 1. Namespace: re-export and binding identity (no conflict with Random) + # ========================================================================= + @testset "re-export & binding identity" begin + @test :rand! in names(AdaptiveArrayPools) + @test :randn! in names(AdaptiveArrayPools) + # Re-exported the SAME binding as Random → no conflict when both loaded + @test AdaptiveArrayPools.rand! === Random.rand! + @test AdaptiveArrayPools.randn! === Random.randn! + end + + # ========================================================================= + # 2. rand! — basic properties (shape, eltype, value range) + # ========================================================================= + @testset "rand! properties" begin + pool = AdaptiveArrayPool() + + @testset "explicit type" begin + v = rand!(pool, Float64, 1000) + @test v isa Array{Float64, 1} + @test length(v) == 1000 + @test all(x -> 0.0 <= x < 1.0, v) + + v32 = rand!(pool, Float32, 500) + @test eltype(v32) == Float32 + @test all(x -> 0.0f0 <= x < 1.0f0, v32) + end + + @testset "default type (Float64)" begin + v = rand!(pool, 1000) + @test v isa Array{Float64, 1} + @test all(x -> 0.0 <= x < 1.0, v) + end + + @testset "multi-dimensional" begin + m = rand!(pool, Float64, 30, 40) + @test size(m) == (30, 40) + @test all(x -> 0.0 <= x < 1.0, m) + + t = rand!(pool, Float32, 4, 5, 6) + @test size(t) == (4, 5, 6) + end + + @testset "tuple form" begin + m = rand!(pool, (8, 9)) + @test size(m) == (8, 9) + @test eltype(m) == Float64 + + m32 = rand!(pool, Float32, (3, 4)) + @test size(m32) == (3, 4) + @test eltype(m32) == Float32 + end + + @testset "returns Array not view" begin + v = rand!(pool, Float64, 10) + @test v isa Array{Float64, 1} + @test !(v isa SubArray) + end + end + + # ========================================================================= + # 3. rand! — collection/range sampling form + # ========================================================================= + @testset "rand! collection form" begin + pool = AdaptiveArrayPool() + + @testset "integer range (dice)" begin + Random.seed!(101) + v = rand!(pool, 1:6, 1000) + @test eltype(v) == Int + @test length(v) == 1000 + @test all(x -> x in 1:6, v) + @test Set(v) == Set(1:6) # all faces appear in 1000 draws + end + + @testset "multi-dim range" begin + m = rand!(pool, 1:6, 20, 30) + @test size(m) == (20, 30) + @test all(x -> x in 1:6, m) + end + + @testset "non-integer collection" begin + chars = ('a', 'b', 'c') + v = rand!(pool, chars, 200) + @test eltype(v) == Char + @test all(x -> x in chars, v) + + floats = [1.5, 2.5, 3.5] + vf = rand!(pool, floats, 200) + @test eltype(vf) == Float64 + @test all(x -> x in floats, vf) + end + + @testset "tuple-dims collection form" begin + m = rand!(pool, 1:6, (5, 5)) + @test size(m) == (5, 5) + @test all(x -> x in 1:6, m) + end + + @testset "returns Array not view" begin + v = rand!(pool, 1:6, 10) + @test v isa Array{Int, 1} + @test !(v isa SubArray) + end + end + + # ========================================================================= + # 4. randn! — basic properties + # ========================================================================= + @testset "randn! properties" begin + pool = AdaptiveArrayPool() + + @testset "explicit type" begin + v = randn!(pool, Float64, 1000) + @test v isa Array{Float64, 1} + @test length(v) == 1000 + @test all(isfinite, v) + + v32 = randn!(pool, Float32, 500) + @test eltype(v32) == Float32 + @test all(isfinite, v32) + end + + @testset "default type (Float64)" begin + v = randn!(pool, 1000) + @test eltype(v) == Float64 + @test all(isfinite, v) + end + + @testset "multi-dim & tuple" begin + m = randn!(pool, Float64, 10, 20) + @test size(m) == (10, 20) + + t = randn!(pool, (4, 5)) + @test size(t) == (4, 5) + @test eltype(t) == Float64 + end + + @testset "returns Array not view" begin + v = randn!(pool, 10) + @test v isa Array{Float64, 1} + @test !(v isa SubArray) + end + end + + @testset "randn! distribution sanity (seeded)" begin + pool = AdaptiveArrayPool() + Random.seed!(20240601) + v = randn!(pool, Float64, 100_000) + m = sum(v) / length(v) + s = sqrt(sum(abs2, v .- m) / (length(v) - 1)) + @test abs(m) < 0.05 # mean ≈ 0 + @test abs(s - 1.0) < 0.05 # std ≈ 1 + end + + # ========================================================================= + # 5. Seed-equivalence with plain Base/Random (the headline property) + # ========================================================================= + @testset "seed equivalence with plain rand/randn" begin + + @testset "rand!(pool, T, n) == rand(T, n)" begin + for T in (Float64, Float32, Int) + Random.seed!(424242); ref = rand(T, 64) + pool = AdaptiveArrayPool() + Random.seed!(424242); got = copy(rand!(pool, T, 64)) + @test ref == got + end + end + + @testset "default eltype rand!(pool, n) == rand(Float64, n)" begin + Random.seed!(1); ref = rand(Float64, 50) + pool = AdaptiveArrayPool() + Random.seed!(1); got = copy(rand!(pool, 50)) + @test ref == got + end + + @testset "multi-dim equivalence" begin + Random.seed!(99); ref = rand(Float64, 7, 11) + pool = AdaptiveArrayPool() + Random.seed!(99); got = copy(rand!(pool, Float64, 7, 11)) + @test ref == got + end + + @testset "collection form rand!(pool, S, n) == rand(S, n)" begin + Random.seed!(7); ref = rand(1:6, 100) + pool = AdaptiveArrayPool() + Random.seed!(7); got = copy(rand!(pool, 1:6, 100)) + @test ref == got + end + + @testset "randn!(pool, n) == randn(Float64, n)" begin + Random.seed!(2024); ref = randn(Float64, 64) + pool = AdaptiveArrayPool() + Random.seed!(2024); got = copy(randn!(pool, 64)) + @test ref == got + end + + @testset "consecutive draws: no RNG over-consumption" begin + # If rand! filled more than n (e.g. the whole backing buffer), the + # SECOND draw would desync from the plain stream. This catches that. + Random.seed!(555) + r1 = rand(Float64, 5); r2 = rand(Float64, 9) + pool = AdaptiveArrayPool() + Random.seed!(555) + ok = @with_pool pool begin + a = rand!(pool, 5) + chk1 = (a == r1) + b = rand!(pool, 9) + chk2 = (b == r2) + chk1 && chk2 + end + @test ok + end + + @testset "@with_pool machinery does not touch the RNG" begin + # Acquiring / zeroing other arrays and entering the scope must not + # consume any random numbers. + Random.seed!(13) + ref = rand(Float64, 6) + pool = AdaptiveArrayPool() + Random.seed!(13) + ok = @with_pool pool begin + z = zeros!(pool, 100) # fill zeros — no RNG + w = acquire!(pool, Float64, 50) # acquire — no RNG + fill!(w, 1.0) + x = rand!(pool, 6) + x == ref + end + @test ok + end + + @testset "scope enter/exit (checkpoint/rewind) is RNG-neutral" begin + Random.seed!(31) + ref1 = rand(Float64, 5); ref2 = rand(Float64, 5) + pool = AdaptiveArrayPool() + Random.seed!(31) + ok1 = @with_pool pool begin + rand!(pool, 5) == ref1 + end + ok2 = @with_pool pool begin + rand!(pool, 5) == ref2 + end + @test ok1 + @test ok2 + end + + @testset "reproducibility & seed sensitivity" begin + pool = AdaptiveArrayPool() + Random.seed!(123); a = copy(rand!(pool, 32)) + Random.seed!(123); b = copy(rand!(pool, 32)) + @test a == b # same seed → same sequence + Random.seed!(456); c = copy(rand!(pool, 32)) + @test a != c # different seed → different sequence + end + end + + # ========================================================================= + # 6. @with_pool integration + # ========================================================================= + @testset "integration with @with_pool" begin + + @testset "rand!/randn! in macro return Array" begin + result = @with_pool pool begin + v = rand!(pool, Float64, 100) + @test v isa Array{Float64, 1} + @test !(v isa SubArray) + g = randn!(pool, 100) + @test g isa Array{Float64, 1} + length(v) + length(g) + end + @test result == 200 + end + + @testset "collection form in macro" begin + result = @with_pool pool begin + v = rand!(pool, 1:6, 100) + @test all(x -> x in 1:6, v) + sum(v) + end + @test 100 <= result <= 600 + end + + @testset "size(x) form (macro NTuple transform)" begin + x1 = rand(10); x2 = rand(5, 8) + @with_pool pool begin + v = rand!(pool, Float64, size(x1)) + @test size(v) == (10,) + m = rand!(pool, size(x2)) + @test size(m) == (5, 8) + d = rand!(pool, 1:6, size(x1)) + @test size(d) == (10,) + @test all(y -> y in 1:6, d) + nothing + end + end + + @testset "rewind reuse gives fresh, seed-equivalent values" begin + pool = AdaptiveArrayPool() + Random.seed!(0) + got = Vector{Float64}[] + for _ in 1:5 + @with_pool pool begin + x = rand!(pool, 4) + push!(got, copy(x)) + nothing + end + end + @test allunique(got) # each iteration fresh (not stuck) + Random.seed!(0) + ref = [rand(Float64, 4) for _ in 1:5] + @test got == ref # rewind-reuse is seed-equivalent + end + + @testset "pool state: n_active increments" begin + pool = AdaptiveArrayPool() + checkpoint!(pool) + rand!(pool, Float64, 10) + randn!(pool, Float64, 10) + @test pool.float64.n_active == 2 + rewind!(pool) + @test pool.float64.n_active == 0 + end + end + + # ========================================================================= + # 7. Zero-allocation (function barrier + warmup), for the 0-alloc forms + # ========================================================================= + @testset "zero allocation" begin + function _alloc_rand() + pool = AdaptiveArrayPool() + for _ in 1:3 + @with_pool pool begin + a = rand!(pool, Float64, 100) + b = randn!(pool, Float64, 100) + c = rand!(pool, 1:6, 100) + sum(a) + sum(b) + sum(c) + end + end + @allocated @with_pool pool begin + a = rand!(pool, Float64, 100) + b = randn!(pool, Float64, 100) + c = rand!(pool, 1:6, 100) + sum(a) + sum(b) + sum(c) + end + end + _alloc_rand(); _alloc_rand() + alloc = _alloc_rand() + println(" rand!/randn!/rand!(1:6) (single iter): $alloc bytes") + @test alloc == 0 + end + + # ========================================================================= + # 8. Escape detection: vars assigned from rand!/randn! are tracked + # ========================================================================= + @testset "escape detection tracks rand!/randn!" begin + vars = _extract_acquired_vars( + quote + r = rand!(pool, 5) + g = randn!(pool, Float64, 5) + d = rand!(pool, 1:6, 5) + end, :pool + ) + @test :r in vars + @test :g in vars + @test :d in vars + end + + # ========================================================================= + # 9. DisabledPool fallback (allocating, plain Array) + # ========================================================================= + @testset "DisabledPool fallback" begin + v = rand!(DISABLED_CPU, Float64, 10) + @test v isa Array{Float64, 1} + @test all(x -> 0.0 <= x < 1.0, v) + + v2 = rand!(DISABLED_CPU, 5, 5) + @test v2 isa Matrix{Float64} + @test size(v2) == (5, 5) + + vd = rand!(DISABLED_CPU, 1:6, 20) + @test vd isa Array{Int, 1} + @test all(x -> x in 1:6, vd) + + g = randn!(DISABLED_CPU, Float64, 10) + @test g isa Array{Float64, 1} + @test all(isfinite, g) + + v3 = rand!(DISABLED_CPU, Float32, (3, 4)) # NTuple fallback + @test v3 isa Array{Float32, 2} + @test size(v3) == (3, 4) + end + + @testset "DisabledPool seed-equivalence" begin + Random.seed!(77); ref = rand(Float64, 16) + Random.seed!(77); got = rand!(DISABLED_CPU, Float64, 16) + @test ref == got + end + + # ========================================================================= + # 10. Error handling + # ========================================================================= + @testset "randn! requires a float eltype" begin + pool = AdaptiveArrayPool() + @test_throws Exception randn!(pool, Int, 5) + end + + # ========================================================================= + # 11. Seed-equivalence across more eltypes (regression guard for the + # headline invariant — Float16/Complex consume the RNG via different + # per-type codepaths than Float64). + # ========================================================================= + @testset "seed equivalence: extra eltypes" begin + @testset "rand! across eltypes" begin + for T in (Float16, ComplexF64, ComplexF32, Bool) + Random.seed!(11); ref = rand(T, 48) + pool = AdaptiveArrayPool() + Random.seed!(11); got = copy(rand!(pool, T, 48)) + @test eltype(got) == T + @test ref == got + end + end + @testset "randn! across float/complex eltypes" begin + for T in (Float32, Float16, ComplexF64, ComplexF32) + Random.seed!(2024); ref = randn(T, 64) + pool = AdaptiveArrayPool() + Random.seed!(2024); got = copy(randn!(pool, T, 64)) + @test eltype(got) == T + @test ref == got + end + end + end + + # ========================================================================= + # 12. Collection-form seed-equivalence: multi-dim, tuple-dims, non-Int + # collections — each is a distinct dispatch path from the 1-D vararg form. + # ========================================================================= + @testset "seed equivalence: collection-form shapes" begin + pool = AdaptiveArrayPool() + + Random.seed!(7); ref = rand(1:6, 7, 11) # multi-dim vararg + Random.seed!(7); got = copy(rand!(pool, 1:6, 7, 11)) + @test ref == got + + Random.seed!(8); ref2 = rand(1:6, (5, 5)) # tuple-dims (separate method) + Random.seed!(8); got2 = copy(rand!(pool, 1:6, (5, 5))) + @test ref2 == got2 + + chars = ('a', 'b', 'c', 'd') # non-Int collection + Random.seed!(9); ref3 = rand(chars, 40) + Random.seed!(9); got3 = copy(rand!(pool, chars, 40)) + @test ref3 == got3 + + vals = [1.5, 2.5, 3.5] # Vector collection + Random.seed!(10); ref4 = rand(vals, 40) + Random.seed!(10); got4 = copy(rand!(pool, vals, 40)) + @test ref4 == got4 + end + + # ========================================================================= + # 13. DisabledPool seed-equivalence: collection + randn! fallbacks + # (distinct methods from the typed rand! fallback). + # ========================================================================= + @testset "DisabledPool seed equivalence: collection & randn!" begin + Random.seed!(77); ref = rand(1:6, 50) + Random.seed!(77); got = rand!(DISABLED_CPU, 1:6, 50) + @test ref == got + + Random.seed!(88); ref2 = rand(1:6, (5, 5)) + Random.seed!(88); got2 = rand!(DISABLED_CPU, 1:6, (5, 5)) + @test ref2 == got2 + + Random.seed!(99); ref3 = randn(Float64, 16) + Random.seed!(99); got3 = randn!(DISABLED_CPU, Float64, 16) + @test ref3 == got3 + end + + # ========================================================================= + # 14. Empty (n=0) boundary: correct empty Array + ZERO RNG consumption. + # ========================================================================= + @testset "empty (n=0) request" begin + pool = AdaptiveArrayPool() + @test length(rand!(pool, Float64, 0)) == 0 + @test length(randn!(pool, 0)) == 0 + @test length(rand!(pool, 1:6, 0)) == 0 + @test eltype(rand!(pool, 1:6, 0)) == Int + + # An empty draw must consume zero random numbers (RNG-neutral). + Random.seed!(1); ref = rand(Float64, 4) + pool2 = AdaptiveArrayPool() + Random.seed!(1) + ok = @with_pool pool2 begin + rand!(pool2, Float64, 0) # empty: draws nothing + rand!(pool2, Float64, 4) == ref + end + @test ok + end + + # ========================================================================= + # 15. Collection-form type registration & reclamation. + # The Int touch is self-recorded inside `_rand_impl!` (NOT by the macro, + # which spuriously registers the default Float64). Guard that path: + # int64 must activate (not float64), and reclaim across @with_pool scopes. + # ========================================================================= + @testset "collection-form Int registration & reclamation" begin + pool = AdaptiveArrayPool() + checkpoint!(pool) + rand!(pool, 1:6, 10) + @test pool.int64.n_active == 1 # Int touch recorded + @test pool.float64.n_active == 0 # NOT the spurious default type + rewind!(pool) + @test pool.int64.n_active == 0 + + # Many @with_pool iterations (typed-lazy path) must not leak the Int pool; + # a regression in the self-record would grow int64.n_active across scopes. + tlp = get_task_local_pool() + base = tlp.int64.n_active + for _ in 1:500 + @with_pool p begin + d = rand!(p, 1:6, 64) + sum(d) + end + end + @test tlp.int64.n_active == base # no residual Int slots after scopes + end + + # ========================================================================= + # 16. Disambiguation: a bare Int-tuple is DIMS, not a sample collection. + # rand!(pool, (8,9)) → 8x9 array; rand!(pool, (8,9), k) → sample from {8,9}. + # ========================================================================= + @testset "tuple-as-dims vs tuple-as-collection" begin + m = rand!(AdaptiveArrayPool(), (8, 9)) # dims → 8x9 matrix + @test size(m) == (8, 9) + @test eltype(m) == Float64 + + v = rand!(AdaptiveArrayPool(), (8, 9), 100) # collection {8,9} → 100 samples + @test length(v) == 100 + @test all(x -> x in (8, 9), v) + end + + # ========================================================================= + # 17. Type stability — every form must @inferred to a concrete Array type + # (type stability is what makes the zero-allocation path possible). + # ========================================================================= + @testset "type stability (@inferred)" begin + pool = AdaptiveArrayPool() + @test (@inferred rand!(pool, Float64, 10)) isa Vector{Float64} + @test (@inferred rand!(pool, 10)) isa Vector{Float64} + @test (@inferred rand!(pool, Float64, 3, 4)) isa Matrix{Float64} + @test (@inferred rand!(pool, Float32, 10)) isa Vector{Float32} + @test (@inferred rand!(pool, 1:6, 10)) isa Vector{Int} # eltype(S) inferred + @test (@inferred rand!(pool, 1:6, (5, 5))) isa Matrix{Int} + @test (@inferred randn!(pool, Float64, 10)) isa Vector{Float64} + @test (@inferred randn!(pool, 10)) isa Vector{Float64} + @test (@inferred randn!(pool, Float32, 3, 4)) isa Matrix{Float32} + end + + # ========================================================================= + # 18. @maybe_with_pool: when MAYBE_POOLING[]=false, rand!/randn! must revert + # to plain rand/randn (plain Array, seed-equivalent) — safe with no pool. + # ========================================================================= + @testset "@maybe_with_pool reverts to rand when MAYBE_POOLING[]=false" begin + old = AdaptiveArrayPools.MAYBE_POOLING[] + AdaptiveArrayPools.MAYBE_POOLING[] = false + try + Random.seed!(5); ref = rand(Float64, 8) + Random.seed!(5) + got = @maybe_with_pool mp begin + copy(rand!(mp, Float64, 8)) + end + @test got isa Vector{Float64} # plain Array, not pool-backed + @test got == ref # seed-equivalent to plain rand + + Random.seed!(6); refd = rand(1:6, 10) + Random.seed!(6) + gotd = @maybe_with_pool mp begin + copy(rand!(mp, 1:6, 10)) + end + @test gotd isa Vector{Int} + @test gotd == refd + + Random.seed!(7); refn = randn(Float64, 8) + Random.seed!(7) + gotn = @maybe_with_pool mp begin + copy(randn!(mp, 8)) + end + @test gotn isa Vector{Float64} + @test gotn == refn + finally + AdaptiveArrayPools.MAYBE_POOLING[] = old + end + + # Sanity: with pooling ENABLED (default), the same call still works. + Random.seed!(5); ref = rand(Float64, 8) + Random.seed!(5) + ok = @maybe_with_pool mp begin + rand!(mp, Float64, 8) == ref + end + @test ok + end +end From e5cecc4cecf4363e64e9603102bc806ec80092e4 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 26 Jun 2026 11:36:43 -0700 Subject: [PATCH 2/5] (feat): Extend rand!/randn! to Metal and CUDA GPU pools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The live GPU pool inherits the core _rand_impl!/_randn_impl! (which fill the acquired CuArray/MtlArray via Random.rand!/randn!, run on-device by CUDA.jl / Metal.jl) — verified on Metal for Float32/Float16/Int32. Each extension adds: - DisabledPool{:cuda}/{:metal} rand!/randn! fallbacks (CUDA.rand/randn, Metal.rand/randn) - A clear ArgumentError for the collection/range form rand!(pool, S, dims): GPU RNGs have no arbitrary-collection sampler, so it is CPU-only. Overriding _rand_impl! covers both the @with_pool macro path and the direct path, for live and disabled pools. Tests: test/metal/test_random.jl (34 pass locally) and test/cuda/test_random.jl (mirror, to run on a CUDA machine). --- docs/design/rand_randn_pool_constructors.md | 34 +++--- ext/AdaptiveArrayPoolsCUDAExt/convenience.jl | 33 ++++++ ext/AdaptiveArrayPoolsMetalExt/convenience.jl | 33 ++++++ test/cuda/runtests.jl | 1 + test/cuda/test_random.jl | 103 ++++++++++++++++++ test/metal/runtests.jl | 1 + test/metal/test_random.jl | 101 +++++++++++++++++ 7 files changed, 292 insertions(+), 14 deletions(-) create mode 100644 test/cuda/test_random.jl create mode 100644 test/metal/test_random.jl diff --git a/docs/design/rand_randn_pool_constructors.md b/docs/design/rand_randn_pool_constructors.md index 79078c0d..ba1f78eb 100644 --- a/docs/design/rand_randn_pool_constructors.md +++ b/docs/design/rand_randn_pool_constructors.md @@ -130,19 +130,26 @@ Mirror the `zeros!` disabled-pool methods (`convenience.jl:439-442`): - Escape of a pool-backed random array — `PoolEscapeError`, via the macro escape lists (§5.5.3), identical to `zeros!`. -## 7. GPU note (out of scope for v1, documented for the follow-up) - -`CuAdaptiveArrayPool{S} <: AbstractArrayPool`, and the CUDA extension does **not** -override `_zeros_impl!` for live pools — it inherits the core impl, which works on -GPU because CUDA.jl overloads `fill!` for `CuArray`. By the same mechanism, the -core `_rand_impl! = Random.rand!(arr)` would inherit GPU support **if** -`Random.rand!(::CuArray)` runs on-device. Therefore the GPU follow-up is mostly: -- `DisabledPool{:cuda}` / `{:metal}` fallback methods (`rand!(::DisabledPool{:cuda}, …) = CUDA.rand(…)`), and -- **verification** that `Random.rand!`/`randn!` on `CuArray`/`MtlArray` are GPU-native - (no scalar fallback under `allowscalar(false)`), with **Metal `randn!` coverage explicitly unconfirmed**. - -The **collection form is intentionally CPU-only** — GPU RNGs have no efficient -arbitrary-collection sampler and would scalar-fallback. +## 7. GPU (Metal + CUDA) — implemented + +`CuAdaptiveArrayPool` / `MetalAdaptiveArrayPool <: AbstractArrayPool`, and the GPU +extensions do **not** override `_rand_impl!`/`_randn_impl!` for live pools — they +**inherit the core impl**, which works on GPU because `Random.rand!`/`randn!` are +overloaded for `CuArray`/`MtlArray` and run on-device. **Verified on Metal** +(Apple Silicon): `Random.rand!`/`randn!` on `MtlArray` are GPU-native for +`Float32`/`Float16`/`Int32`. So the live GPU pool gets the typed/default forms for +free, exactly like `zeros!`. + +GPU-specific code (per backend, in the extension): +- `DisabledPool{:cuda}`/`{:metal}` fallbacks (`Metal.rand`/`Metal.randn`, + `CUDA.rand`/`CUDA.randn`). +- **Collection-form rejection**: `rand!(pool, S, dims)` is **CPU-only** — GPU RNGs have + no arbitrary-collection sampler (`Random.rand!(gpuarray, 1:6)` scalar-indexes the + device array, which is disallowed). Overriding `_rand_impl!(::GPUPool, ::_SampleColl, …)` + raises a clear `ArgumentError` (covers the macro path and the direct path, live + disabled). + +Tests: `test/metal/test_random.jl` (run locally — **34 pass**) and +`test/cuda/test_random.jl` (mirror, to run on a CUDA machine). ## 8. Testing (`test/test_convenience.jl`, mirroring the `zeros!` testsets) @@ -166,7 +173,6 @@ arbitrary-collection sampler and would scalar-fallback. ## 10. Explicitly out of scope (YAGNI) -- GPU parity (separate PR — see §7). - `randexp!` (trivial later add). - `rng`-argument forms (`rand!(pool, rng, dims)`) — would complicate macro type extraction. - `randperm!` / `randcycle!` / `shuffle!` (different semantics — permutation/in-place on a user array). diff --git a/ext/AdaptiveArrayPoolsCUDAExt/convenience.jl b/ext/AdaptiveArrayPoolsCUDAExt/convenience.jl index d239266e..bee08196 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/convenience.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/convenience.jl @@ -76,3 +76,36 @@ AdaptiveArrayPools.default_eltype(::DisabledPool{:cuda}) = Float32 @inline AdaptiveArrayPools.acquire_view!(::DisabledPool{:cuda}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = CuArray{T, N}(undef, dims) @inline AdaptiveArrayPools.acquire_view!(::DisabledPool{:cuda}, x::CuArray) = CUDA.similar(x) @inline AdaptiveArrayPools.acquire_view!(::DisabledPool{:cuda}, x::AbstractArray) = CuArray{eltype(x)}(undef, size(x)) + +# ============================================================================== +# rand! / randn! (CUDA) +# ============================================================================== +# Live pool: the core `_rand_impl!`/`_randn_impl!` are inherited unchanged — they +# call `Random.rand!`/`randn!` on the acquired CuArray, which CUDA.jl runs on the +# GPU. Only the DisabledPool fallbacks and the collection-form rejection are needed. + +# --- rand!/randn! for DisabledPool{:cuda} (GPU-native CUDA RNG) --- +@inline AdaptiveArrayPools.rand!(::DisabledPool{:cuda}, ::Type{T}, dims::Vararg{Int, N}) where {T, N} = CUDA.rand(T, dims...) +@inline AdaptiveArrayPools.rand!(p::DisabledPool{:cuda}, dims::Vararg{Int, N}) where {N} = CUDA.rand(AdaptiveArrayPools.default_eltype(p), dims...) +@inline AdaptiveArrayPools.rand!(::DisabledPool{:cuda}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = CUDA.rand(T, dims...) +@inline AdaptiveArrayPools.rand!(p::DisabledPool{:cuda}, dims::NTuple{N, Int}) where {N} = CUDA.rand(AdaptiveArrayPools.default_eltype(p), dims...) + +@inline AdaptiveArrayPools.randn!(::DisabledPool{:cuda}, ::Type{T}, dims::Vararg{Int, N}) where {T, N} = CUDA.randn(T, dims...) +@inline AdaptiveArrayPools.randn!(p::DisabledPool{:cuda}, dims::Vararg{Int, N}) where {N} = CUDA.randn(AdaptiveArrayPools.default_eltype(p), dims...) +@inline AdaptiveArrayPools.randn!(::DisabledPool{:cuda}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = CUDA.randn(T, dims...) +@inline AdaptiveArrayPools.randn!(p::DisabledPool{:cuda}, dims::NTuple{N, Int}) where {N} = CUDA.randn(AdaptiveArrayPools.default_eltype(p), dims...) + +# --- Collection/range sampling is CPU-only --- +# GPU RNGs have no efficient arbitrary-collection sampler; `Random.rand!(gpuarray, S)` +# scalar-indexes the device array (disallowed). Reject with a clear error for both the +# live and disabled CUDA pools. Overriding `_rand_impl!` covers both the @with_pool +# macro path AND the direct public path (core `rand!(pool, S, …)` delegates to it). +@noinline _cuda_no_collection() = throw( + ArgumentError( + "rand!(pool, collection/range, dims) is CPU-only: CUDA GPU pools have no " * + "arbitrary-collection sampler. Use rand!(pool, T, dims) for uniform GPU randoms." + ) +) +const _CudaAnyPool = Union{CuAdaptiveArrayPool, DisabledPool{:cuda}} +@inline AdaptiveArrayPools._rand_impl!(::_CudaAnyPool, ::AdaptiveArrayPools._SampleColl, ::Int, ::Vararg{Int, N}) where {N} = _cuda_no_collection() +@inline AdaptiveArrayPools._rand_impl!(::_CudaAnyPool, ::AdaptiveArrayPools._SampleColl, ::NTuple{N, Int}) where {N} = _cuda_no_collection() diff --git a/ext/AdaptiveArrayPoolsMetalExt/convenience.jl b/ext/AdaptiveArrayPoolsMetalExt/convenience.jl index 852a8b22..24b53745 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/convenience.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/convenience.jl @@ -76,3 +76,36 @@ AdaptiveArrayPools.default_eltype(::DisabledPool{:metal}) = Float32 @inline AdaptiveArrayPools.acquire_view!(::DisabledPool{:metal}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = MtlArray{T, N}(undef, dims) @inline AdaptiveArrayPools.acquire_view!(::DisabledPool{:metal}, x::MtlArray) = Metal.similar(x) @inline AdaptiveArrayPools.acquire_view!(::DisabledPool{:metal}, x::AbstractArray) = MtlArray{eltype(x)}(undef, size(x)) + +# ============================================================================== +# rand! / randn! (Metal) +# ============================================================================== +# Live pool: the core `_rand_impl!`/`_randn_impl!` are inherited unchanged — they +# call `Random.rand!`/`randn!` on the acquired MtlArray, which Metal.jl runs on the +# GPU. Only the DisabledPool fallbacks and the collection-form rejection are needed. + +# --- rand!/randn! for DisabledPool{:metal} (GPU-native Metal RNG) --- +@inline AdaptiveArrayPools.rand!(::DisabledPool{:metal}, ::Type{T}, dims::Vararg{Int, N}) where {T, N} = Metal.rand(T, dims...) +@inline AdaptiveArrayPools.rand!(p::DisabledPool{:metal}, dims::Vararg{Int, N}) where {N} = Metal.rand(AdaptiveArrayPools.default_eltype(p), dims...) +@inline AdaptiveArrayPools.rand!(::DisabledPool{:metal}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = Metal.rand(T, dims...) +@inline AdaptiveArrayPools.rand!(p::DisabledPool{:metal}, dims::NTuple{N, Int}) where {N} = Metal.rand(AdaptiveArrayPools.default_eltype(p), dims...) + +@inline AdaptiveArrayPools.randn!(::DisabledPool{:metal}, ::Type{T}, dims::Vararg{Int, N}) where {T, N} = Metal.randn(T, dims...) +@inline AdaptiveArrayPools.randn!(p::DisabledPool{:metal}, dims::Vararg{Int, N}) where {N} = Metal.randn(AdaptiveArrayPools.default_eltype(p), dims...) +@inline AdaptiveArrayPools.randn!(::DisabledPool{:metal}, ::Type{T}, dims::NTuple{N, Int}) where {T, N} = Metal.randn(T, dims...) +@inline AdaptiveArrayPools.randn!(p::DisabledPool{:metal}, dims::NTuple{N, Int}) where {N} = Metal.randn(AdaptiveArrayPools.default_eltype(p), dims...) + +# --- Collection/range sampling is CPU-only --- +# GPU RNGs have no efficient arbitrary-collection sampler; `Random.rand!(gpuarray, S)` +# scalar-indexes the device array (disallowed). Reject with a clear error for both the +# live and disabled Metal pools. Overriding `_rand_impl!` covers both the @with_pool +# macro path AND the direct public path (core `rand!(pool, S, …)` delegates to it). +@noinline _metal_no_collection() = throw( + ArgumentError( + "rand!(pool, collection/range, dims) is CPU-only: Metal GPU pools have no " * + "arbitrary-collection sampler. Use rand!(pool, T, dims) for uniform GPU randoms." + ) +) +const _MetalAnyPool = Union{MetalAdaptiveArrayPool, DisabledPool{:metal}} +@inline AdaptiveArrayPools._rand_impl!(::_MetalAnyPool, ::AdaptiveArrayPools._SampleColl, ::Int, ::Vararg{Int, N}) where {N} = _metal_no_collection() +@inline AdaptiveArrayPools._rand_impl!(::_MetalAnyPool, ::AdaptiveArrayPools._SampleColl, ::NTuple{N, Int}) where {N} = _metal_no_collection() diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl index 6a813f27..c5974fbb 100644 --- a/test/cuda/runtests.jl +++ b/test/cuda/runtests.jl @@ -44,6 +44,7 @@ else include("test_nway_cache.jl") include("test_display.jl") include("test_convenience.jl") + include("test_random.jl") include("test_disabled_pool.jl") include("test_cuda_safety.jl") include("test_runtime_mutation.jl") diff --git a/test/cuda/test_random.jl b/test/cuda/test_random.jl new file mode 100644 index 00000000..6ddfd7bb --- /dev/null +++ b/test/cuda/test_random.jl @@ -0,0 +1,103 @@ +# CUDA rand! / randn! tests +# ========================= +# The live CUDA pool inherits the core _rand_impl!/_randn_impl! (which call +# Random.rand!/randn! on the acquired CuArray — CUDA.jl runs these on-device). +# Only the typed/default forms are supported on GPU; the collection/range form is +# CPU-only (no GPU arbitrary-collection sampler) and must raise a clear error. +# +# Mirrors test/metal/test_random.jl. Run on a CUDA-capable machine to verify parity. + +using Random + +@testset "CUDA rand!/randn!" begin + pool = get_task_local_cuda_pool() + checkpoint!(pool) + + @testset "rand! default type is Float32" begin + v = rand!(pool, 100) + @test v isa CuArray{Float32} + @test length(v) == 100 + va = Array(v) + @test all(x -> 0.0f0 <= x < 1.0f0, va) + + m = rand!(pool, 8, 8) + @test m isa CuArray{Float32, 2} + @test size(m) == (8, 8) + end + + @testset "rand! explicit type" begin + v32 = rand!(pool, Float32, 50) + @test v32 isa CuArray{Float32} + @test all(x -> 0.0f0 <= x < 1.0f0, Array(v32)) + + v16 = rand!(pool, Float16, 50) + @test v16 isa CuArray{Float16} + + vi = rand!(pool, Int32, 50) + @test vi isa CuArray{Int32} + end + + @testset "rand! tuple form" begin + m = rand!(pool, (4, 5)) + @test m isa CuArray{Float32, 2} + @test size(m) == (4, 5) + + m32 = rand!(pool, Float32, (3, 4)) + @test m32 isa CuArray{Float32, 2} + @test size(m32) == (3, 4) + end + + @testset "randn! default + explicit" begin + g = randn!(pool, 100) + @test g isa CuArray{Float32} + @test all(isfinite, Array(g)) + + g32 = randn!(pool, Float32, 4, 5) + @test g32 isa CuArray{Float32, 2} + @test size(g32) == (4, 5) + @test all(isfinite, Array(g32)) + + gt = randn!(pool, (6, 7)) + @test size(gt) == (6, 7) + end + + @testset "returns CuArray (not a view)" begin + v = rand!(pool, Float32, 10) + @test v isa CuArray + @test !(v isa SubArray) + end + + @testset "collection/range form is rejected (CPU-only)" begin + @test_throws ArgumentError rand!(pool, 1:6, 10) + @test_throws ArgumentError rand!(pool, 1:6, (5, 5)) + @test_throws ArgumentError rand!(pool, Int32(1):Int32(6), 8) + end + + rewind!(pool) + + @testset "DisabledPool{:cuda} fallbacks (GPU-native)" begin + DC = ext.DISABLED_CUDA + + v = rand!(DC, Float32, 16) + @test v isa CuArray{Float32} + @test all(x -> 0.0f0 <= x < 1.0f0, Array(v)) + + v2 = rand!(DC, 5, 5) # default Float32 + @test v2 isa CuArray{Float32, 2} + @test size(v2) == (5, 5) + + vt = rand!(DC, Float32, (3, 4)) # NTuple form + @test vt isa CuArray{Float32, 2} + @test size(vt) == (3, 4) + + g = randn!(DC, Float32, 16) + @test g isa CuArray{Float32} + @test all(isfinite, Array(g)) + + gt = randn!(DC, (2, 3)) # default Float32 NTuple + @test size(gt) == (2, 3) + + # collection rejected on the disabled pool too + @test_throws ArgumentError rand!(DC, 1:6, 10) + end +end diff --git a/test/metal/runtests.jl b/test/metal/runtests.jl index a327908f..006c523e 100644 --- a/test/metal/runtests.jl +++ b/test/metal/runtests.jl @@ -44,6 +44,7 @@ else include("test_allocation.jl") include("test_display.jl") include("test_convenience.jl") + include("test_random.jl") include("test_disabled_pool.jl") include("test_metal_safety.jl") include("test_runtime_mutation.jl") diff --git a/test/metal/test_random.jl b/test/metal/test_random.jl new file mode 100644 index 00000000..89dc8d1f --- /dev/null +++ b/test/metal/test_random.jl @@ -0,0 +1,101 @@ +# Metal rand! / randn! tests +# ========================== +# The live Metal pool inherits the core _rand_impl!/_randn_impl! (which call +# Random.rand!/randn! on the acquired MtlArray — Metal.jl runs these on-device). +# Only the typed/default forms are supported on GPU; the collection/range form is +# CPU-only (no GPU arbitrary-collection sampler) and must raise a clear error. + +using Random + +@testset "Metal rand!/randn!" begin + pool = get_task_local_metal_pool() + checkpoint!(pool) + + @testset "rand! default type is Float32" begin + v = rand!(pool, 100) + @test v isa MtlArray{Float32} + @test length(v) == 100 + va = Array(v) + @test all(x -> 0.0f0 <= x < 1.0f0, va) + + m = rand!(pool, 8, 8) + @test m isa MtlArray{Float32, 2} + @test size(m) == (8, 8) + end + + @testset "rand! explicit type" begin + v32 = rand!(pool, Float32, 50) + @test v32 isa MtlArray{Float32} + @test all(x -> 0.0f0 <= x < 1.0f0, Array(v32)) + + v16 = rand!(pool, Float16, 50) + @test v16 isa MtlArray{Float16} + + vi = rand!(pool, Int32, 50) + @test vi isa MtlArray{Int32} + end + + @testset "rand! tuple form" begin + m = rand!(pool, (4, 5)) + @test m isa MtlArray{Float32, 2} + @test size(m) == (4, 5) + + m32 = rand!(pool, Float32, (3, 4)) + @test m32 isa MtlArray{Float32, 2} + @test size(m32) == (3, 4) + end + + @testset "randn! default + explicit" begin + g = randn!(pool, 100) + @test g isa MtlArray{Float32} + @test all(isfinite, Array(g)) + + g32 = randn!(pool, Float32, 4, 5) + @test g32 isa MtlArray{Float32, 2} + @test size(g32) == (4, 5) + @test all(isfinite, Array(g32)) + + gt = randn!(pool, (6, 7)) + @test size(gt) == (6, 7) + end + + @testset "returns MtlArray (not a view)" begin + v = rand!(pool, Float32, 10) + @test v isa MtlArray + @test !(v isa SubArray) + end + + @testset "collection/range form is rejected (CPU-only)" begin + @test_throws ArgumentError rand!(pool, 1:6, 10) + @test_throws ArgumentError rand!(pool, 1:6, (5, 5)) + @test_throws ArgumentError rand!(pool, Int32(1):Int32(6), 8) + end + + rewind!(pool) + + @testset "DisabledPool{:metal} fallbacks (GPU-native)" begin + DM = ext.DISABLED_METAL + + v = rand!(DM, Float32, 16) + @test v isa MtlArray{Float32} + @test all(x -> 0.0f0 <= x < 1.0f0, Array(v)) + + v2 = rand!(DM, 5, 5) # default Float32 + @test v2 isa MtlArray{Float32, 2} + @test size(v2) == (5, 5) + + vt = rand!(DM, Float32, (3, 4)) # NTuple form + @test vt isa MtlArray{Float32, 2} + @test size(vt) == (3, 4) + + g = randn!(DM, Float32, 16) + @test g isa MtlArray{Float32} + @test all(isfinite, Array(g)) + + gt = randn!(DM, (2, 3)) # default Float32 NTuple + @test size(gt) == (2, 3) + + # collection rejected on the disabled pool too + @test_throws ArgumentError rand!(DM, 1:6, 10) + end +end From 4bbfd8335162b15271bcd09867dbf1df35c56ebb Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 26 Jun 2026 14:26:23 -0700 Subject: [PATCH 3/5] (fix): use real checkpoint! in GPU dynamic-type test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test/cuda/test_extension.jl simulated a @with_pool scope by setting pool._current_depth = 2 directly, but skipped the lockstep growth of the per-depth _touched_has_others vector that a real checkpoint! performs. get_typed_pool! then wrote _touched_has_others[2] into a 1-element vector — masked by @inbounds normally, but a BoundsError under Pkg.test's --check-bounds=yes (seen on NERSC/CI). Use the real checkpoint!(pool), exactly as the Metal test already did, and assert _touched_has_others == [false, true] in both the CUDA and Metal tests so the depth flag is verified rather than just executed. Test-only change. --- test/cuda/test_extension.jl | 3 ++- test/metal/test_extension.jl | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/test/cuda/test_extension.jl b/test/cuda/test_extension.jl index a8cfd2bd..247ff373 100644 --- a/test/cuda/test_extension.jl +++ b/test/cuda/test_extension.jl @@ -83,11 +83,12 @@ end @testset "Checkpoint auto-init for dynamic types" begin pool = CuAdaptiveArrayPool() - pool._current_depth = 2 # Simulate inside @with_pool scope + checkpoint!(pool) # Properly enter depth 2 (also grows _touched_has_others) tp = get_typed_pool!(pool, UInt16) @test tp._checkpoint_n_active == [0, 0] @test tp._checkpoint_depths == [0, 2] + @test pool._touched_has_others == [false, true] # depth-2 marked as "has dynamic types" end end diff --git a/test/metal/test_extension.jl b/test/metal/test_extension.jl index 633835f2..69abd6ef 100644 --- a/test/metal/test_extension.jl +++ b/test/metal/test_extension.jl @@ -98,6 +98,7 @@ end tp = get_typed_pool!(pool, UInt16) @test tp._checkpoint_n_active == [0, 0] @test tp._checkpoint_depths == [0, 2] + @test pool._touched_has_others == [false, true] # depth-2 marked as "has dynamic types" end end From f3f8273d10697f948cc8d11c06be41c33d108602 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 26 Jun 2026 14:26:33 -0700 Subject: [PATCH 4/5] (test): cover remaining rand!/randn! NTuple + qualified-name forms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codecov flagged 13 uncovered patch lines — all never-called method bodies: - convenience.jl: randn! typed-NTuple public, the two _randn_impl! NTuple delegations (reached via the macro size(x) transform), and the rand!/randn! DisabledPool NTuple fallbacks - macros.jl: the qualified-name AdaptiveArrayPools.rand!/randn! → _impl! transform branches Add two testsets exercising every NTuple form (direct + size(x) macro path + DisabledPool) and a qualified-name call inside @with_pool. All 13 lines now hit. Also refresh the design doc: status -> Implemented, GPU-parity row -> done (addresses the stale-doc points from the Copilot PR review). --- docs/design/rand_randn_pool_constructors.md | 4 +- test/test_random.jl | 45 +++++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/docs/design/rand_randn_pool_constructors.md b/docs/design/rand_randn_pool_constructors.md index ba1f78eb..52ecd8e7 100644 --- a/docs/design/rand_randn_pool_constructors.md +++ b/docs/design/rand_randn_pool_constructors.md @@ -1,6 +1,6 @@ # Design: `rand!` / `randn!` pool constructors -**Status:** Approved design, pending implementation +**Status:** Implemented (CPU + Metal + CUDA) **Branch:** `feature/rand-randn-pool-constructors` **Author:** Min-Gu Yoo @@ -30,7 +30,7 @@ values, preserving the package's zero-allocation guarantee inside `@with_pool`. | Namespace | **`import Random: rand!, randn!` then `export rand!, randn!`** (re-export) | Makes `rand!` a true peer of `zeros!` — works with just `using AdaptiveArrayPools`. Re-exporting the *same binding* means **no conflict warning** if the user also does `using Random`. | | Distributions (v1) | **`rand!` (uniform) + `randn!` (normal)** | Covers the overwhelming majority of use. `randexp!` deferred (low demand, trivial to add later via same machinery). | | Collection/range form | **Include `rand!(pool, S, dims)`** (CPU-only) | This is the real "random integers in a range" case (`rand!(pool, 1:6, n)`). More useful than full-range `rand!(pool, Int, n)`. | -| GPU parity | **Deferred to a follow-up PR** | Live GPU pools inherit the core impl via dispatch (see §7), so it is cheap, but Metal's `randn!`/`rand!` MPS coverage needs verification. Kept out of v1 to ship CPU cleanly. | +| GPU parity | **Implemented (Metal + CUDA)** | Live GPU pools inherit the core impl via dispatch (see §7); each extension adds only `DisabledPool` fallbacks + a CPU-only-collection guard. Verified on Metal. | ## 3. Why this is "wide but shallow" diff --git a/test/test_random.jl b/test/test_random.jl index 85dc8c12..a0f65c9a 100644 --- a/test/test_random.jl +++ b/test/test_random.jl @@ -618,4 +618,49 @@ using AdaptiveArrayPools: _extract_acquired_vars end @test ok end + + # ========================================================================= + # 19. Coverage of the remaining NTuple method forms (typed-NTuple randn!, + # the _randn_impl! NTuple delegations via the macro size(x) path, and + # the DisabledPool NTuple fallbacks). + # ========================================================================= + @testset "NTuple method-form coverage" begin + pool = AdaptiveArrayPool() + + # randn! public typed-NTuple form (direct) + g = randn!(pool, Float32, (3, 4)) + @test g isa Matrix{Float32} + @test size(g) == (3, 4) + + # _randn_impl! NTuple delegations are reached via the macro's size(x) transform + x1 = rand(7) + x2 = rand(3, 5) + @with_pool pool begin + a = randn!(pool, Float64, size(x2)) # _randn_impl!(pool, T, NTuple) + @test size(a) == (3, 5) + b = randn!(pool, size(x1)) # _randn_impl!(pool, NTuple) default + @test size(b) == (7,) + nothing + end + + # DisabledPool NTuple fallbacks + @test rand!(DISABLED_CPU, (3, 4)) isa Matrix{Float64} # rand! default NTuple + @test randn!(DISABLED_CPU, Float32, (3, 4)) isa Matrix{Float32} # randn! typed NTuple + @test randn!(DISABLED_CPU, (3, 4)) isa Matrix{Float64} # randn! default NTuple + end + + # ========================================================================= + # 20. Qualified-name calls inside @with_pool exercise the macro's + # `Module.rand!`/`Module.randn!` → `_impl!` transform branches. + # ========================================================================= + @testset "qualified-name rand!/randn! in @with_pool" begin + result = @with_pool pool begin + v = AdaptiveArrayPools.rand!(pool, Float64, 10) + g = AdaptiveArrayPools.randn!(pool, 10) + @test v isa Vector{Float64} + @test g isa Vector{Float64} + length(v) + length(g) + end + @test result == 20 + end end From e13fcef3b6096f9137a95e013e5e2635fac83dc4 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Fri, 26 Jun 2026 14:30:11 -0700 Subject: [PATCH 5/5] (chore): untrack internal design doc docs/design/rand_randn_pool_constructors.md is an internal/local working note, not part of the shipped package. Remove it from version control (kept locally, untracked) so it no longer appears in the PR. --- docs/design/rand_randn_pool_constructors.md | 180 -------------------- 1 file changed, 180 deletions(-) delete mode 100644 docs/design/rand_randn_pool_constructors.md diff --git a/docs/design/rand_randn_pool_constructors.md b/docs/design/rand_randn_pool_constructors.md deleted file mode 100644 index 52ecd8e7..00000000 --- a/docs/design/rand_randn_pool_constructors.md +++ /dev/null @@ -1,180 +0,0 @@ -# Design: `rand!` / `randn!` pool constructors - -**Status:** Implemented (CPU + Metal + CUDA) -**Branch:** `feature/rand-randn-pool-constructors` -**Author:** Min-Gu Yoo - -## 1. Goal - -Add random-array convenience constructors to the pool API, peers of `zeros!` / -`ones!`: - -```julia -@with_pool pool begin - v = rand!(pool, 100) # Vector{Float64}, U[0,1) - m = rand!(pool, Float32, 8, 8) # Matrix{Float32}, U[0,1) - d = rand!(pool, 1:6, 10) # Vector{Int}, each ∈ 1:6 (dice) - g = randn!(pool, 100) # Vector{Float64}, N(0,1) -end -``` - -These acquire a right-sized array from the pool and fill it in place with random -values, preserving the package's zero-allocation guarantee inside `@with_pool`. - -## 2. Key decisions (settled during brainstorming) - -| Decision | Choice | Rationale | -|---|---|---| -| Dependency | **`Random` as a hard `[deps]`** | `Random` is a stdlib (in the sysimage); `using Random` adds ~0 load cost. The usual reason to use a weakdep/extension (avoid forcing a heavy dep — cf. CUDA/Metal) does not apply. | -| Why not an extension | The convenience functions are **wired into `@with_pool` at 3 sites**; an extension would force a core stub + `GlobalRef` to a conditionally-defined impl and yield cryptic `MethodError`s when `Random` isn't loaded. Hard-dep keeps `_rand_impl!` next to `_zeros_impl!`. | -| Namespace | **`import Random: rand!, randn!` then `export rand!, randn!`** (re-export) | Makes `rand!` a true peer of `zeros!` — works with just `using AdaptiveArrayPools`. Re-exporting the *same binding* means **no conflict warning** if the user also does `using Random`. | -| Distributions (v1) | **`rand!` (uniform) + `randn!` (normal)** | Covers the overwhelming majority of use. `randexp!` deferred (low demand, trivial to add later via same machinery). | -| Collection/range form | **Include `rand!(pool, S, dims)`** (CPU-only) | This is the real "random integers in a range" case (`rand!(pool, 1:6, n)`). More useful than full-range `rand!(pool, Int, n)`. | -| GPU parity | **Implemented (Metal + CUDA)** | Live GPU pools inherit the core impl via dispatch (see §7); each extension adds only `DisabledPool` fallbacks + a CPU-only-collection guard. Verified on Metal. | - -## 3. Why this is "wide but shallow" - -The implementation body is ~2 lines per function (`acquire!` → `Random.rand!`). -The work is **parallelism with `zeros!`**: every site that already knows about -`zeros!` must learn `rand!` / `randn!`. This is a consequence of the package -centralizing checkpoint / escape-analysis / type-registration logic in the -macro — new convenience functions are cheap to write but must be *registered* in -the macro. The spec below is therefore primarily a **checklist of mirror-sites**. - -## 4. Public API surface (CPU, v1) - -Mirrors `zeros!`'s four forms, plus the collection form for `rand!`: - -```julia -# Uniform -rand!(pool, dims...) # default_eltype(pool) -rand!(pool, T::Type, dims...) # typed -rand!(pool, dims::NTuple{N,Int}) # tuple form (for macro size(x) handling) -rand!(pool, T::Type, dims::NTuple{N,Int}) -rand!(pool, S, dims...) # sample from collection S; eltype = eltype(S) -rand!(pool, S, dims::NTuple{N,Int}) - -# Normal (T must be AbstractFloat; non-float T inherits Random.randn!'s error) -randn!(pool, dims...) -randn!(pool, T::Type, dims...) -randn!(pool, dims::NTuple{N,Int}) -randn!(pool, T::Type, dims::NTuple{N,Int}) -``` - -**Default eltype** comes from `default_eltype(pool)` (CPU: `Float64`), so -`rand!(pool, 10)` → `U[0,1)` Float64 and `randn!(pool, 10)` → `N(0,1)` Float64. - -**Collection form semantics**: `rand!(pool, S, dims...)` acquires an array of -`eltype(S)` and fills it via `Random.rand!(arr, S)`. Examples: -`rand!(pool, 1:6, 10)` (Int dice), `rand!(pool, ('a','b','c'), 5)` (Char). - -## 5. Components & files to change - -### 5.1 `Project.toml` -- Add `Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"` to `[deps]`. -- Add `Random = "1"` to `[compat]`. - -### 5.2 `src/AdaptiveArrayPools.jl` (main module) -- `import Random: rand!, randn!` (brings names in scope for extending + re-export). -- `import Random` (so `rand`, `randn` allocating fallbacks are reachable for `DisabledPool`). -- Add `rand!, randn!` to the convenience-functions `export` line (next to `zeros!, ones!, ...`). - -### 5.3 `src/convenience.jl` (next to `zeros!` / `_zeros_impl!`) -Public wrappers (each does `_record_type_touch!` + `_set_pending_callsite!` + impl, -identical structure to `zeros!`): -- `rand!(pool::AbstractArrayPool, ...)` — all six forms above. -- `randn!(pool::AbstractArrayPool, ...)` — four forms. - -Internal impls (backend-agnostic — rely on `Random.rand!`/`randn!` dispatch on the array type): -- `_rand_impl!(pool, T, dims...) = (arr = _acquire_impl!(pool, T, dims...); Random.rand!(arr); arr)` -- `_rand_impl!(pool, dims...) = _rand_impl!(pool, default_eltype(pool), dims...)` -- `_rand_impl!(pool, S, dims...) = (arr = _acquire_impl!(pool, eltype(S), dims...); Random.rand!(arr, S); arr)` *(collection)* -- NTuple overloads for each. -- `_randn_impl!(...)` analogous to `_rand_impl!` minus the collection form, using `Random.randn!`. - -**Dispatch disambiguation (collection vs typed vs default):** the collection -argument `S` is constrained to sampleable, non-`Type`, non-`Integer` collections — -`S::Union{AbstractArray, Tuple, AbstractRange, AbstractSet, AbstractString, AbstractDict}` -— so `rand!(pool, 1:6, 10)` picks the collection method while `rand!(pool, 10)` -and `rand!(pool, Int, 10)` pick the dims/typed methods. This mirrors how Base -separates `rand(coll, n)` from `rand(T, n)`. - -### 5.4 `src/convenience.jl` — `DisabledPool{:cpu}` fallbacks -Mirror the `zeros!` disabled-pool methods (`convenience.jl:439-442`): -- `rand!(::DisabledPool{:cpu}, T, dims...) = rand(T, dims...)` -- `rand!(p::DisabledPool{:cpu}, dims...) = rand(default_eltype(p), dims...)` -- `rand!(::DisabledPool{:cpu}, S, dims...) = rand(S, dims...)` *(collection)* -- `randn!(::DisabledPool{:cpu}, T, dims...) = randn(T, dims...)`, default + NTuple forms. - -### 5.5 `src/macros.jl` — the three mirror-sites -1. **Type extraction** `_extract_acquire_types` (~L1143-1189): add `:rand!` / `:randn!`. - - Typed form: read the type token like `zeros!`/`ones!`. - - Default form: use `default_eltype(pool)` like `zeros!`. - - **Collection form (NEW code path):** emit `eltype(S)` as a runtime expression - (constant-folds for literal ranges) rather than reading a literal token. -2. **Call → impl transform** (~L1483-1537, GlobalRefs at L1491-1492, impl-name list at L1584): - - Add `_RAND_IMPL_REF = GlobalRef(@__MODULE__, :_rand_impl!)` and `_RANDN_IMPL_REF`. - - Map `:rand! → _RAND_IMPL_REF`, `:randn! → _RANDN_IMPL_REF` (both the `fn ==` and `QuoteNode` branches). - - Add `:_rand_impl!`, `:_randn_impl!` to the impl-name list (~L1584). - - The transform only swaps the function name; the collection arg passes through unchanged. -3. **Escape-detection symbol lists** (~L1943-1973, and the lists feeding L2032/L2371): - - Add `:rand!`, `:randn!` so `x = rand!(pool, 10); return x` triggers `PoolEscapeError`, - matching `zeros!` behavior. - -## 6. Error handling - -- `randn!(pool, Int, n)` — inherits `Random.randn!`'s error for non-float eltypes - (no special handling; documented). -- Collection form with an eltype the pool can't store — surfaces through the - normal `acquire!` type path (`_record_type_touch!(pool, eltype(S))`). -- Escape of a pool-backed random array — `PoolEscapeError`, via the macro escape - lists (§5.5.3), identical to `zeros!`. - -## 7. GPU (Metal + CUDA) — implemented - -`CuAdaptiveArrayPool` / `MetalAdaptiveArrayPool <: AbstractArrayPool`, and the GPU -extensions do **not** override `_rand_impl!`/`_randn_impl!` for live pools — they -**inherit the core impl**, which works on GPU because `Random.rand!`/`randn!` are -overloaded for `CuArray`/`MtlArray` and run on-device. **Verified on Metal** -(Apple Silicon): `Random.rand!`/`randn!` on `MtlArray` are GPU-native for -`Float32`/`Float16`/`Int32`. So the live GPU pool gets the typed/default forms for -free, exactly like `zeros!`. - -GPU-specific code (per backend, in the extension): -- `DisabledPool{:cuda}`/`{:metal}` fallbacks (`Metal.rand`/`Metal.randn`, - `CUDA.rand`/`CUDA.randn`). -- **Collection-form rejection**: `rand!(pool, S, dims)` is **CPU-only** — GPU RNGs have - no arbitrary-collection sampler (`Random.rand!(gpuarray, 1:6)` scalar-indexes the - device array, which is disallowed). Overriding `_rand_impl!(::GPUPool, ::_SampleColl, …)` - raises a clear `ArgumentError` (covers the macro path and the direct path, live + disabled). - -Tests: `test/metal/test_random.jl` (run locally — **34 pass**) and -`test/cuda/test_random.jl` (mirror, to run on a CUDA machine). - -## 8. Testing (`test/test_convenience.jl`, mirroring the `zeros!` testsets) - -- `rand!`: length/size & eltype correct for default, typed (`Float32`, `Int`), and - multi-dim forms; **all values ∈ [0,1)** for float eltypes. -- `rand!(pool, 1:6, n)`: `eltype == Int`, **all values ∈ 1:6**; a Char/tuple collection case. -- `randn!`: length/size & eltype correct; values finite (statistical sanity, not exact equality). -- **Zero-allocation** inside `@with_pool` (reuse the existing alloc-test helper / `test_zero_allocation.jl` pattern). -- **Escape detection**: `x = rand!(pool, 10); return x` triggers `PoolEscapeError` - (mirror the existing `test_compile_escape.jl` style). -- **`DisabledPool` fallback**: returns a plain `Array` of correct type/size. -- **Determinism**: seed via `Random.seed!` (no `rng` arg in v1) and assert reproducible values. -- **Aqua** (`test_aqua.jl`): confirm no new ambiguities, no piracy flag (methods are - on our `AbstractArrayPool` type → not piracy), deps/compat consistent with the new `Random` entry. - -## 9. Documentation - -- Docstrings on `rand!` / `randn!` mirroring the `zeros!` docstring (signatures, - default-eltype note, collection-form example, `See also`). -- README / docs convenience-function section: add `rand!` / `randn!` next to `zeros!` / `ones!`. - -## 10. Explicitly out of scope (YAGNI) - -- `randexp!` (trivial later add). -- `rng`-argument forms (`rand!(pool, rng, dims)`) — would complicate macro type extraction. -- `randperm!` / `randcycle!` / `shuffle!` (different semantics — permutation/in-place on a user array). -- `bitrand` / `rand!(pool, Bit, …)` — tied to the deferred `Bit`/`BitArray` handling. -- GPU collection sampling.