datadeps: Make some Array allocations eagerly freeable

ext/DistributionsExt.jl

@@ -7,9 +7,39 @@ else
 end
 
 using Dagger, Random
+import Dagger: chunks, imap!, InOut, randfork
 
+# Default-RNG path: each task uses Julia's task-local default RNG, so concurrent
+# chunk tasks do not share mutable RNG state.
 Random.rand!(s::Sampleable, A::DArray{T}) where T = map!(_ -> rand(s), A)
-Random.rand!(rng::AbstractRNG, s::Sampleable{Univariate}, A::DArray{T}) where T = map!(_ -> rand(rng, s), A)
-Random.rand!(rng::AbstractRNG, s::Sampleable{ArrayLikeVariate{M}}, A::DArray{T}) where {M, T} = map!(_ -> rand(rng, s), A)
+
+# Explicit-RNG paths must fork a distinct RNG per chunk (and bind it in `let`),
+# matching `Random.rand!(::AbstractRNG, ::DArray)`. Otherwise `map!` would capture
+# one shared RNG in a single closure and race under multithreading — corrupting
+# RNG state and yielding out-of-range / NaN values.
+function Random.rand!(rng::AbstractRNG, s::Sampleable{Univariate}, A::DArray{T}) where T
+    part_sz = prod(map(length, first(A.subdomains).indexes))
+    Dagger.spawn_datadeps() do
+        for Ac in chunks(A)
+            rng = randfork(rng, part_sz)
+            let rng = rng, s = s
+                Dagger.@spawn imap!(InOut(_ -> rand(rng, s)), InOut(Ac))
+            end
+        end
+    end
+    return A
+end
+function Random.rand!(rng::AbstractRNG, s::Sampleable{ArrayLikeVariate{M}}, A::DArray{T}) where {M,T}
+    part_sz = prod(map(length, first(A.subdomains).indexes))
+    Dagger.spawn_datadeps() do
+        for Ac in chunks(A)
+            rng = randfork(rng, part_sz)
+            let rng = rng, s = s
+                Dagger.@spawn imap!(InOut(_ -> rand(rng, s)), InOut(Ac))
+            end
+        end
+    end
+    return A
+end
 
 end # module DistributionsExt

ext/MPIExt.jl

@@ -1669,7 +1669,7 @@ function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_pro
     if child_types === nothing
         # Leaf: transfer to the destination, sharing via the aliased-object cache
         return aliased_object!(cache, w) do w
-            return mpi_endpoint_transfer(accel, from_proc, to_proc, from_space, to_space, w)
+            return Dagger.libc_backed(mpi_endpoint_transfer(accel, from_proc, to_proc, from_space, to_space, w))
         end
     end
 
@@ -2131,4 +2131,4 @@ accel_matches_proc(accel::MPIAcceleration, proc::MPIClusterProc) = true
 accel_matches_proc(accel::MPIAcceleration, proc::MPIProcessor) = true
 accel_matches_proc(accel::MPIAcceleration, proc) = false
 
-end # module MPIExt
+end # module MPIExt

src/Dagger.jl

@@ -90,6 +90,7 @@ include("submission.jl")
 abstract type MemorySpace end
 include("utils/memory-span.jl")
 include("utils/interval_tree.jl")
+include("utils/libc-array.jl")
 include("memory-spaces.jl")
 include("acceleration.jl")
 

src/array/stencil.jl

@@ -1239,15 +1239,23 @@ macro stencil(orig_ex)
             end
         end))
 
-        # 2d. After spawn_datadeps completes all fill and stencil tasks, fetch the filled
-        # HaloArray Chunks and store them in the cache (replacing any alloc DTasks from 2a).
-        # spawn_datadeps already waited for everything, so fetch is instant.
+        # 2d. After spawn_datadeps completes, resolve any alloc DTasks in the cache to
+        # their origin Chunks. Do *not* store `fetch(fill_task)`: Datadeps may have run
+        # fill on a libc-backed remote slot copy and then `unsafe_free!`'d that copy at
+        # region end, so the fill task's return value can dangle. The origin Chunk (from
+        # alloc_halo / a prior iteration) receives the updated data via remainder sync
+        # and is what must stay in the cache for the next iteration.
         for read_var in read_vars
             if read_var in keys(neighborhoods)
                 syms = cache_sym_map[read_var]
