diff --git a/ext/DistributionsExt.jl b/ext/DistributionsExt.jl index 8402c9701..e13eb6584 100644 --- a/ext/DistributionsExt.jl +++ b/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 diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 69ed8c9f5..9274496b4 100644 --- a/ext/MPIExt.jl +++ b/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 \ No newline at end of file +end # module MPIExt diff --git a/src/Dagger.jl b/src/Dagger.jl index 1a4099bee..81881716e 100644 --- a/src/Dagger.jl +++ b/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") diff --git a/src/array/stencil.jl b/src/array/stencil.jl index 858445d77..0f3af7f22 100644 --- a/src/array/stencil.jl +++ b/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 diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 725e1094e..dc6668486 100644 --- a/src/datadeps/aliasing.jl +++ b/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 diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index c78b96d83..617f7ea8b 100644 --- a/src/memory-spaces.jl +++ b/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 diff --git a/src/utils/libc-array.jl b/src/utils/libc-array.jl new file mode 100644 index 000000000..928b3c2b4 --- /dev/null +++ b/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