diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index f8de25146..997eaa943 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -658,6 +658,79 @@ function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(cache.a end end +""" + ArenaRef{T,N} + +The handle of a Datadeps slot carved out of a `SlotArena`: an `Array{T,N}` of +size `dims` at address `addr` on worker `owner`, inside arena `arena`. Unlike a +`DRef`, the planner makes one without asking `owner`, since it already knows the +arena's base address: so it knows the slot's aliasing before the slot exists. +""" +struct ArenaRef{T,N} + owner::Int + arena::UInt + addr::UInt + dims::NTuple{N,Int} +end + +""" + SlotArena + +One block of memory a region reserved on a worker for its slots, which the +planner hands out by bumping `used`. It is freed whole, once, when the region +ends (see `RegionSlotArenas`). +""" +mutable struct SlotArena + const space::CPURAMMemorySpace + const id::UInt + const base::UInt + const capacity::Int + used::Int +end + +""" + RegionSlotArenas + +Every slot arena a `spawn_datadeps` region reserved, as `(owner, id)`, which it +frees once all of its tasks have finished, failed or not: right away when the +region succeeds, or in the background when a task fails +(`free_region_arenas_when_done!`). Freeing an arena from a task that depends on +its slots' users would leak it whenever one of them fails, since the scheduler +fails a task's dependents outright; the per-slot buffers this replaces were +reclaimed by their finalizers in that case, but an arena has none. +""" +mutable struct RegionSlotArenas + const lock::ReentrantLock + const arenas::Vector{Tuple{Int,UInt}} + # The region's `wait_all` queue's tasks, which every task it launches joins + tasks::Union{Vector{DTask},Nothing} +end +RegionSlotArenas() = RegionSlotArenas(ReentrantLock(), Tuple{Int,UInt}[], nothing) +const DATADEPS_REGION_ARENAS = ScopedValue{Union{RegionSlotArenas,Nothing}}(nothing) + +""" + SlotPool + +The memory a region allocates its slots from, ahead of need. `generate_slot!`'s +direct path (see `direct_uninit_slot`) takes each dense CPU slot from an arena +(`SlotArena`) on its destination, computing its address and aliasing itself, +instead of paying a round trip per slot -- most of what was left of a flat +launch loop, where one allocation costs ~0.3 ms on an idle worker but ~5 ms in +the middle of a region, queued behind the region's copies. + +For AOT plans `prefill_slot_pool!` reserves exactly what the plan needs, with +one call per destination worker. Anything else reserves arenas on first need, +growing from `SLOT_ARENA_MIN_BYTES` to `SLOT_ARENA_BYTES`. Each arena is freed +whole when the region ends (`RegionSlotArenas`), so a region no longer spawns a +free task per slot. The pool is only enabled inside a `spawn_datadeps` region, +which is what frees its arenas. +""" +mutable struct SlotPool + enabled::Bool + arenas::Dict{CPURAMMemorySpace,Vector{SlotArena}} +end +SlotPool() = SlotPool(false, Dict{CPURAMMemorySpace,Vector{SlotArena}}()) + # N.B. Declared `mutable` (despite no field ever being rebound) purely for # identity: an immutable struct this large is re-boxed on every dynamic call # that passes it, and planning threads it through hundreds of such calls per @@ -753,6 +826,9 @@ mutable struct DataDepsState scratch_deps::Vector{DataDepsTaskDependency} scratch_remote::Vector{Any} + # Slot buffers allocated ahead of need (see `SlotPool`) + slot_pool::SlotPool + function DataDepsState() arg_to_chunk = IdDict{Any,Any}() sparse_arg_wrap = IdDict{Any,Any}() @@ -781,7 +857,7 @@ mutable struct DataDepsState return new(arg_to_chunk, sparse_arg_wrap, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_current, arg_overlaps, ainfo_backing_chunk, supports_inplace_cache, ainfo_cache, ainfos_lookup, ainfos_overlaps, ainfos_owner, ainfos_readers, - TaskArgInfo[], DataDepsTaskDependency[], Any[]) + TaskArgInfo[], DataDepsTaskDependency[], Any[], SlotPool()) end end @@ -1437,8 +1513,20 @@ function generate_slot!(state::DataDepsState, dest_space, data) # original for its aliasing key, exactly as the general path does. aliased_object!(Returns(data), aliased_object_cache, data)::Chunk else - moved = with(DATADEPS_THUNK_ID=>tid) do - remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + direct = direct_uninit_slot(current_acceleration(), data, from_proc, to_proc, orig_space, dest_space) + moved = if direct !== nothing + # What `move_rewrap` would do on the source worker, done from here: + # the object cache lives with the planner, and the slot is only + # allocated, so the source worker has nothing to contribute. + T_dest, header = direct + slot = aliased_object!(aliased_object_cache, data) do _ + alloc_slot_with_aliasing!(state.slot_pool, to_proc, dest_space, T_dest, header) + end + slot::Chunk + else + with(DATADEPS_THUNK_ID=>tid) do + remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + end end if timing move_ns = time_ns() - t0 @@ -1497,6 +1585,273 @@ function remotecall_endpoint_alloc(accel::DistributedAcceleration, to_proc, to_s end end +""" + direct_uninit_slot(accel, data, from_proc, to_proc, from_space, to_space) + -> Union{Tuple{Type,Any},Nothing} + +The type and `alloc_header` of the slot `move_rewrap` would conjure +uninitialized for `data` in `to_space`, when the planning process can tell both +from `data`'s handle alone; `nothing` when only the source worker can. + +Creating a slot through `move_rewrap` runs on the source worker, which then +calls back to the planner for each object cache lookup (the cache lives with +the planner) and on to the destination to allocate: ~6 round trips in series +per slot. When the slot is only allocated, none of that needs the source +worker, so `generate_slot!` does the lookups locally and allocates with one +round trip (`alloc_slot_with_aliasing!`). That was ~90% of a flat-mode launch +loop: a 4096^2 matmul with 512^2 tiles on 4 Distributed workers spent 2.8 s of +a 3.4 s loop creating 461 slots under RoundRobin. + +Only Distributed qualifies: under MPI the cache is replicated and every rank +takes part in each step. +""" +direct_uninit_slot(::Acceleration, data, from_proc, to_proc, from_space, to_space) = nothing +function direct_uninit_slot(::DistributedAcceleration, data::Chunk, from_proc, to_proc, from_space, to_space) + T = chunktype(data) + isconcretetype(T) && T <: DenseArray || return nothing + # A wrapper is rebuilt from its children on the destination + move_rewrap_child_types(T) === nothing || return nothing + T_dest = move_type(from_proc, to_proc, T) + slot_may_be_uninit(from_space, to_space, T_dest) || return nothing + # `tochunk` records `domain(x)`, whose ranges span `size(x)`: what + # `alloc_header` would compute from the array itself. + dom = domain(data) + dom isa ArrayDomain || return nothing + header = map(length, dom.indexes) + header isa alloc_header_type(T_dest) || return nothing + return T_dest, header +end + +""" + alloc_slot_with_aliasing!(pool::SlotPool, to_proc, to_space, ::Type{T}, header) -> Chunk + +`remotecall_endpoint_alloc`, but bringing the new slot's aliasing back with it +and memoizing it, so that planning does not pay another round trip to ask. +With `pool` enabled, a dense CPU slot instead comes out of an arena on its +destination, with no round trip at all once the arena is reserved. +""" +function alloc_slot_with_aliasing!(pool::SlotPool, to_proc, to_space, ::Type{T}, header) where T + if pool.enabled && _arena_slot_type(T, to_space) + return _arena_slot!(pool, to_proc, to_space, T, header) + end + wid = root_worker_id(to_proc) + slot, ainfo = if wid == myid() + _alloc_slot_with_aliasing(to_proc, to_space, T, header) + else + remotecall_fetch(_alloc_slot_with_aliasing, wid, to_proc, to_space, T, header) + end + memoize_ainfo!(ainfo_memo_key(slot, identity), ainfo) + return slot +end +function _alloc_slot_with_aliasing(to_proc, to_space, ::Type{T}, header) where T + x = alloc_uninit(to_space, T, header) + return tochunk(x, to_proc, to_space), aliasing_unwrapped(x, identity) +end + +_arena_slot_type(::Type{<:Array{E}}, ::CPURAMMemorySpace) where {E} = isbitstype(E) +_arena_slot_type(::Type, ::MemorySpace) = false + +# Arenas hand out slots at this alignment. When no plan says how much a region +# needs, they are reserved between these sizes (see `_arena_with_room!`). +const SLOT_ARENA_ALIGN = 64 +const SLOT_ARENA_MIN_BYTES = 64 * 2^10 +const SLOT_ARENA_BYTES = 64 * 2^20 +_arena_aligned(n::Int) = cld(n, SLOT_ARENA_ALIGN) * SLOT_ARENA_ALIGN +_arena_slot_bytes(::Type{<:Array{E}}, header) where {E} = sizeof(E) * prod(header) + +function _arena_slot!(pool::SlotPool, to_proc, to_space::CPURAMMemorySpace, ::Type{T}, header::Dims{N}) where {T<:Array,N} + E = eltype(T) + nbytes = _arena_slot_bytes(T, header) + arena = _arena_with_room!(pool, to_space, nbytes) + offset = arena.used + arena.used += _arena_aligned(nbytes) + ref = ArenaRef{E,N}(root_worker_id(to_space), arena.id, arena.base + offset, header) + # What `tochunk(x, to_proc, to_space)` would make of the array `unwrap(ref)` + slot = Chunk{T,typeof(ref),typeof(to_proc),AnyScope,typeof(to_space)}( + T, ArrayDomain([1:l for l in header]), ref, to_proc, AnyScope(), to_space) + # What `aliasing(::Array)` would compute on the owner + ainfo = ContiguousAliasing(MemorySpan{CPURAMMemorySpace}(RemotePtr{Cvoid}(arena.base + offset, to_space), nbytes)) + memoize_ainfo!(ainfo_memo_key(slot, identity), ainfo) + return slot +end + +function _arena_with_room!(pool::SlotPool, space::CPURAMMemorySpace, nbytes::Int) + arenas = get!(Vector{SlotArena}, pool.arenas, space) + prev = 0 + if !isempty(arenas) + arena = last(arenas) + arena.capacity - arena.used >= nbytes && return arena + prev = arena.capacity + end + # On demand, start small and double: an arena reserved for a few small + # slots would be mostly waste. + want = clamp(max(2prev, 8 * _arena_aligned(nbytes)), SLOT_ARENA_MIN_BYTES, SLOT_ARENA_BYTES) + reserve_slot_arenas!(pool, Dict(space => max(want, _arena_aligned(nbytes)))) + return last(pool.arenas[space]) +end + +"Reserve an arena of `bytes[space]` on each `space`'s worker, concurrently." +function reserve_slot_arenas!(pool::SlotPool, bytes::Dict{CPURAMMemorySpace,Int}) + region = DATADEPS_REGION_ARENAS[]::RegionSlotArenas + spaces = collect(keys(bytes)) + results = Vector{Tuple{UInt,UInt}}(undef, length(spaces)) + @sync for (j, space) in enumerate(spaces) + Threads.@spawn begin + wid = root_worker_id(space) + id, base = wid == myid() ? _reserve_slot_arena(bytes[space]) : + remotecall_fetch(_reserve_slot_arena, wid, bytes[space]) + # Registered as soon as it exists, so that the region frees it even + # if reserving another worker's arena fails + @lock region.lock push!(region.arenas, (wid, id)) + results[j] = (id, base) + end + end + for (j, space) in enumerate(spaces) + id, base = results[j] + push!(get!(Vector{SlotArena}, pool.arenas, space), SlotArena(space, id, base, bytes[space], 0)) + end + return +end + +# Worker side: the arenas this process holds, by id +const SLOT_ARENAS = Dict{UInt,Ptr{Cvoid}}() +const SLOT_ARENAS_LOCK = ReentrantLock() +const SLOT_ARENA_NEXT_ID = Threads.Atomic{UInt}(1) +function _reserve_slot_arena(nbytes::Int) + p = Libc.malloc(nbytes) + p == C_NULL && throw(OutOfMemoryError()) + id = Threads.atomic_add!(SLOT_ARENA_NEXT_ID, UInt(1)) + @lock SLOT_ARENAS_LOCK SLOT_ARENAS[id] = p + return id, UInt(p) +end +function _free_slot_arenas!(ids::Vector{UInt}) + for id in ids + p = @lock SLOT_ARENAS_LOCK pop!(SLOT_ARENAS, id) + Libc.free(p) + end + return +end + +root_worker_id(h::ArenaRef) = h.owner +datasize(h::ArenaRef{T}) where {T} = sizeof(T) * prod(h.dims) +function unwrap(h::ArenaRef{T,N}) where {T,N} + @assert h.owner == myid() "ArenaRef $h is not owned by this process: $(h.owner) != $(myid())" + # Not registered with `_libc_register!`: the arena owns this memory, so + # `unsafe_free!` of the slot alone is a no-op and the arena is freed whole. + # Wrapping needs only the address, not the arena, so a slot outliving its + # region still answers `aliasing` (logs hold slots), like a freed + # `alloc_libc_array` buffer does; reading it then is a use-after-free either way. + return unsafe_wrap(Array{T,N}, Ptr{T}(h.addr), h.dims; own=false) +end +fetch_handle(h::ArenaRef; uniform::Bool=false) = + h.owner == myid() ? unwrap(h) : remotecall_fetch(unwrap, h.owner, h) +MemPool.poolget(h::ArenaRef) = fetch_handle(h) +unwrappable(::ArenaRef) = true +move(from_proc::Processor, to_proc::Processor, x::ArenaRef) = move(from_proc, to_proc, fetch_handle(x)) +move(to_proc::Processor, d::ArenaRef) = move(OSProc(d.owner), to_proc, d) + +""" + prefill_slot_pool!(state, pairs, schedule) + +Enable `state.slot_pool` and reserve on each destination worker an arena with +room for every slot that `schedule` (an AOT plan) will make `generate_slot!` +create through its direct path: one per argument per destination space it is +placed in, counted the way `direct_uninit_slot` decides them. Arguments it +cannot see yet (`DTask`s still running) get arenas reserved on first need. +""" +function prefill_slot_pool!(state::DataDepsState, pairs::Vector{DTaskPair}, schedule::Dict{DTask,Processor}) + accel = current_acceleration() + accel isa DistributedAcceleration || return + # Arenas are only safe where a region will free them + DATADEPS_REGION_ARENAS[] === nothing && return + state.slot_pool.enabled = true + bytes = Dict{CPURAMMemorySpace,Int}() + seen = Set{Tuple{UInt,MemorySpace}}() + for pair in pairs + proc = get(schedule, pair.task, nothing) + proc === nothing && continue + _count_spec_slots!(bytes, seen, accel, pair.spec, only(memory_spaces(proc))) + end + isempty(bytes) || reserve_slot_arenas!(state.slot_pool, bytes) + return +end +# A function barrier: `pair.spec` is abstractly typed (cf. `_collect_spec_remote_chunks!`) +_count_spec_slots!(bytes, seen, accel, spec::DTaskSpec, to_space) = + _foreach_arg(_arg -> _count_arg_slot!(bytes, seen, accel, value(_arg), to_space), spec.fargs) +function _count_arg_slot!(bytes, seen, accel, v, to_space) + arg, deps = unwrap_inout(v) + if arg isa DTask + isready(arg) || return + arg = fetch(arg; raw=true) + end + arg isa Chunk || return + from_space = memory_space(arg) + direct = direct_uninit_slot(accel, arg, first(processors(from_space)), + first(processors(to_space)), from_space, to_space) + direct === nothing && return + T, header = direct + _arena_slot_type(T, to_space) || return + id = (_identity_hash(arg), to_space) + id in seen && return + push!(seen, id) + bytes[to_space] = get(bytes, to_space, 0) + _arena_aligned(_arena_slot_bytes(T, header)) + return +end + +"Whether `slot` lives in a slot arena, which its region frees whole (no free task of its own)." +freed_with_its_arena(slot) = slot isa Chunk && slot.handle isa ArenaRef + +""" + free_region_arenas!(region::RegionSlotArenas) + +Free every arena `region` reserved, with one call per worker. Only valid once +all of the region's tasks have finished (see `spawn_datadeps`). Never throws: +it runs in a `finally`, where an error would hide the region's own; an arena +that cannot be freed (its worker died, say) is reported and dropped. +""" +function free_region_arenas!(region::RegionSlotArenas) + arenas = @lock region.lock copy(region.arenas) + isempty(arenas) && return + by_owner = Dict{Int,Vector{UInt}}() + for (wid, id) in arenas + push!(get!(Vector{UInt}, by_owner, wid), id) + end + @sync for (wid, ids) in by_owner + Threads.@spawn try + wid == myid() ? _free_slot_arenas!(ids) : remotecall_wait(_free_slot_arenas!, wid, ids) + catch err + @warn "Could not free $(length(ids)) Datadeps slot arena(s) on worker $wid" exception=err + end + end + @lock region.lock empty!(region.arenas) + return +end + +""" + free_region_arenas_when_done!(region::RegionSlotArenas) + +`free_region_arenas!`, from a background task, once every task `region` +launched has finished: for a region that failed, whose `wait_all` rethrew the +first task error without waiting for the rest, which may still be using their +slots. The error reaches the caller as soon as it did without arenas. +""" +function free_region_arenas_when_done!(region::RegionSlotArenas) + @lock(region.lock, isempty(region.arenas)) && return + tasks = region.tasks === nothing ? DTask[] : copy(region.tasks) + errormonitor(Threads.@spawn begin + for task in tasks + # A task that was never launched has nothing to wait for + istaskstarted(task) || continue + try + wait(task) + catch + end + end + free_region_arenas!(region) + end) + return +end + function remotecall_endpoint_transfer(f, accel::DistributedAcceleration, from_proc, to_proc, from_space, to_space, data) wid = root_worker_id(to_proc) if wid == myid() diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index b0f364294..bfa065d2f 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -1446,6 +1446,9 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end filtered end + # Arenas for this partition's slots, in the spaces its plan uses (see + # `SlotPool`); freed by the epilogue of `distribute_tasks_hierarchical!`. + prefill_slot_pool!(state, partition_pairs, schedule) # N.B. If this partition throws partway through (e.g. from # `distribute_task!`), any of our vertices that haven't yet been @@ -1546,6 +1549,8 @@ function schedule_partitions_sequential!(queue::DataDepsTaskQueue, local_scopes = Vector{AbstractScope}(undef, n_partitions) proc_to_scope_lfus = [BasicLFUCache{Processor,AbstractScope}(1024) for _ in 1:n_partitions] shared_state = DataDepsState() + # Arenas for the region's slots (see `SlotPool`); freed by the epilogue + prefill_slot_pool!(shared_state, seen_tasks, schedule) write_num = 1 # Uniform execution inverts `BatchedEnqueueQueue`'s latency/throughput @@ -2148,6 +2153,8 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat is_original(obj_cache, remote_space, ainfo) && continue haskey(freed, remote_arg) && continue freed[remote_arg] = nothing + # A slot in an arena is freed with it, when the region ends + freed_with_its_arena(remote_arg) && continue free_syncdeps = Set{ThunkSyncdep}() buf_ainfo = stored_value_ainfo(obj_cache, remote_space, ainfo) gather_free_syncdeps!(state, remote_space, buf_ainfo, remote_arg, diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index a0230f7d6..54d521126 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -76,7 +76,27 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, if !aliasing throw(ArgumentError("Aliasing analysis is no longer optional")) end + # The region's slot arenas (see `SlotPool`) are freed once every task it + # launched has finished, since a failed task would never let a dependent + # free task run. When the region succeeds, `wait_all` has waited for them + # all; when a task fails, it rethrows right away, possibly with others + # still running, so the arenas are freed in the background once those end. + arenas = RegionSlotArenas() + ok = false + try + result = with(DATADEPS_REGION_ARENAS => arenas) do + _spawn_datadeps_region(f, scheduler, launch_wait, hierarchical) + end + ok = true + return result + finally + ok ? free_region_arenas!(arenas) : free_region_arenas_when_done!(arenas) + end +end +function _spawn_datadeps_region(f, scheduler, launch_wait, hierarchical) wait_all(; check_errors=true) do + # Every task the region launches joins this queue + DATADEPS_REGION_ARENAS[].tasks = (get_options(:task_queue)::WaitAllQueue).tasks scheduler = something(scheduler, DATADEPS_SCHEDULER[], RoundRobinScheduler()) launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool hierarchical = something(hierarchical, DATADEPS_HIERARCHICAL[], true)::Bool @@ -138,6 +158,122 @@ function _spec_has_region_dtask_arg(spec::DTaskSpec, region_uids::Set{UInt}) return false end +""" + prefetch_chunk_aliasing!(pairs::Vector{DTaskPair}) + +Resolve the aliasing of every remote `Chunk` argument in `pairs` with one +`remotecall_fetch` per owning worker, issued concurrently, and seed the region's +aliasing memo with the answers. + +Without this, the flat path asks each owner one chunk at a time, the first time +`dag_add_task!` or the launch loop (`aliasing!`) needs a chunk: the memo only +deduplicates a chunk within the region, so a 512-task matmul over 128 remote +tiles paid ~128 sequential round trips, whatever the scheduler. +`build_aliasing_parallel` already batches this way for the hierarchical path. + +Only each chunk's `DRef` travels, and only its aliasing comes back: sending the +`Chunk` (and getting it back in a `Pair`, as `_compute_aliasing_batch` does) +made the exchange ~3x slower, since every `Chunk` carries its type, domain, +processor and scope, and each `DRef` crossing the wire is refcounted again. + +Skips arguments already in the memo, `DTask`s that have not finished (in-region +producers, which fetching would block on), and uniform execution (MPI), whose +`aliasing` is a collective that must stay in submission order on every rank. +""" +function prefetch_chunk_aliasing!(pairs::Vector{DTaskPair}) + memo = CHUNK_AINFO_MEMO[] + memo === nothing && return + uniform_execution(current_acceleration()) && return + candidates = Tuple{Int,DRef,Any,UInt}[] # (owner, handle, dep_mod, memo key) + seen = Set{UInt}() + for pair in pairs + _collect_spec_remote_chunks!(candidates, seen, pair.spec) + end + isempty(candidates) && return + batches = Dict{Int,ChunkAliasingBatch}() + @lock memo.lock for (wid, handle, dep_mod, key) in candidates + haskey(memo.entries, key) && continue + _batch_push!(get!(ChunkAliasingBatch, batches, wid), handle, dep_mod, key) + end + isempty(batches) && return + wids = collect(keys(batches)) + results = Vector{Any}(undef, length(wids)) + @sync for (j, wid) in enumerate(wids) + b = batches[wid] + Threads.@spawn results[j] = remotecall_fetch(_compute_chunk_aliasing_batch, wid, + b.handles, b.which, b.dep_mods) + end + # Seed on the region's own task: the spawned tasks above do not see its + # task-local acceleration, which is part of every memo key. + for (j, wid) in enumerate(wids) + ks = batches[wid].keys + rs = results[j] + @assert length(rs) == length(ks) "prefetch_chunk_aliasing!: $(length(rs)) results for $(length(ks)) args (wid=$wid)" + for i in eachindex(ks) + memoize_ainfo!(ks[i], rs[i]) + end + end + return +end + +# One owner's share of a prefetch. Each `DRef` is listed once, even when the +# region needs it under several `dep_mod`s (`which` points into `handles`): +# MemPool's `DRef` serializer ignores `serialize_cycle_header`'s answer, so the +# same `DRef` twice in one message corrupts it, and the owner errors or hangs. +struct ChunkAliasingBatch + handles::Vector{DRef} + handle_idx::Dict{DRef,Int} + which::Vector{Int} + dep_mods::Vector{Any} + keys::Vector{UInt} +end +ChunkAliasingBatch() = ChunkAliasingBatch(DRef[], Dict{DRef,Int}(), Int[], Any[], UInt[]) +function _batch_push!(b::ChunkAliasingBatch, handle::DRef, dep_mod, key::UInt) + i = get!(b.handle_idx, handle) do + push!(b.handles, handle) + return length(b.handles) + end + push!(b.which, i) + push!(b.dep_mods, dep_mod) + push!(b.keys, key) + return +end + +# Runs on the owner of every handle: what `aliasing(::Chunk, dep_mod)` computes +# there, for a whole batch. `map` rather than an `AbstractAliasing[...]`: when +# every answer has the same type, as for the tiles of an array, the result is +# concretely typed and goes back as one block instead of element by element. +function _compute_chunk_aliasing_batch(handles::Vector{DRef}, which::Vector{Int}, dep_mods::Vector{Any}) + values = map(unwrap, handles) + return map(i -> aliasing_unwrapped(values[which[i]], dep_mods[i]), eachindex(which)) +end + +# A function barrier: `pair.spec` is abstractly typed, and reading `fargs` off +# it dynamically cost more than the rest of collecting the candidates. +_collect_spec_remote_chunks!(candidates, seen, spec::DTaskSpec) = + _foreach_arg(_arg -> _collect_remote_chunk_args!(candidates, seen, value(_arg)), spec.fargs) + +# A function barrier, like `_push_argspecs!`: `value(_arg)` is not always inferrable. +function _collect_remote_chunk_args!(candidates, seen, v) + arg, deps = unwrap_inout(v) + if arg isa DTask + isready(arg) || return + arg = fetch(arg; raw=true) + end + arg isa Chunk || return + # Any other handle is left to `aliasing(::Chunk, dep_mod)`, as before. + arg.handle isa DRef || return + wid = root_worker_id(arg.processor) + wid == myid() && return + for (dep_mod, _, _) in deps + key = ainfo_memo_key(arg, dep_mod) + key in seen && continue + push!(seen, key) + push!(candidates, (wid, arg.handle, dep_mod, key)) + end + return +end + """ datadeps_build_schedule!(scheduler, pairs, all_procs, all_scope; region_uids=nothing, cache=true) -> (dag_spec, schedule) @@ -197,11 +333,10 @@ function datadeps_build_schedule!(scheduler::DataDepsScheduler, found, stale = _schedule_cache_find(scheduler, schedule_cache, dag_spec) if found !== nothing @dagdebug nothing :spawn_datadeps "Found matching DAG spec!" + # `dag_add_task!` recorded each id's task; looking it up there + # rather than scanning `pairs` keeps a cache hit linear in tasks. for (id, proc) in found.id_to_proc - uid = dag_spec.id_to_uid[id] - task_idx = findfirst(spec_task -> spec_task.task.uid == uid, pairs) - task = pairs[task_idx].task - schedule[task] = proc + schedule[dag_spec.id_to_task[id]] = proc end cache_hit = true end @@ -332,6 +467,10 @@ function distribute_tasks!(queue::DataDepsTaskQueue) upper_queue = get_options(:task_queue) + # Both planning and launching need every argument's aliasing; ask for the + # remote chunks' in one batch per owner instead of one round trip each. + prefetch_chunk_aliasing!(queue.seen_tasks) + # Plan the whole region ahead of time, if the scheduler supports it. # Schedulers that don't define `datadeps_schedule_dag_aot!` leave `schedule` # empty and every task falls through to JIT placement below. A planner is @@ -351,6 +490,10 @@ function distribute_tasks!(queue::DataDepsTaskQueue) write_num = 1 proc_to_scope_lfu = BasicLFUCache{Processor,AbstractScope}(1024) launch_start = time_ns() + # Allocate the slots the plan will need up front, one call per worker, so + # the loop below does not wait on one allocation per slot. Counted in the + # launch time, which is what it replaces. + prefill_slot_pool!(state, queue.seen_tasks, schedule) for pair in queue.seen_tasks spec = pair.spec task = pair.task @@ -431,6 +574,8 @@ function distribute_tasks!(queue::DataDepsTaskQueue) is_original(obj_cache, remote_space, ainfo) && continue haskey(freed, remote_arg) && continue freed[remote_arg] = nothing + # A slot in an arena is freed with it, when the region ends + freed_with_its_arena(remote_arg) && continue free_syncdeps = Set{ThunkSyncdep}() buf_ainfo = stored_value_ainfo(obj_cache, remote_space, ainfo) gather_free_syncdeps!(state, remote_space, buf_ainfo, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 1ea158550..1b047373d 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -7,48 +7,16 @@ Base.isempty(spec::DAGSpec) = length(spec) == 0 function dag_add_task!(dspec::DAGSpec, tspec::DTaskSpec, task::DTask) # Check if this task depends on any other tasks within the DAG, - # which we are not yet ready to handle - for (idx, _arg) in enumerate(tspec.fargs) - arg, deps = unwrap_inout(value(_arg)) - for (dep_mod, readdep, writedep) in deps - if arg isa DTask - if arg.uid in keys(dspec.uid_to_id) - # Within-DAG dependency, bail out - return false - end - end - end - end + # which we are not yet ready to handle (a within-DAG dependency: bail out) + _any_arg(_arg -> _depends_on_dag_task(dspec, value(_arg)), tspec.fargs) && return false add_vertex!(dspec.g) id = nv(dspec.g) # Record function signature dspec.id_to_functype[id] = chunktype(tspec.fargs[1]) - argtypes = DatadepsArgSpec[] - for (idx, _arg) in enumerate(tspec.fargs) - arg, deps = unwrap_inout(value(_arg)) - pos = raw_position(_arg) - for (dep_mod, readdep, writedep) in deps - if arg isa DTask - #= TODO: Re-enable this when we can handle within-DAG dependencies - if arg.uid in keys(dspec.uid_to_id) - # Within-DAG dependency - arg_id = dspec.uid_to_id[arg.uid] - push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing())) - add_edge!(dspec.g, arg_id, id) - continue - end - =# - - # External DTask, so fetch this and track it as a raw value - arg = fetch(arg; raw=true) - end - ainfo = aliasing(arg, dep_mod) - # FIXME: Generate syncdeps and add edges - push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, ainfo)) - end - end + argtypes = sizehint!(DatadepsArgSpec[], length(tspec.fargs)) + _foreach_arg(_arg -> _push_argspecs!(argtypes, raw_position(_arg), value(_arg)), tspec.fargs) dspec.id_to_argtypes[id] = argtypes dspec.id_to_scope[id] = @something(tspec.options.compute_scope, tspec.options.scope, @@ -62,6 +30,55 @@ function dag_add_task!(dspec::DAGSpec, tspec::DTaskSpec, task::DTask) return true end + +# Visit a task's arguments. A typed spec's arguments are a heterogeneous `Tuple`, +# whose elements a `for` loop sees as `Any`; `map` unrolls it, so each call is +# specialized on that argument's type (cf. `map_or_ntuple` in `distribute_task!`). +# `f::F` forces specializing on the closure: Julia does not specialize on a +# function argument that is only passed along, so `map` was a dynamic call whose +# result was boxed. +_foreach_arg(f::F, fargs::Vector{Argument}) where {F} = foreach(f, fargs) +_foreach_arg(f::F, fargs::Tuple) where {F} = (map(f, fargs); nothing) +_any_arg(f::F, fargs::Vector{Argument}) where {F} = any(f, fargs) +_any_arg(f::F, fargs::Tuple) where {F} = any(map(f, fargs)) + +# Function barriers for `dag_add_task!`: `value(_arg)` is not inferrable for an +# untyped spec, so each argument is handled by a method specialized on its +# concrete type. Destructuring `unwrap_inout`'s result and iterating its deps on +# an `Any` cost a generic `indexed_iterate`/`getfield` dispatch per dependency, +# and was most of what was left of building a `DAGSpec` once aliasing was +# prefetched. +function _depends_on_dag_task(dspec::DAGSpec, v) + arg, deps = unwrap_inout(v) + arg isa DTask || return false + isempty(deps) && return false + return haskey(dspec.uid_to_id, arg.uid) +end +function _push_argspecs!(argtypes::Vector{DatadepsArgSpec}, pos, v) + arg, deps = unwrap_inout(v) + for (dep_mod, readdep, writedep) in deps + if arg isa DTask + #= TODO: Re-enable this when we can handle within-DAG dependencies + if arg.uid in keys(dspec.uid_to_id) + # Within-DAG dependency + arg_id = dspec.uid_to_id[arg.uid] + push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing())) + add_edge!(dspec.g, arg_id, id) + continue + end + =# + + # External DTask, so fetch this and track it as a raw value + arg = fetch(arg; raw=true) + end + # FIXME: Generate syncdeps and add edges + _push_argspec!(argtypes, pos, arg, dep_mod) + end + return +end +# A second barrier: the result of `fetch(...; raw=true)` is not inferrable either. +_push_argspec!(argtypes::Vector{DatadepsArgSpec}, pos, arg, dep_mod) = + push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, aliasing(arg, dep_mod))) function dag_build_edges!(dag_spec::DAGSpec) # Naively build edges based on exact argument comparisons (no aliasing) arg_writes = Dict{Any,Int}() diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 5cd345b06..5a8c60112 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -712,7 +712,7 @@ function aliasing(x::Chunk, T) if root_worker_id(x.processor) == myid() return aliasing_unwrapped(unwrap(x), T) end - @assert x.handle isa DRef + @assert x.handle isa Union{DRef,ArenaRef} return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T aliasing_unwrapped(unwrap(x), T) end diff --git a/test/datadeps.jl b/test/datadeps.jl index 41fd9e800..23b5215ce 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -81,6 +81,73 @@ end end end +@testset "Flat prefetch of chunk aliasing" begin + # Every remote chunk the tasks name, directly or as a finished DTask's + # result, is resolved in one batch per owner and memoized under the key + # `aliasing` itself would look up. Local chunks are left to `aliasing`, and + # unfinished DTasks too: fetching one would block. Each remote chunk is + # named under two dep_mods, so its owner's batch must carry its `DRef` once + # (the same `DRef` twice in one message used to hang the owner). + remote = Chunk[remotecall_fetch(Dagger.tochunk, wid, zeros(4, 4)) + for wid in procs() if wid != myid()] + done = Dagger.@spawn scope=Dagger.scope(worker=last(procs())) zeros(4, 4) + wait(done) + done_chunk = fetch(done; raw=true) + local_chunk = Dagger.tochunk(zeros(4, 4)) + unlaunched = Dagger.DTask(rand(UInt), Dagger.ThunkFuture(), Dagger.DTaskMetadata(Any)) + task = Dagger.DTask(rand(UInt), Dagger.ThunkFuture(), Dagger.DTaskMetadata(Any)) + pair(fargs) = Dagger.DTaskPair(Dagger.DTaskSpec(fargs, Dagger.Options()), task) + pairs = [pair(Dagger.Argument[Dagger.Argument(1, identity), Dagger.Argument(2, In(c)), + Dagger.Argument(3, Deps(c, InOut(UpperTriangular)))]) + for c in remote] + # A typed spec, as `@spawn` builds for concretely-typed arguments + push!(pairs, pair((Dagger.TypedArgument(1, identity), Dagger.TypedArgument(2, InOut(done)), + Dagger.TypedArgument(3, In(local_chunk)), Dagger.TypedArgument(4, In(unlaunched))))) + expected = Dict{Tuple{Chunk,Any},AbstractAliasing}((c, dep_mod) => aliasing(c, dep_mod) + for c in remote for dep_mod in (identity, UpperTriangular)) + # With no other workers (`-p 0`) the finished task ran here, so its result + # is local and left alone, like every other chunk + if Dagger.root_worker_id(done_chunk.processor) != myid() + expected[(done_chunk, identity)] = aliasing(done_chunk, identity) + end + memo = Dagger.ChunkAinfoMemo() + Dagger.with(Dagger.CHUNK_AINFO_MEMO => memo) do + Dagger.prefetch_chunk_aliasing!(pairs) + @test length(memo.entries) == length(expected) + for ((c, dep_mod), ainfo) in expected + key = Dagger.ainfo_memo_key(c, dep_mod) + # A cache miss must fail, not silently recompute the same answer. + cached = Dagger.memoized_ainfo(() -> error("prefetch did not memoize"), key) + @test cached == ainfo + end + end +end + +@testset "Slot arenas are freed when a region fails" begin + # A region frees its slot arenas once all of its tasks are done, in the + # background when one fails. Freeing them from a task depending on the + # slots' users leaked them whenever one failed, since the scheduler fails a + # task's dependents outright. + if nprocs() >= 3 + arenas_alive() = sum(w -> remotecall_fetch(() -> length(Dagger.SLOT_ARENAS), w), procs()) + arenas_made() = sum(w -> remotecall_fetch(() -> Int(Dagger.SLOT_ARENA_NEXT_ID[]), w), procs()) + chunks = [remotecall_fetch(Dagger.tochunk, w, rand(16, 16)) for w in procs()] + made = arenas_made() + @test_throws Exception Dagger.with(Dagger.DATADEPS_SCHEDULER => Dagger.RoundRobinScheduler(), + Dagger.DATADEPS_HIERARCHICAL => false) do + Dagger.spawn_datadeps() do + for _ in 1:2, c in chunks + Dagger.@spawn (x -> (x .+= 1; nothing))(InOut(c)) + end + Dagger.@spawn (x -> error("failing on purpose"))(InOut(first(chunks))) + end + end + # RoundRobin moved tasks off their chunks' owners, so arenas were used + @test arenas_made() > made + @test timedwait(() -> arenas_alive() == 0, 60; pollint=0.5) === :ok + end +end + @testset "Memory Aliasing" begin A = rand(4) a = Dagger.aliasing(A)