(feat): rand! / randn! — random-array pool constructors (CPU + Metal + CUDA)

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"

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
 

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

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

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

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, "<direct rand! call>")
+    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, "<direct rand! call>")
+    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, "<direct rand! call>")
+    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, "<direct rand! call>")
+    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, "<direct rand! call>")
+    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, "<direct rand! call>")
+    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, "<direct randn! call>")
+    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, "<direct randn! call>")
+    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, "<direct randn! call>")
+    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, "<direct randn! call>")
+    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)