From 81baea0e5853784a2278e26967b934d02fc003ec Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Wed, 14 Oct 2026 10:49:14 -0300 Subject: [PATCH] Datadeps: Allocate uninitialized slots with one round trip Creating a slot went through move_rewrap on the source worker, which calls back to the planning process for every lookup in the aliased object cache and then on to the destination to allocate, about seven round trips in series per slot, and in flat mode that was most of the launch loop: a 4096x4096 matmul with 512x512 tiles on 4 workers spent 2.8 s of a 3.4 s loop creating 461 slots under RoundRobin. When the slot is only going to be allocated, because slot_may_be_uninit lets the copy-to phase fill it, the source worker has nothing to add, so generate_slot! now does the cache lookups itself and allocates with a single round trip to the destination that also brings back the new slot's aliasing. This only applies under Distributed and to dense, non-wrapper arrays whose type and size the chunk already tells, and everything else keeps going through move_rewrap. The RoundRobin launch loop on that matmul gets about 30-40% faster. --- src/datadeps/aliasing.jl | 74 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 72 insertions(+), 2 deletions(-) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index f8de25146..99a0a074b 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1437,8 +1437,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!(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 +1509,64 @@ 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!(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. +""" +function alloc_slot_with_aliasing!(to_proc, to_space, ::Type{T}, header) where T + 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 + 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()