+                @gensym cache_entry
                 push!(final_ex.args, quote
                     for $chunk_idx in $CartesianIndices($chunks($read_var))
-                        $(syms.inner_cache_var)[($chunk_idx, $(syms.halo_width_var))] = fetch($(syms.fill_tasks)[$chunk_idx]; raw=true)
+                        $(syms.cache_key_var) = ($chunk_idx, $(syms.halo_width_var))
+                        $cache_entry = $(syms.inner_cache_var)[$(syms.cache_key_var)]
+                        if $cache_entry isa $DTask
+                            $(syms.inner_cache_var)[$(syms.cache_key_var)] = fetch($cache_entry; raw=true)
+                        end
                     end
                 end)
             end

src/datadeps/aliasing.jl

@@ -1003,7 +1003,7 @@ function move_rewrap(accel, cache::AliasedObjectCache, from_proc::Processor, to_
         # Leaf: transfer the value, sharing via the aliased-object cache
         return aliased_object!(cache, data) do data
             return remotecall_endpoint_transfer(accel, from_proc, to_proc, from_space, to_space, data) do accel, from_proc, to_proc, from_space, to_space, data
-                return tochunk(move(from_proc, to_proc, data), to_proc, to_space)
+                return tochunk(libc_backed(move(from_proc, to_proc, data)), to_proc, to_space)
             end
         end
     end

src/memory-spaces.jl

@@ -651,4 +651,10 @@ unsafe_free!(x::Chunk) = remotecall_fetch(root_worker_id(x), x) do x
     return
 end
 unsafe_free!(x::DTask) = unsafe_free!(fetch(x; move_value=false, unwrap=false))
+# CPU `Array`s can only be freed if we allocated their backing memory via
+# `Libc.malloc` (see `alloc_libc_array`); freeing arbitrary Julia-allocated
+# `Array`s is not possible (and would be unsafe). `_libc_finalize!` only frees
+# allocations present in our objectid registry, so this is a no-op for plain
+# `Array`s.
+unsafe_free!(x::Array) = _libc_finalize!(x)
 unsafe_free!(x) = nothing # Do nothing by default

src/utils/libc-array.jl

@@ -0,0 +1,150 @@
+# Libc-backed `Array` allocations
+#
+# Memory allocated by Julia's built-in `Array` allocator cannot be explicitly
+# freed; we must wait for the GC to reclaim it. This makes `unsafe_free!` a
+# no-op for CPU `Array`s, in contrast to GPU arrays (which support eager,
+# explicit freeing). To give Datadeps the ability to eagerly free large CPU
+# buffers, we allocate the backing memory with `Libc.malloc` and wrap it in an
+# `Array` via `unsafe_wrap`. Every allocation is recorded in a per-process
+# registry keyed by `objectid`, so `unsafe_free!` only ever calls `Libc.free` on
+# memory we are certain we own (never on a user's plain `Array`).
+#
+# The registry must *not* be keyed by the malloc pointer alone: after an eager
+# `unsafe_free!`, the Julia `Array` object still holds a dangling data pointer
+# and its finalizer may run later. If malloc reuses that address for a newer
+# Libc-backed array in the meantime, a pointer-keyed finalizer would steal the
+# new registration and free live memory (heap corruption / use-after-free).
+# Keying by `objectid` keeps the eager-free and finalizer paths tied to the
+# specific `Array` instance.
+
+const LIBC_ARRAY_LOCK = Threads.SpinLock()
+# objectid(A) => (ptr, nbytes)
+const LIBC_ARRAY_ALLOCS = Dict{UInt,Tuple{Ptr{Cvoid},Int}}()
+
+# Bookkeeping for benchmarking/diagnostics (bytes).
+const LIBC_ARRAY_TOTAL_BYTES = Ref{Int}(0)
+const LIBC_ARRAY_LIVE_BYTES = Ref{Int}(0)
+const LIBC_ARRAY_PEAK_BYTES = Ref{Int}(0)
+const LIBC_ARRAY_NUM_ALLOCS = Ref{Int}(0)
+
+function _libc_register!(A::Array, ptr::Ptr{Cvoid}, nbytes::Integer)
+    id = objectid(A)
+    @lock LIBC_ARRAY_LOCK begin
+        @assert !haskey(LIBC_ARRAY_ALLOCS, id) "Array $(id) is already registered as Libc-backed"
+        LIBC_ARRAY_ALLOCS[id] = (ptr, Int(nbytes))
+        LIBC_ARRAY_TOTAL_BYTES[] += nbytes
+        LIBC_ARRAY_NUM_ALLOCS[] += 1
+        LIBC_ARRAY_LIVE_BYTES[] += nbytes
+        LIBC_ARRAY_PEAK_BYTES[] = max(LIBC_ARRAY_PEAK_BYTES[], LIBC_ARRAY_LIVE_BYTES[])
+    end
+    return
+end
+
+# Returns the registered `(ptr, nbytes)` after removing it, or `nothing`.
+function _libc_unregister!(A::Array)
+    id = objectid(A)
+    @lock LIBC_ARRAY_LOCK begin
+        if haskey(LIBC_ARRAY_ALLOCS, id)
+            ptr_nbytes = LIBC_ARRAY_ALLOCS[id]
+            delete!(LIBC_ARRAY_ALLOCS, id)
+            LIBC_ARRAY_LIVE_BYTES[] -= ptr_nbytes[2]
+            return ptr_nbytes
+        end
+        return nothing
+    end
+end
+
+"""
+    libc_array_stats() -> NamedTuple
+
+Return diagnostics about Libc-backed `Array` allocations made by Dagger on the
+current process: `total_bytes` (cumulative bytes ever allocated), `live_bytes`
+(currently-live bytes), `peak_bytes` (high-water mark of live bytes), and
+`num_allocs` (number of allocations).
+"""
+libc_array_stats() = @lock LIBC_ARRAY_LOCK (;
+    total_bytes = LIBC_ARRAY_TOTAL_BYTES[],
+    live_bytes = LIBC_ARRAY_LIVE_BYTES[],
+    peak_bytes = LIBC_ARRAY_PEAK_BYTES[],
+    num_allocs = LIBC_ARRAY_NUM_ALLOCS[],
+)
+
+"""
+    reset_libc_array_stats!()
+
+Reset the cumulative/peak counters reported by [`libc_array_stats`](@ref).
+`live_bytes` is left untouched (it reflects genuinely-live allocations).
+"""
+function reset_libc_array_stats!()
+    @lock LIBC_ARRAY_LOCK begin
+        LIBC_ARRAY_TOTAL_BYTES[] = 0
+        LIBC_ARRAY_NUM_ALLOCS[] = 0
+        LIBC_ARRAY_PEAK_BYTES[] = LIBC_ARRAY_LIVE_BYTES[]
+    end
+    return
+end
+
+"""
+    is_libc_allocated(A) -> Bool
+
+Return `true` if `A` is a CPU `Array` whose backing memory was allocated by
+Dagger via `Libc.malloc` (and can therefore be eagerly freed by
+[`unsafe_free!`](@ref)).
+"""
+function is_libc_allocated(A::Array)
+    id = objectid(A)
+    @lock LIBC_ARRAY_LOCK return haskey(LIBC_ARRAY_ALLOCS, id)
+end
+is_libc_allocated(@nospecialize(x)) = false
+
+"""
+    alloc_libc_array(T, dims...) -> Array{T}
+
+Allocate an `Array{T}` whose backing memory comes from `Libc.malloc`, allowing
+it to be eagerly freed with [`unsafe_free!`](@ref). `T` must be an `isbits` type,
+since `Libc`-managed memory cannot safely hold GC-tracked references. A
+finalizer is attached as a safety net, so the memory is still reclaimed if
+`unsafe_free!` is never called.
+"""
+alloc_libc_array(::Type{T}, dims::Integer...) where {T} =
+    alloc_libc_array(T, convert(Dims, dims))
+function alloc_libc_array(::Type{T}, dims::Dims{N}) where {T,N}
+    isbitstype(T) || throw(ArgumentError("alloc_libc_array requires an isbits element type, got $T"))
+    nbytes = prod(dims) * sizeof(T)
+    if nbytes == 0
+        # `Libc.malloc(0)` is implementation-defined; fall back to a normal
+        # (empty) Array, which has nothing to free anyway.
+        return Array{T,N}(undef, dims)
+    end
+    ptr = Libc.malloc(nbytes)
+    ptr == C_NULL && throw(OutOfMemoryError())
+    A = unsafe_wrap(Array{T,N}, Ptr{T}(ptr), dims; own=false)
+    _libc_register!(A, ptr, nbytes)
+    finalizer(_libc_finalize!, A)
+    return A
+end
+
+function _libc_finalize!(A::Array)
+    ptr_nbytes = _libc_unregister!(A)
+    if ptr_nbytes !== nothing
+        Libc.free(ptr_nbytes[1])
+    end
+    return
+end
+
+"""
+    libc_backed(x)
+
+If `x` is a CPU `Array` with an `isbits` element type that is not already
+Libc-backed, return a Libc-backed copy of it (which can later be eagerly freed
+via [`unsafe_free!`](@ref)); otherwise return `x` unchanged. Used by Datadeps to
+make its internal buffer copies freeable.
+"""
+function libc_backed(A::Array{T,N}) where {T,N}
+    isbitstype(T) || return A
+    is_libc_allocated(A) && return A
+    B = alloc_libc_array(T, size(A))
+    copyto!(B, A)
+    return B
+end
+libc_backed(@nospecialize(x)) = x