From 76fd60c3adf07e58518b63ae155121af60eefbd6 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 19 Jul 2026 09:57:19 -0700 Subject: [PATCH 1/6] MPI: SPMD Datadeps support, adapted to the NG scheduler Adds MPI acceleration for Dagger with a focus on SPMD Datadeps: MPIAcceleration/MPIRef, deterministic rank-uniform task placement, processor grids, multi-span (KA) copies, move_rewrap header+children protocol, and per-backend GPU MPI/IPC hooks. All MPI.jl-dependent code lives in `ext/MPIExt.jl` (MPI is a weakdep). Scheduler integration is ported onto the new NG scheduler: `schedule_one!` filters/orders processors via the acceleration (`compatible_processors(accel, ...)`, `select_processors_uniform!`) and ignores rank-local capacity/occupancy under uniform execution; `do_task` sets the task acceleration, sequences argument moves (`schedule_argument_move`), and binds moved values via `bind_moved_argument`. This collapses the earlier MPI branch (which evolved over many commits: MPIRPC added then dropped, `src/mpi.jl` moved to the extension, etc.) into its end state, rebased onto current master. The datadeps buffer-free path keeps the upstream `gather_free_syncdeps!` while adding the MPI-uniform task `tag`; `system_uuid` and other changes that landed upstream separately are taken from master. Co-authored-by: Felipe Tome Co-authored-by: yanzin00 Co-authored-by: Cursor --- Project.toml | 3 + ext/CUDAExt.jl | 127 ++- ext/IntelExt.jl | 127 ++- ext/MPIExt.jl | 2134 ++++++++++++++++++++++++++++++++++++ ext/MetalExt.jl | 92 +- ext/OpenCLExt.jl | 51 +- ext/ROCExt.jl | 133 ++- src/Dagger.jl | 12 + src/acceleration.jl | 158 +++ src/array/alloc.jl | 72 +- src/array/darray.jl | 128 +-- src/array/linalg.jl | 30 +- src/array/lu.jl | 2 + src/array/map-reduce.jl | 23 +- src/array/mul.jl | 10 +- src/array/sort.jl | 20 - src/array/trsm.jl | 2 +- src/chunks.jl | 88 +- src/datadeps/aliasing.jl | 305 ++++-- src/datadeps/chunkview.jl | 77 +- src/datadeps/queue.jl | 83 +- src/datadeps/remainders.jl | 69 +- src/dtask.jl | 8 +- src/gpu.jl | 112 ++ src/lib/domain-blocks.jl | 2 + src/memory-spaces.jl | 115 +- src/mutable.jl | 41 + src/options.jl | 9 + src/processor.jl | 6 + src/procgrid.jl | 150 +++ src/queue.jl | 3 +- src/sch/Sch.jl | 149 ++- src/sch/util.jl | 14 +- src/scopes.jl | 2 +- src/shard.jl | 89 ++ src/submission.jl | 47 +- src/thunk.jl | 63 +- src/tochunk.jl | 145 +++ src/types/acceleration.jl | 3 + src/types/chunk.jl | 33 + src/types/memory-space.jl | 1 + src/types/processor.jl | 2 + src/types/scope.jl | 1 + src/utils/chunks.jl | 10 +- src/utils/dagdebug.jl | 7 +- src/utils/haloarray.jl | 36 +- src/utils/span_copy.jl | 338 ++++++ src/weakchunk.jl | 32 + 48 files changed, 4555 insertions(+), 609 deletions(-) create mode 100644 ext/MPIExt.jl create mode 100644 src/acceleration.jl create mode 100644 src/mutable.jl create mode 100644 src/procgrid.jl create mode 100644 src/shard.jl create mode 100644 src/tochunk.jl create mode 100644 src/types/acceleration.jl create mode 100644 src/types/chunk.jl create mode 100644 src/types/memory-space.jl create mode 100644 src/types/processor.jl create mode 100644 src/types/scope.jl create mode 100644 src/utils/span_copy.jl create mode 100644 src/weakchunk.jl diff --git a/Project.toml b/Project.toml index 68b410991..6b4a2e9f2 100644 --- a/Project.toml +++ b/Project.toml @@ -40,6 +40,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" @@ -55,6 +56,7 @@ GraphVizSimpleExt = "Colors" IntelExt = "oneAPI" JSON3Ext = "JSON3" LinuxPerfExt = "LinuxPerf" +MPIExt = "MPI" MetalExt = "Metal" OpenCLExt = "OpenCL" PlotsExt = ["DataFrames", "Plots"] @@ -78,6 +80,7 @@ Graphs = "1" JSON3 = "1" KernelAbstractions = "0.9" LinuxPerf = "0.4.2" +MPI = "0.20.22" MacroTools = "0.5" MemPool = "0.4.17" Metal = "1.1" diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 2f4bff0e5..813763ce1 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -145,13 +145,20 @@ end function Dagger.move!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} sync_with_context(from_space) with_context!(to_space) - copyto!(to, from) + if Dagger.needs_multi_span_copy(to_space, from_space, to, from) + Dagger.multi_span_move!(to, from) + else + copyto!(to, from) + end return end # Out-of-place HtoD function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x) with_context(to_proc) do + if x isa DenseArray && isbitstype(eltype(x)) + Dagger.pin_buffer!(:CUDA, x) + end arr = adapt(CuArray, x) CUDA.synchronize() return arr @@ -163,6 +170,9 @@ function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Chunk) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) with_context(to_proc) do + if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) + Dagger.pin_buffer!(:CUDA, cpu_data) + end arr = adapt(CuArray, cpu_data) CUDA.synchronize() return arr @@ -184,7 +194,8 @@ end function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x) with_context(from_proc) do CUDA.synchronize() - _x = adapt(Array, x) + _x = x isa DenseArray && isbitstype(eltype(x)) ? + Dagger.pinned_host_array(x) : adapt(Array, x) CUDA.synchronize() return _x end @@ -201,8 +212,7 @@ end function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x::CuArray{T,N}) where {T,N} with_context(from_proc) do CUDA.synchronize() - _x = Array{T,N}(undef, size(x)) - copyto!(_x, x) + _x = Dagger.pinned_host_array(x) CUDA.synchronize() return _x end @@ -225,7 +235,7 @@ function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x CUDA.synchronize() return to_arr end - elseif Dagger.system_uuid(from_proc.owner) == Dagger.system_uuid(to_proc.owner) && from_proc.device_uuid == to_proc.device_uuid + elseif Dagger.same_node(Dagger.current_acceleration(), from_proc.owner, to_proc.owner) && from_proc.device_uuid == to_proc.device_uuid # Same node, we can use IPC ipc_handle, eT, shape = remotecall_fetch(from_proc.owner, x) do x arr = unwrap(x) @@ -254,13 +264,14 @@ function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x return arr end else - # Different node, use DtoH, serialization, HtoD + # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x return with_context(from_proc) do - Array(unwrap(x)) + Dagger.pinned_host_array(unwrap(x)) end end return with_context(to_proc) do + Dagger.pin_buffer!(:CUDA, host_copy) return CuArray(host_copy) end end @@ -280,9 +291,10 @@ function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x end else host_copy = with_context(from_proc) do - return Array(x) + return Dagger.pinned_host_array(x) end return with_context(to_proc) do + Dagger.pin_buffer!(:CUDA, host_copy) return CuArray(host_copy) end end @@ -429,6 +441,105 @@ function Dagger.to_scope(::Val{:cuda_gpus}, sc::NamedTuple) end Dagger.scope_key_precedence(::Val{:cuda_gpus}) = 1 +# MPI (SPMD) integration: aliasing spans broadcast from an owner rank must be +# stamped with that rank so same-device addresses on different ranks never +# falsely alias (every rank has myid() == 1 under SPMD) +Dagger.mpi_remap_space(space::CUDAVRAMMemorySpace, owner::Int) = + CUDAVRAMMemorySpace(owner, space.device, space.device_uuid) +# Acceleration-free memory space of a raw value, for chunk record labeling +Dagger.value_memory_space(x::CuArray) = Dagger.memory_space(x) + +# MPI data plane: pass CuArrays to MPI directly when the library is CUDA-aware +# (DAGGER_MPI_GPU_DIRECT=0/1 forces the decision; detection is best-effort) +const MPI_GPU_DIRECT = Ref{Union{Nothing,Bool}}(nothing) +function mpi_gpu_direct_enabled() + v = MPI_GPU_DIRECT[] + v !== nothing && return v + env = get(ENV, "DAGGER_MPI_GPU_DIRECT", "") + v = if !isempty(env) + something(tryparse(Bool, env), false) + else + Dagger.mpi_library_gpu_aware(:CUDA) + end + MPI_GPU_DIRECT[] = v + return v +end +Dagger.mpi_device_direct(x::CUDA.StridedCuArray) = mpi_gpu_direct_enabled() +# Device-wide sync: producers/consumers run on other tasks' streams, so a +# current-stream synchronize is not enough at the communication boundary +# (device_synchronize yields to the Julia scheduler while waiting) +Dagger.mpi_device_sync(x::CuArray) = CUDA.device_synchronize() +Dagger.mpi_device_sync(::CUDAVRAMMemorySpace) = CUDA.device_synchronize() + +# Page-lock host staging buffers (~2x DtoH/HtoD bandwidth); CUDA.pin +# registers a GC finalizer that unregisters the memory +Dagger.gpu_memory_kind(::CuArray) = :CUDA +Dagger.gpu_memory_kind(::CUDAVRAMMemorySpace) = :CUDA +Dagger.pin_buffer!(::Val{:CUDA}, buf::DenseArray) = (CUDA.pin(buf); nothing) + +# Same-node device IPC: pool-backed CuArrays are not cuIpc-exportable, so the +# sender stages into a direct (cuMemAlloc) allocation and ships its 64-byte +# handle; the receiver maps it and copies device-to-device. +# DAGGER_IPC=0 disables the path. +const GPU_IPC = Ref{Union{Nothing,Bool}}(nothing) +function ipc_enabled() + v = GPU_IPC[] + v !== nothing && return v + v = something(tryparse(Bool, get(ENV, "DAGGER_IPC", "true")), true) + GPU_IPC[] = v + return v +end +Dagger.ipc_eligible(::CUDAVRAMMemorySpace, ::CUDAVRAMMemorySpace) = ipc_enabled() + +# The driver API wrappers live in CUDA itself on 5.x and in CUDACore +# (DeviceMemory's home module) on 6.x +const CUDADRV = isdefined(CUDA, :cuIpcGetMemHandle) ? CUDA : + parentmodule(CUDA.DeviceMemory) + +struct CuIpcInfo{T,N} + handle::CUDADRV.CUipcMemHandle + shape::Dims{N} +end +function Dagger.ipc_export(value::CUDA.StridedCuArray{T,N}) where {T,N} + # Activate the source device — MPI execute! may have bound the dest context + with_context!(Dagger.memory_space(value)) + buf = CUDADRV.alloc(CUDA.DeviceMemory, sizeof(value)) + staged = unsafe_wrap(CuArray, convert(CUDA.CuPtr{T}, buf.ptr), size(value)) + copyto!(staged, value) + CUDA.device_synchronize() + handle_ref = Ref{CUDADRV.CUipcMemHandle}() + CUDADRV.cuIpcGetMemHandle(handle_ref, convert(CUDA.CuPtr{Nothing}, buf.ptr)) + return CuIpcInfo{T,N}(handle_ref[], size(value)), buf +end +Dagger.ipc_release!(buf::CUDA.DeviceMemory) = CUDADRV.free(buf) +function _ipc_open(info::CuIpcInfo{T,N}) where {T,N} + ptr_ref = Ref{CUDA.CuPtr{Cvoid}}() + if isdefined(CUDADRV, :cuIpcOpenMemHandle_v2) + CUDADRV.cuIpcOpenMemHandle_v2(ptr_ref, info.handle, CUDADRV.CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS) + else + CUDADRV.cuIpcOpenMemHandle(ptr_ref, info.handle, CUDADRV.CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS) + end + src = unsafe_wrap(CuArray, convert(CUDA.CuPtr{T}, ptr_ref[]), info.shape) + return src, ptr_ref[] +end +function Dagger.ipc_copyto!(dest::CUDA.StridedCuArray{T,N}, info::CuIpcInfo{T,N}) where {T,N} + @assert size(dest) == info.shape "IPC shape mismatch: $(size(dest)) != $(info.shape)" + with_context!(Dagger.memory_space(dest)) + src, raw = _ipc_open(info) + try + copyto!(dest, src) + CUDA.device_synchronize() + finally + CUDA.cuIpcCloseMemHandle(raw) + end + return dest +end +function Dagger.ipc_materialize(info::CuIpcInfo{T,N}) where {T,N} + # Caller must have set the destination device context (ipc_copyto! does) + dest = CuArray{T,N}(undef, info.shape) + return Dagger.ipc_copyto!(dest, info) +end + const DEVICES = Dict{Int, CuDevice}() const CONTEXTS = Dict{Int, CuContext}() const STREAMS = Dict{Int, CuStream}() diff --git a/ext/IntelExt.jl b/ext/IntelExt.jl index e5af5e64d..87bb7138f 100644 --- a/ext/IntelExt.jl +++ b/ext/IntelExt.jl @@ -31,7 +31,7 @@ Dagger.get_parent(proc::oneArrayDeviceProc) = Dagger.OSProc(proc.owner) Dagger.root_worker_id(proc::oneArrayDeviceProc) = proc.owner Base.show(io::IO, proc::oneArrayDeviceProc) = print(io, "oneArrayDeviceProc(worker $(proc.owner), device $(proc.device_id))") -Dagger.short_name(proc::oneArrayDeviceProc) = "W: $(proc.owner), oneAPI: $(proc.device)" +Dagger.short_name(proc::oneArrayDeviceProc) = "W: $(proc.owner), oneAPI: $(proc.device_id)" Dagger.@gpuproc(oneArrayDeviceProc, oneArray) "Represents the memory space of a single Intel GPU's VRAM." @@ -55,6 +55,16 @@ function Dagger.aliasing(x::oneArray{T}) where T return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end +# MPI (SPMD) integration: stamp owning rank on VRAM spaces +Dagger.mpi_remap_space(space::IntelVRAMMemorySpace, owner::Int) = + IntelVRAMMemorySpace(owner, space.device_id) +Dagger.value_memory_space(x::oneArray) = Dagger.memory_space(x) + +# Staging-pool key; Level Zero has no CUDA-style pin of arbitrary host buffers +Dagger.gpu_memory_kind(::oneArray) = :oneAPI +Dagger.gpu_memory_kind(::IntelVRAMMemorySpace) = :oneAPI +Dagger.pin_buffer!(::Val{:oneAPI}, buf::DenseArray) = nothing + function Dagger.unsafe_free!(x::oneArray) oneAPI.unsafe_free!(x) return @@ -161,13 +171,20 @@ end function Dagger.move!(to_space::IntelVRAMMemorySpace, from_space::IntelVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} sync_with_context(from_space) with_context!(to_space) - copyto!(to, from) + if Dagger.needs_multi_span_copy(to_space, from_space, to, from) + Dagger.multi_span_move!(to, from) + else + copyto!(to, from) + end return end # Out-of-place HtoD function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x) with_context(to_proc) do + if x isa DenseArray && isbitstype(eltype(x)) + Dagger.pin_buffer!(:oneAPI, x) + end arr = adapt(oneArray, x) oneAPI.synchronize() return arr @@ -179,6 +196,9 @@ function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::Chunk) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) with_context(to_proc) do + if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) + Dagger.pin_buffer!(:oneAPI, cpu_data) + end arr = adapt(oneArray, cpu_data) oneAPI.synchronize() return arr @@ -192,7 +212,8 @@ end function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::CPUProc, x) with_context(from_proc) do oneAPI.synchronize() - _x = adapt(Array, x) + _x = x isa DenseArray && isbitstype(eltype(x)) ? + Dagger.pinned_host_array(x) : adapt(Array, x) oneAPI.synchronize() return _x end @@ -209,8 +230,7 @@ end function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::CPUProc, x::oneArray{T,N}) where {T,N} with_context(from_proc) do oneAPI.synchronize() - _x = Array{T,N}(undef, size(x)) - copyto!(_x, x) + _x = Dagger.pinned_host_array(x) oneAPI.synchronize() return _x end @@ -229,18 +249,34 @@ function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::oneArrayDeviceProc, with_context(oneAPI.synchronize, from_proc) return _ensure_device(from_arr, to_proc.device_id) else - # Different node, use DtoH, serialization, HtoD - host_copy = remotecall_fetch(from_proc.owner, x) do x - Array(unwrap(x)) + # Different node, use DtoH, serialization, HtoD (pinned host staging) + host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x + return with_context(from_proc) do + Dagger.pinned_host_array(unwrap(x)) + end end return with_context(to_proc) do - oneArray(host_copy) + Dagger.pin_buffer!(:oneAPI, host_copy) + return oneArray(host_copy) end end end function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::oneArrayDeviceProc, x::oneArray) - return _ensure_device(x, to_proc.device_id) + if from_proc == to_proc + with_context(oneAPI.synchronize, from_proc) + return x + elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) + return _ensure_device(x, to_proc.device_id) + else + host_copy = with_context(from_proc) do + return Dagger.pinned_host_array(x) + end + return with_context(to_proc) do + Dagger.pin_buffer!(:oneAPI, host_copy) + return oneArray(host_copy) + end + end end # Adapt generic functions @@ -382,6 +418,77 @@ function Dagger.to_scope(::Val{:intel_gpus}, sc::NamedTuple) end Dagger.scope_key_precedence(::Val{:intel_gpus}) = 1 +# MPI data plane: pass oneArrays to MPI directly when the library is Level Zero / +# Intel-MPI aware (DAGGER_MPI_GPU_DIRECT=0/1 forces; else mpi_library_gpu_aware) +const MPI_GPU_DIRECT = Ref{Union{Nothing,Bool}}(nothing) +function mpi_gpu_direct_enabled() + v = MPI_GPU_DIRECT[] + v !== nothing && return v + env = get(ENV, "DAGGER_MPI_GPU_DIRECT", "") + v = if !isempty(env) + something(tryparse(Bool, env), false) + else + Dagger.mpi_library_gpu_aware(:oneAPI) + end + MPI_GPU_DIRECT[] = v + return v +end +Dagger.mpi_device_direct(x::oneAPI.oneStridedArray) = mpi_gpu_direct_enabled() +Dagger.mpi_device_sync(x::oneArray) = oneAPI.synchronize() +Dagger.mpi_device_sync(::IntelVRAMMemorySpace) = oneAPI.synchronize() + +# Same-node device IPC via Level Zero zeMem*IpcHandle. Stage into a dedicated +# device_alloc (not the array pool) so the handle stays valid across processes. +# DAGGER_IPC=0 disables the path. +const GPU_IPC = Ref{Union{Nothing,Bool}}(nothing) +function ipc_enabled() + v = GPU_IPC[] + v !== nothing && return v + v = something(tryparse(Bool, get(ENV, "DAGGER_IPC", "true")), true) + GPU_IPC[] = v + return v +end +Dagger.ipc_eligible(::IntelVRAMMemorySpace, ::IntelVRAMMemorySpace) = ipc_enabled() + +const oneL0 = oneAPI.oneL0 + +struct ZeIpcInfo{T,N} + handle::oneL0.ze_ipc_mem_handle_t + shape::Dims{N} + bytesize::Int +end +function Dagger.ipc_export(value::oneAPI.oneStridedArray{T,N}) where {T,N} + # Stage into a fresh device allocation so the IPC handle outlives the source + with_context!(Dagger.memory_space(value)) + staged = oneArray{T,N}(undef, size(value)) + copyto!(staged, value) + oneAPI.synchronize() + handle_ref = Ref{oneL0.ze_ipc_mem_handle_t}() + oneL0.zeMemGetIpcHandle(oneAPI.context(), pointer(staged), handle_ref) + return ZeIpcInfo{T,N}(handle_ref[], size(value), sizeof(value)), staged +end +Dagger.ipc_release!(staged::oneArray) = oneAPI.unsafe_free!(staged) +function Dagger.ipc_copyto!(dest::oneAPI.oneStridedArray{T,N}, info::ZeIpcInfo{T,N}) where {T,N} + @assert size(dest) == info.shape "IPC shape mismatch: $(size(dest)) != $(info.shape)" + with_context!(Dagger.memory_space(dest)) + ctx = oneAPI.context() + dev = oneAPI.device() + ptr_ref = Ref{Ptr{Cvoid}}() + oneL0.zeMemOpenIpcHandle(ctx, dev, info.handle, UInt32(0), ptr_ref) + try + src_ptr = reinterpret(oneL0.ZePtr{T}, ptr_ref[]) + Base.unsafe_copyto!(ctx, dev, pointer(dest), src_ptr, length(dest)) + oneAPI.synchronize() + finally + oneL0.zeMemCloseIpcHandle(ctx, ptr_ref[]) + end + return dest +end +function Dagger.ipc_materialize(info::ZeIpcInfo{T,N}) where {T,N} + dest = oneArray{T,N}(undef, info.shape) + return Dagger.ipc_copyto!(dest, info) +end + const DEVICES = Dict{Int, ZeDevice}() const DRIVERS = Dict{Int, ZeDriver}() const CONTEXTS = IdDict{ZeDriver, ZeContext}() diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl new file mode 100644 index 000000000..69ed8c9f5 --- /dev/null +++ b/ext/MPIExt.jl @@ -0,0 +1,2134 @@ +module MPIExt + +import Dagger +import Dagger: Thunk, Chunk, Processor, MemorySpace +import Dagger: @dagdebug, @opcounter + +# Types and generic functions from Dagger that this extension either +# extends with new methods (for MPI-specific types) or calls directly. +import Dagger: + AbstractAliasing, accelerate!, accel_matches_proc, aliased_object!, + AliasedObjectCache, AliasedObjectCacheStore, aliasing, bind_moved_argument, + chunktype, ChunkView, check_uniform, check_uniformity!, CHECK_UNIFORMITY, + cleanup_tasks_accel!, constrain, CPURAMMemorySpace, + current_acceleration, CyclicProcGrid, datasize, default_enabled, + default_memory_space, default_processor, default_procgrid, + finalize_acceleration!, execute!, fetch_handle, fire_order_key, get_parent, + get_processors, gpu_kernel_backend, gpu_memory_kind, + initialize_acceleration!, InvalidScope, ipc_copyto!, ipc_eligible, + ipc_export, ipc_materialize, IPC_MIN_BYTES, ipc_release!, is_local, + istask, LockedObject, memory_space, memory_spaces, MemorySpan, + memory_spans, move, move!, move_rewrap, move_rewrap_build, DATADEPS_THUNK_ID, + mpi_device_direct, mpi_device_sync, mpi_library_gpu_aware, mpi_remap_space, + myid, + move_rewrap_child_types, move_rewrap_header_mode, move_rewrap_parts, + move_rewrap_result_type, move_type, multi_span_gather!, + multi_span_scatter!, next_id, NoAliasing, Options, OSProc, + post_stage_array_chunks!, processors, ProcessScope, proc_in_scope, + remotecall_endpoint_toplevel, RemotePtr, root_worker_id, same_node, + schedule_argument_move, sch_handle, select_processors_uniform!, set_key_stored!, + set_stored!, short_name, span_len, stage_acquire!, stage_release!, + stage_to_host!, system_uuid, task_processor_preference, ThreadProc, to_tag, + tochunk, tochunk_pset, uniform_execution, UnknownAliasing, unwrap, + unwrap_weak_checked, value_memory_space, WeakChunk, with_context!, + _with_default_acceleration + +import MemPool +import MemPool: DRef, poolget, poolset + +# `Base.ScopedValues` only exists on Julia 1.11+; on 1.10 (LTS) fall back to the +# `ScopedValues` package (a Dagger dependency), matching `src/Dagger.jl`. +if !isdefined(Base, :ScopedValues) + import ScopedValues: ScopedValue, @with, with +else + import Base.ScopedValues: ScopedValue, @with, with +end + +import TaskLocalValues: TaskLocalValue + +import SparseArrays: SparseMatrixCSC + +import Random +import Serialization + +using MPI + +function check_uniform(value::Integer, original=value) + CHECK_UNIFORMITY[] && uniform_execution() || return true + comm = MPI.COMM_WORLD + rank = MPI.Comm_rank(comm) + matched = compare_all(value, comm) + if !matched + if rank == 0 + Core.print("[$rank] Found non-uniform value!\n") + end + Core.print("[$rank] value=$value, original=$original\n") + throw(ArgumentError("Non-uniform value")) + end + MPI.Barrier(comm) + return matched +end + +# MPI tag reserved for `compare_all` / `check_uniform` P2P only. It must not +# collide with ordinary Dagger message tags on the same `comm` (which come from +# `to_tag` = thunk/ref IDs, and which `remotecall_endpoint_toplevel` broadcast +# metadata uses `0` for), or ranks steal each other's messages and hang until +# `mpi_deadlock_detect` fires. +# +# Use the top of the MPI tag space (`MPI.tag_ub()`): `take_ref_id!` asserts every +# ordinary tag is strictly `< MPI.tag_ub()`, so this value provably cannot +# collide with one. `tag_ub()` is a valid tag (the MPI standard allows tags in +# `0:MPI_TAG_UB`) and is a library constant identical on every rank, which the +# cross-rank uniformity sends/receives require. (Computed lazily rather than as a +# load-time `const`, since MPI is not yet initialized at module load.) +compare_all_mpi_tag() = UInt32(MPI.tag_ub()) + +function compare_all(value, comm) + rank = MPI.Comm_rank(comm) + size = MPI.Comm_size(comm) + tag = compare_all_mpi_tag() + for i in 0:(size-1) + if i != rank + send_yield(value, comm, i, tag) + end + end + match = true + for i in 0:(size-1) + if i != rank + other_value = recv_yield(comm, i, tag) + if value != other_value + match = false + end + end + end + return match +end + +""" + sync_mpi_rng!(comm) + +Broadcast a seed from rank 0 and `Random.seed!` it on every rank so default +RNGs are coherent under SPMD MPI acceleration. +""" +function sync_mpi_rng!(comm::MPI.Comm) + seed = Ref{UInt64}(0) + if MPI.Comm_rank(comm) == 0 + seed[] = rand(Random.RandomDevice(), UInt64) + end + MPI.Bcast!(seed, 0, comm) + Random.seed!(seed[]) + return seed[] +end + +""" + check_mpi_rng_coherent!(comm) + +Probe `rand(UInt64)` on every rank and throw if any rank differs. Always runs +(independent of `CHECK_UNIFORMITY[]`) because coherent RNGs are required for +SPMD correctness after MPI acceleration init. +""" +function check_mpi_rng_coherent!(comm::MPI.Comm) + probe = rand(UInt64) + if !compare_all(probe, comm) + throw(ArgumentError("Non-uniform RNG under MPI acceleration")) + end + return true +end + +struct MPIAcceleration <: Dagger.Acceleration + comm::MPI.Comm +end +MPIAcceleration() = MPIAcceleration(MPI.COMM_WORLD) + +function aliasing(accel::MPIAcceleration, x::Chunk, T) + handle = x.handle::MPIRef + @assert accel.comm == handle.comm "MPIAcceleration comm mismatch" + tag = to_tag() + check_uniform(tag) + rank = MPI.Comm_rank(accel.comm) + if handle.rank == rank + ainfo = _with_default_acceleration() do + aliasing(x, T) + end + ainfo = mpi_remap_ainfo(ainfo, handle.rank) + @opcounter :aliasing_bcast_send_yield + ainfo = bcast_yield(accel.comm, handle.rank, tag, ainfo) + else + ainfo = bcast_yield(accel.comm, handle.rank, tag) + end + check_uniform(ainfo) + return ainfo +end + +default_processor(accel::MPIAcceleration) = MPIOSProc(accel.comm, 0) +default_processor(accel::MPIAcceleration, x) = MPIOSProc(accel.comm, 0) +default_processor(accel::MPIAcceleration, x::Chunk) = MPIOSProc(x.handle.comm, x.handle.rank) +default_processor(accel::MPIAcceleration, x::Function) = MPIOSProc(accel.comm, MPI.Comm_rank(accel.comm)) +default_processor(accel::MPIAcceleration, T::Type) = MPIOSProc(accel.comm, MPI.Comm_rank(accel.comm)) +uniform_execution(accel::MPIAcceleration) = true + +# Children / uniform-processor caches are shared across reusable scheduler +# tasks. Protect both with one lock; never mutate them unlocked. +const MPI_PROC_CACHE_LOCK = Threads.ReentrantLock() +const MPIClusterProcChildren = Dict{MPI.Comm, Set{Processor}}() +const MPIUniformProcsCache = Dict{MPI.Comm, Vector{Processor}}() + +struct MPIClusterProc <: Processor + comm::MPI.Comm + function MPIClusterProc(comm::MPI.Comm) + ensure_children!(comm) + return new(comm) + end +end + +Dagger.Sch.init_proc(state, proc::MPIClusterProc, log_sink) = Dagger.Sch.init_proc(state, MPIOSProc(proc.comm), log_sink) + +MPIClusterProc() = MPIClusterProc(MPI.COMM_WORLD) + +# Populate `MPIClusterProcChildren[comm]` once. Safe on the scheduler hot path; +# does not clear the uniform-processor cache on every lookup. +function ensure_children!(comm::MPI.Comm) + # `@lock` (not `lock(...) do`) so the early `return` on a cache hit actually + # returns from `ensure_children!`; inside a `do` closure it would only + # return from the closure and we'd needlessly re-fetch the children below. + @lock MPI_PROC_CACHE_LOCK begin + haskey(MPIClusterProcChildren, comm) && return MPIClusterProcChildren[comm] + end + # Fetch outside the lock: `get_processors(OSProc())` may remotecall / take + # OSPROC_PROCESSOR_CACHE and must not run under MPI_PROC_CACHE_LOCK. + children = get_processors(OSProc()) + @lock MPI_PROC_CACHE_LOCK begin + if !haskey(MPIClusterProcChildren, comm) + MPIClusterProcChildren[comm] = children + # Invalidate only this communicator's uniform listing + delete!(MPIUniformProcsCache, comm) + end + return MPIClusterProcChildren[comm] + end +end + +# Force-refresh children for `comm` (e.g. after processor callbacks change). +function populate_children!(comm::MPI.Comm) + children = get_processors(OSProc()) + lock(MPI_PROC_CACHE_LOCK) do + MPIClusterProcChildren[comm] = children + delete!(MPIUniformProcsCache, comm) + end + return children +end + +function uniform_mpi_processors(accel::MPIAcceleration) + # Ensure children first so the `get!` callback never mutates this Dict + # mid-insertion (via MPIClusterProc / ensure_children!). + ensure_children!(accel.comm) + lock(MPI_PROC_CACHE_LOCK) do + get!(MPIUniformProcsCache, accel.comm) do + children = MPIClusterProcChildren[accel.comm] + procs = Processor[] + for i in 0:(MPI.Comm_size(accel.comm)-1) + for innerProc in children + push!(procs, MPIProcessor(innerProc, accel.comm, i)) + end + end + # Default placement only targets default-enabled (CPU) processors; GPU + # processors are opt-in via explicit scopes + filter!(default_enabled, procs) + select_processors_uniform!(procs, accel) + procs + end + end +end + +struct MPIOSProc <: Processor + comm::MPI.Comm + rank::Int +end + +function MPIOSProc(comm::MPI.Comm) + rank = MPI.Comm_rank(comm) + return MPIOSProc(comm, rank) +end + +function MPIOSProc() + return MPIOSProc(MPI.COMM_WORLD) +end + +ProcessScope(p::MPIOSProc) = ProcessScope(myid()) + +function check_uniform(proc::MPIOSProc, original=proc) + return check_uniform(hash(MPIOSProc), original) && + check_uniform(proc.rank, original) +end + +function memory_spaces(proc::MPIOSProc) + children = get_processors(proc) + spaces = Set{MemorySpace}() + for proc in children + for space in memory_spaces(proc) + push!(spaces, space) + end + end + return spaces +end + +struct MPIProcessScope <: Dagger.AbstractScope + comm::MPI.Comm + rank::Int +end + +Base.isless(::MPIProcessScope, ::MPIProcessScope) = false +Base.isless(::MPIProcessScope, ::Dagger.ProcessScope) = true +Base.isless(::MPIProcessScope, ::Dagger.NodeScope) = true +Base.isless(::MPIProcessScope, ::Dagger.UnionScope) = true +Base.isless(::MPIProcessScope, ::Dagger.TaintScope) = true +Base.isless(::MPIProcessScope, ::Dagger.AnyScope) = true +Base.isless(::Dagger.ProcessScope, ::MPIProcessScope) = false +Base.isless(::Dagger.NodeScope, ::MPIProcessScope) = false +constrain(x::MPIProcessScope, y::MPIProcessScope) = + x == y ? y : InvalidScope(x, y) +# Under MPI, every processor reports root_worker_id == myid(), so a local +# ProcessScope intersects any MPIProcessScope as the more specific rank scope +constrain(x::Dagger.ProcessScope, y::MPIProcessScope) = + x.wid == myid() ? y : InvalidScope(x, y) +constrain(x::Dagger.NodeScope, y::MPIProcessScope) = + x == Dagger.NodeScope() ? y : Dagger.InvalidScope(x, y) + +Base.isless(::Dagger.ExactScope, ::MPIProcessScope) = true +constrain(x::MPIProcessScope, y::Dagger.ExactScope) = + (y.processor isa MPIProcessor && + y.processor.comm == x.comm && + y.processor.rank == x.rank) ? y : Dagger.InvalidScope(x, y) + +function enclosing_scope(proc::MPIOSProc) + return MPIProcessScope(proc.comm, proc.rank) +end + +function Dagger.to_scope(::Val{:mpi_rank}, sc::NamedTuple) + if sc.mpi_rank == Colon() + return Dagger.to_scope(Val{:mpi_ranks}(), merge(sc, (;mpi_ranks=Colon()))) + else + @assert sc.mpi_rank isa Integer "Expected a single MPI rank for :mpi_rank, got $(sc.mpi_rank)\nConsider using :mpi_ranks instead." + return Dagger.to_scope(Val{:mpi_ranks}(), merge(sc, (;mpi_ranks=[sc.mpi_rank]))) + end +end +Dagger.scope_key_precedence(::Val{:mpi_rank}) = 2 +function Dagger.to_scope(::Val{:mpi_ranks}, sc::NamedTuple) + comm = get(sc, :mpi_comm, MPI.COMM_WORLD) + # `:mpi_ranks` is `Colon()` to select every rank, otherwise an iterable of + # rank indices. (`mpi_rank=r` funnels through here as `mpi_ranks=[r]`.) + if sc.mpi_ranks == Colon() + ranks = 0:(MPI.Comm_size(comm)-1) + else + ranks = sc.mpi_ranks + end + scopes = Dagger.ExactScope[] + for rank in ranks + procs = Dagger.get_processors(MPIOSProc(comm, rank)) + rank_scope = MPIProcessScope(comm, rank) + for proc in procs + proc_scope = Dagger.ExactScope(proc) + constrain(proc_scope, rank_scope) isa Dagger.InvalidScope && continue + push!(scopes, proc_scope) + end + end + return Dagger.UnionScope(scopes) +end +Dagger.scope_key_precedence(::Val{:mpi_ranks}) = 2 + +struct MPIProcessor{P<:Processor} <: Processor + innerProc::P + comm::MPI.Comm + rank::Int +end +proc_in_scope(proc::Processor, scope::MPIProcessScope) = false +proc_in_scope(proc::MPIProcessor, scope::MPIProcessScope) = + proc.comm == scope.comm && proc.rank == scope.rank + +function check_uniform(proc::MPIProcessor, original=proc) + return check_uniform(hash(MPIProcessor), original) && + check_uniform(proc.rank, original) && + # Compare the logical identity (kind + ordinal), not the raw struct: + # device handles/UUIDs of GPU processors are rank-local + check_uniform(hash(short_name(proc.innerProc)), original) +end + +Dagger.iscompatible_func(::MPIProcessor, opts, ::Any) = true +Dagger.iscompatible_arg(::MPIProcessor, opts, ::Any) = true + +# Under uniform (SPMD) execution every rank must run each task on the rank it +# was deterministically scheduled to; local work stealing would diverge +Dagger.Sch.stealing_permitted(::MPIProcessor) = false + +default_enabled(proc::MPIProcessor) = default_enabled(proc.innerProc) + +root_worker_id(proc::MPIProcessor) = myid() +root_worker_id(proc::MPIOSProc) = myid() +root_worker_id(proc::MPIClusterProc) = myid() + +get_parent(proc::MPIClusterProc) = proc +get_parent(proc::MPIOSProc) = MPIClusterProc(proc.comm) +get_parent(proc::MPIProcessor) = MPIOSProc(proc.comm, proc.rank) + +# Rank-deterministic keys: default fire_order_key uses root_worker_id, which +# is always myid() for MPI processors and is not comparable across ranks. +function fire_order_key(proc::MPIClusterProc) + wid = root_worker_id(proc) + return (wid, wid) +end +fire_order_key(proc::MPIOSProc) = (proc.rank, 0) +fire_order_key(proc::MPIProcessor) = (get_parent(proc).rank, proc.rank) + +function task_processor_preference(accel::MPIAcceleration, task::Thunk, state) + # N.B. inputs[1] is the function, whose chunk (when wrapped) is deliberately + # owned by the local rank on every rank — only data arguments carry + # rank-uniform ownership + for arg in task.inputs[2:end] + arg_value = unwrap_weak_checked(Dagger.value(arg)) + arg_chunk = if istask(arg_value) + Dagger.Sch.load_result(state, arg_value) + elseif arg_value isa Chunk + arg_value + else + nothing + end + if arg_chunk isa Chunk && arg_chunk.handle isa MPIRef + return arg_chunk.handle.rank + end + end + return nothing +end + +short_name(proc::MPIProcessor) = "(MPI: $(proc.rank), $(short_name(proc.innerProc)))" + +function get_processors(mosProc::MPIOSProc) + children = ensure_children!(mosProc.comm) + mpiProcs = Set{Processor}() + for proc in children + push!(mpiProcs, MPIProcessor(proc, mosProc.comm, mosProc.rank)) + end + return mpiProcs +end + +#TODO: non-uniform ranking through MPI groups +#TODO: use a lazy iterator +function get_processors(proc::MPIClusterProc) + children = ensure_children!(proc.comm) + result = Set{Processor}() + for i in 0:(MPI.Comm_size(proc.comm)-1) + for innerProc in children + push!(result, MPIProcessor(innerProc, proc.comm, i)) + end + end + return result +end + +struct MPIMemorySpace{S<:MemorySpace} <: MemorySpace + innerSpace::S + comm::MPI.Comm + rank::Int +end + +function Base.:(==)(a::MPIMemorySpace, b::MPIMemorySpace) + return a.innerSpace == b.innerSpace && + a.comm == b.comm && + a.rank == b.rank +end +Base.hash(space::MPIMemorySpace, h::UInt=UInt(0)) = + hash(space.innerSpace, hash(space.comm.val, hash(space.rank, hash(MPIMemorySpace, h)))) + +function check_uniform(space::MPIMemorySpace, original=space) + return check_uniform(space.rank, original) && + # Compare the logical identity (kind + ordinal), not the raw struct: + # device handles/UUIDs of GPU memory spaces are rank-local + check_uniform(hash(short_name(space.innerSpace)), original) +end + +short_name(space::MPIMemorySpace) = "(MPI: $(space.rank), $(short_name(space.innerSpace)))" + +default_processor(space::MPIMemorySpace) = MPIOSProc(space.comm, space.rank) +default_memory_space(accel::MPIAcceleration) = MPIMemorySpace(CPURAMMemorySpace(myid()), accel.comm, 0) + +default_memory_space(accel::MPIAcceleration, x) = MPIMemorySpace(value_memory_space(x), accel.comm, 0) +default_memory_space(accel::MPIAcceleration, x::Chunk) = MPIMemorySpace(CPURAMMemorySpace(myid()), x.handle.comm, x.handle.rank) +default_memory_space(accel::MPIAcceleration, x::Function) = MPIMemorySpace(CPURAMMemorySpace(myid()), accel.comm, MPI.Comm_rank(accel.comm)) +default_memory_space(accel::MPIAcceleration, T::Type) = MPIMemorySpace(CPURAMMemorySpace(myid()), accel.comm, MPI.Comm_rank(accel.comm)) + +function memory_spaces(proc::MPIClusterProc) + rawMemSpace = Set{MemorySpace}() + for rnk in 0:(MPI.Comm_size(proc.comm) - 1) + for innerSpace in memory_spaces(OSProc()) + push!(rawMemSpace, MPIMemorySpace(innerSpace, proc.comm, rnk)) + end + end + return rawMemSpace +end + +function memory_spaces(proc::MPIProcessor) + rawMemSpace = Set{MemorySpace}() + for innerSpace in memory_spaces(proc.innerProc) + push!(rawMemSpace, MPIMemorySpace(innerSpace, proc.comm, proc.rank)) + end + return rawMemSpace +end + +root_worker_id(mem_space::MPIMemorySpace) = myid() + +function processors(memSpace::MPIMemorySpace) + rawProc = Processor[] + for innerProc in processors(memSpace.innerSpace) + push!(rawProc, MPIProcessor(innerProc, memSpace.comm, memSpace.rank)) + end + # Return a deterministically-ordered vector rather than a `Set`: callers pick + # the owner processor with `first(processors(space))`, and that choice must be + # identical on every rank for SPMD uniformity. `Set` iterates in hash order, + # which is not guaranteed to agree across separately-launched ranks (an + # `MPIProcessor` embeds an identity-hashed `MPI.Comm`, so its hash — and thus + # the set's iteration order — can differ per process). Sort by the same + # rank-uniform identity (`short_name` of the inner processor) used by + # `check_uniform`. + sort!(rawProc; by=p->short_name(p.innerProc)) + return rawProc +end + +struct MPIRefID + tid::UInt32 + generic::Bool + id::UInt32 + function MPIRefID(tid, generic, id) + @assert tid > 0 || generic "Invalid MPIRefID: tid=$tid, generic=$generic, id=$id" + return new(tid, generic, id) + end +end +Base.hash(id::MPIRefID, h::UInt=UInt(0)) = + hash(id.tid, hash(id.generic, hash(id.id, hash(MPIRefID, h)))) + +function check_uniform(ref::MPIRefID, original=ref) + return check_uniform(ref.tid, original) && + check_uniform(ref.generic, original) && + check_uniform(ref.id, original) +end + +function to_tag() + if Dagger.in_task() + # Tag is already assigned + opts = Dagger.get_tls().task_spec.options + tag = opts.tag + if tag !== nothing + return tag + end + # Tasks spawned outside the datadeps queue may not have a tag option; + # derive a uniform tag from the thunk ID instead + return to_tag(take_ref_id!()) + end + + # Generate a tag based on the TID + @assert !Dagger.Sch.SCHED_MOVE[] "We should not create a tag during Sch move" + return to_tag(take_ref_id!()) +end +to_tag(id::MPIRefID) = id.generic ? id.id : id.tid + +# DATADEPS_THUNK_ID (current task uid for deterministic MPIRefID generation) is owned by +# core Dagger (src/datadeps/queue.jl) and imported below; core datadeps sets it +# and this extension reads it. +# Private internal value for tracking TID-based ID generations +# N.B. Sub-counters are keyed by TID: every rank allocates IDs for a given TID +# in the same deterministic order (sequential SPMD planning, or the fixed +# tochunk sequence within a thunk), so the resulting MPIRefIDs are uniform. +const _MPIREF_TID = LockedObject(Dict{Int,Int}()) + +mutable struct MPIRef + comm::MPI.Comm + rank::Int + size::Int + innerRef::Union{DRef, Nothing} + id::MPIRefID +end +Base.hash(ref::MPIRef, h::UInt=UInt(0)) = hash(ref.rank, hash(ref.id, hash(MPIRef, h))) +root_worker_id(ref::MPIRef) = myid() + +function check_uniform(ref::MPIRef, original=ref) + return check_uniform(ref.rank, original) && + check_uniform(ref.id, original) +end + +function unwrap(handle::MPIRef) + @assert handle.rank == MPI.Comm_rank(handle.comm) "MPIRef $handle is not owned by this rank: $(handle.rank) != $(MPI.Comm_rank(handle.comm))" + return unwrap(handle.innerRef) +end + +# An MPIRef's data is only inspectable on the rank that owns it; other ranks +# must not attempt to `unwrap`/`aliasing` it during (rank-uniform) planning. +Dagger.aliasing_available(x::Chunk{<:Any,<:MPIRef}) = + x.handle.rank == MPI.Comm_rank(x.handle.comm) + +to_tag(ref::MPIRef) = to_tag(ref.id) + +move(from_proc::Processor, to_proc::Processor, x::MPIRef) = + move(from_proc, to_proc, poolget(x; uniform=uniform_execution())) + +# Read the size carried on the MPIRef itself: `innerRef` only exists on the +# owning rank, so querying it there would report 0 everywhere else. Non-owning +# ranks hold a placeholder (size 0) since they never see the value at +# construction; this residual per-rank difference is harmless because `datasize` +# only feeds cost-based placement, which uniform (MPI) execution overrides via +# `select_processors_uniform!` (see estimate_task_costs! / schedule!). +datasize(x::MPIRef) = x.size + +next_ref_sub_id!(tid) = + lock(_MPIREF_TID) do counters + counters[tid] = get(counters, tid, 0) + 1 + end + +# Forget completed tasks' ref-ID sub-counters: a finished task generates no +# further IDs, and without cleanup the table grows with every task for the +# whole session. Reclaim at wait_all (datadeps region end) after execution — +# planning (DATADEPS_THUNK_ID=>uid) and in-task tochunk share the same TID key (eager +# uid == thunk id), so reclaiming between planning and execution would collide +# on sub-ids. Rank-local and timing-independent: only handed-out values matter +# for uniformity, and TIDs never recur (monotonic shared counter). Deferral +# does not exhaust MPI tags (tags track monotonic TIDs / hashed (tid,sub)). +mpi_cleanup_tid(tid::Integer) = + lock(_MPIREF_TID) do counters + delete!(counters, Int(tid)) + end + +function mpi_cleanup_tids!(tids) + lock(_MPIREF_TID) do counters + for tid in tids + delete!(counters, Int(tid)) + end + end +end + +cleanup_tasks_accel!(::MPIAcceleration, tasks) = + mpi_cleanup_tids!(Iterators.map(t -> Int(t.uid), tasks)) + +function schedule_argument_move(::MPIAcceleration, thunk_id::Integer, f::Function) + with(DATADEPS_THUNK_ID => Int64(thunk_id)) do + f() + end + return nothing +end + +function take_ref_id!() + tid = 0 + generic = 0 + id = 0 + if Dagger.in_task() + tid = sch_handle().thunk_id.id + id = next_ref_sub_id!(tid) + elseif DATADEPS_THUNK_ID[] != 0 + tid = DATADEPS_THUNK_ID[] + id = next_ref_sub_id!(tid) + else + if current_task() !== Base.roottask + throw(ConcurrencyViolationError("Attempted to generate generic MPIRefID in a multi-threaded context")) + end + generic = true + id = next_id() # Abuse the TID counter for generic IDs + check_uniform(id) + end + # Only the tag-relevant component must fit in an MPI tag + # N.B. Skip when MPI is uninitialized (e.g. pure-Distributed callers) + @assert !MPI.Initialized() || (generic == true ? id : tid) < MPI.tag_ub() + return MPIRefID(tid, generic, id) +end + +#TODO: partitioned scheduling with comm bifurcation +function tochunk_pset(x, space::MPIMemorySpace; device=nothing, force_nonlocal=false, kwargs...) + @assert space.comm == MPI.COMM_WORLD "$(space.comm) != $(MPI.COMM_WORLD)" + local_rank = MPI.Comm_rank(space.comm) + Mid = take_ref_id!() + if local_rank != space.rank || force_nonlocal + return MPIRef(space.comm, space.rank, 0, nothing, Mid) + else + # type= is for Chunk metadata only; MemPool.poolset does not accept it + pset_kw = (; (k => v for (k, v) in pairs(kwargs) if k !== :type)...) + return MPIRef(space.comm, space.rank, sizeof(x), poolset(x; device, pset_kw...), Mid) + end +end + +const DEADLOCK_DETECT = TaskLocalValue{Bool}(()->true) +const DEADLOCK_WARN_PERIOD = TaskLocalValue{Float64}(()->10.0) +const DEADLOCK_TIMEOUT_PERIOD = TaskLocalValue{Float64}(()->120.0) +const RECV_WAITING = LockedObject(Dict{Tuple{MPI.Comm, Int, Int}, Base.Event}()) + +# Envelope for the out-of-place raw-bytes MPI path: serialize a small +# descriptor, then send contiguous bitstype buffers. Types register via +# inplace_mpi_parts / inplace_mpi_alloc / inplace_mpi_build (same style as +# move_rewrap_parts / move_rewrap_build). +struct InplaceInfo + type::DataType + header # Any MPI-serializable reconstruction metadata +end + +# Extensibility defaults: unsupported → full Julia serialize +inplace_mpi_parts(x) = nothing +# -> (parts::Tuple, header) where each part is a DenseArray with bitstype eltype +inplace_mpi_alloc(::Type{T}, header) where T = + error("inplace_mpi_alloc not defined for $T") +inplace_mpi_build(::Type{T}, parts, header) where T = + error("inplace_mpi_build not defined for $T") + +# DenseArray: single contiguous buffer + shape header +function inplace_mpi_parts(A::DenseArray) + isbitstype(eltype(A)) || return nothing + return ((A,), size(A)) +end +function inplace_mpi_alloc(::Type{T}, shape::Tuple) where {T<:DenseArray} + # Always materialize on the host: the sender host-stages device arrays, + # and callers convert to the destination space's native storage + return (Array{eltype(T)}(undef, shape),) +end +inplace_mpi_build(::Type{<:DenseArray}, (A,), ::Tuple) = A + +# SparseMatrixCSC: three vectors + (m, n, lengths) header; preserve Ti +function inplace_mpi_parts(S::SparseMatrixCSC) + isbitstype(eltype(S)) || return nothing + return ((S.colptr, S.rowval, S.nzval), + (S.m, S.n, length(S.colptr), length(S.rowval), length(S.nzval))) +end +function inplace_mpi_alloc(::Type{T}, header::Tuple) where {T<:SparseMatrixCSC} + Tv, Ti = eltype(T), T.parameters[2] + _, _, ncolptr, nrowval, nnzval = header + return (Vector{Ti}(undef, ncolptr), + Vector{Ti}(undef, nrowval), + Vector{Tv}(undef, nnzval)) +end +function inplace_mpi_build(::Type{T}, (colptr, rowval, nzval), header::Tuple) where {T<:SparseMatrixCSC} + m, n, _, _, _ = header + Tv, Ti = eltype(T), T.parameters[2] + return SparseMatrixCSC{Tv,Ti}(m, n, colptr, rowval, nzval) +end + +function supports_inplace_mpi(value) + if value isa DenseArray && isbitstype(eltype(value)) + return true + else + return false + end +end + +# Whether the loaded MPI library is known to support device buffers of `kind` +# (GPU extensions cannot depend on MPI, so they query through this). +# `mpi_device_direct`/`mpi_device_sync`/`mpi_remap_space` defaults live in core +# Dagger (src/gpu.jl); this method refines `mpi_library_gpu_aware`. +function mpi_library_gpu_aware(kind::Symbol) + MPI.Initialized() || return false + kind === :CUDA && return MPI.has_cuda() + kind === :ROC && return MPI.has_rocm() + # MPI.jl has no has_ze/has_level_zero; Intel GPU-aware MPI is opt-in via + # Intel MPI + I_MPI_OFFLOAD, or JULIA_MPI_HAS_ONEAPI / DAGGER_MPI_GPU_DIRECT + if kind === :oneAPI + env = get(ENV, "JULIA_MPI_HAS_ONEAPI", "") + !isempty(env) && return something(tryparse(Bool, env), false) + offload = get(ENV, "I_MPI_OFFLOAD", "") + isempty(offload) && return false + something(tryparse(Int, offload), 0) <= 0 && return false + impl = try + MPI.MPI_LIBRARY + catch + "" + end + return occursin(r"(?i)intel|impi", String(impl)) + end + return false +end + +### Same-node device IPC fast path (MPI handle transport) +# Eligibility and export/import live in gpu.jl / GPU extensions. Here we only +# exchange the small handle over MPI P2P and wait for an ack so the sender can +# release its staging allocation. Host staging for the non-IPC path uses +# STAGE_POOLS / stage_* from gpu.jl. + +# The sender keeps its staging alive until the receiver acknowledges the +# device-to-device copy; the ack wait polls via Improbe+yield (recv_yield) +function mpi_ipc_send(value, comm, dest, tag) + info, token = ipc_export(value) + try + send_yield(info, comm, dest, tag) + ack = recv_yield(comm, dest, tag) + @assert ack === true "Unexpected IPC ack: $(repr(ack))" + finally + ipc_release!(token) + end + return +end +function mpi_ipc_recv!(dest, comm, src, tag) + info = recv_yield(comm, src, tag) + ipc_copyto!(dest, info) + send_yield(true, comm, src, tag) + return dest +end +function mpi_ipc_recv_materialize(comm, src, tag) + info = recv_yield(comm, src, tag) + value = ipc_materialize(info) + send_yield(true, comm, src, tag) + return value +end +function recv_yield!(buffer, comm, src, tag) + rank = MPI.Comm_rank(comm) + #Core.println("buffer recv: $buffer, type of buffer: $(typeof(buffer)), is in place? $(supports_inplace_mpi(buffer))") + if !supports_inplace_mpi(buffer) + return recv_yield(comm, src, tag), false + end + #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Starting recv! from [$src]") + + # Ensure no other receiver is waiting + our_event = Base.Event() + @label retry + other_event = lock(RECV_WAITING) do waiting + if haskey(waiting, (comm, src, tag)) + waiting[(comm, src, tag)] + else + waiting[(comm, src, tag)] = our_event + nothing + end + end + if other_event !== nothing + #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Waiting for other receiver...") + wait(other_event) + @goto retry + end + + buffer = recv_yield_inplace!(buffer, comm, rank, src, tag) + + lock(RECV_WAITING) do waiting + delete!(waiting, (comm, src, tag)) + notify(our_event) + end + + return buffer, true + +end + +function recv_yield(comm, src, tag) + rank = MPI.Comm_rank(comm) + #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Starting recv from [$src]") + + # Ensure no other receiver is waiting + our_event = Base.Event() + @label retry + other_event = lock(RECV_WAITING) do waiting + if haskey(waiting, (comm, src, tag)) + waiting[(comm, src, tag)] + else + waiting[(comm, src, tag)] = our_event + nothing + end + end + if other_event !== nothing + #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Waiting for other receiver...") + wait(other_event) + @goto retry + end + #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Receiving...") + + type = nothing + @label receive + value = recv_yield_serialized(comm, rank, src, tag) + if value isa InplaceInfo + value = recv_yield_inplace(value, comm, rank, src, tag) + end + + lock(RECV_WAITING) do waiting + delete!(waiting, (comm, src, tag)) + notify(our_event) + end + return value +end + +# Device-resident dense array: receive directly when the MPI library is +# GPU-aware, otherwise through a host staging buffer +function recv_yield_inplace!(array::DenseArray, comm, my_rank, their_rank, tag) + if mpi_device_direct(array) + mpi_device_sync(array) + return _recv_yield_inplace_raw!(array, comm, my_rank, their_rank, tag) + end + kind = gpu_memory_kind(array) + buf = stage_acquire!(kind, sizeof(array)) + GC.@preserve buf begin + host = unsafe_wrap(Array, Ptr{eltype(array)}(pointer(buf)), size(array)) + _recv_yield_inplace_raw!(host, comm, my_rank, their_rank, tag) + with_context!(memory_space(array)) + copyto!(array, host) + end + # The upload runs on this task's stream; consumers run on other + # tasks/streams, so it must be visible before we hand the buffer over + mpi_device_sync(array) + stage_release!(kind, buf) + return array +end +recv_yield_inplace!(array::Array, comm, my_rank, their_rank, tag) = + _recv_yield_inplace_raw!(array, comm, my_rank, their_rank, tag) +function _recv_yield_inplace_raw!(array, comm, my_rank, their_rank, tag) + time_start = time_ns() + detect = DEADLOCK_DETECT[] + warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) + timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + + while true + (got, msg, stat) = MPI.Improbe(their_rank, tag, comm, MPI.Status) + if got + if MPI.Get_error(stat) != MPI.SUCCESS + error("recv_yield failed with error $(MPI.Get_error(stat))") + end + count = MPI.Get_count(stat, UInt8) + @assert count == sizeof(array) "recv_yield_inplace: expected $(sizeof(array)) bytes, got $count" + buf = MPI.Buffer(array) + req = MPI.Imrecv!(buf, msg) + __wait_for_request(req, comm, my_rank, their_rank, tag, "recv_yield", "recv") + return array + end + warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, my_rank, tag, "recv", their_rank) + yield() + end +end + +function recv_yield_inplace(_value::InplaceInfo, comm, my_rank, their_rank, tag) + T = _value.type + parts = inplace_mpi_alloc(T, _value.header) + parts = map(p -> recv_yield_inplace!(p, comm, my_rank, their_rank, tag), parts) + return inplace_mpi_build(T, parts, _value.header) +end + +function recv_yield_serialized(comm, my_rank, their_rank, tag) + time_start = time_ns() + detect = DEADLOCK_DETECT[] + warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) + timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + + while true + (got, msg, stat) = MPI.Improbe(their_rank, tag, comm, MPI.Status) + if got + if MPI.Get_error(stat) != MPI.SUCCESS + error("recv_yield failed with error $(MPI.Get_error(stat))") + end + count = MPI.Get_count(stat, UInt8) + buf = Array{UInt8}(undef, count) + req = MPI.Imrecv!(MPI.Buffer(buf), msg) + __wait_for_request(req, comm, my_rank, their_rank, tag, "recv_yield", "recv") + return MPI.deserialize(buf) + end + warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, my_rank, tag, "recv", their_rank) + yield() + end +end + +send_yield!(value, comm, dest, tag) = + _send_yield(value, comm, dest, tag; inplace=true) +send_yield(value, comm, dest, tag) = + _send_yield(value, comm, dest, tag; inplace=false) +function _send_yield(value, comm, dest, tag; inplace::Bool) + rank = MPI.Comm_rank(comm) + + #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Starting send to [$dest]: $(typeof(value)), is support inplace? $(supports_inplace_mpi(value))") + if inplace && supports_inplace_mpi(value) + if value isa Array + send_yield_inplace(value, comm, rank, dest, tag) + elseif mpi_device_direct(value) + # GPU-aware MPI: hand the device buffer to MPI directly + mpi_device_sync(value) + send_yield_inplace(value, comm, rank, dest, tag) + else + # Device-resident dense array: stage through a pooled (pinned) + # host buffer so the raw-bytes protocol works without GPU-aware + # MPI. Producers may have enqueued work on other streams, so sync + # before reading. + mpi_device_sync(value) + host, buf, kind = stage_to_host!(value) + GC.@preserve buf begin + send_yield_inplace(host, comm, rank, dest, tag) + end + stage_release!(kind, buf) + end + else + send_yield_serialized(value, comm, rank, dest, tag) + end +end + +function send_yield_inplace(value, comm, my_rank, their_rank, tag) + @opcounter :send_yield_inplace + req = MPI.Isend(value, comm; dest=their_rank, tag) + __wait_for_request(req, comm, my_rank, their_rank, tag, "send_yield", "send") +end + +# Send one contiguous part, staging device arrays through host when needed +function send_yield_inplace_part(part, comm, my_rank, their_rank, tag) + if part isa Array + send_yield_inplace(part, comm, my_rank, their_rank, tag) + elseif part isa DenseArray + # Device-resident dense array: send bytes directly under GPU-aware + # MPI, or host-staged otherwise (receiver materializes on the host) + mpi_device_sync(part) + if mpi_device_direct(part) + send_yield_inplace(part, comm, my_rank, their_rank, tag) + else + host, buf, kind = stage_to_host!(part) + GC.@preserve buf begin + send_yield_inplace(host, comm, my_rank, their_rank, tag) + end + stage_release!(kind, buf) + end + else + send_yield_inplace(part, comm, my_rank, their_rank, tag) + end +end + +function send_yield_serialized(value, comm, my_rank, their_rank, tag) + @opcounter :send_yield_serialized + parts_hdr = inplace_mpi_parts(value) + if parts_hdr !== nothing + parts, header = parts_hdr + send_yield_serialized(InplaceInfo(typeof(value), header), comm, my_rank, their_rank, tag) + for part in parts + send_yield_inplace_part(part, comm, my_rank, their_rank, tag) + end + else + req = MPI.isend(value, comm; dest=their_rank, tag) + __wait_for_request(req, comm, my_rank, their_rank, tag, "send_yield", "send") + end +end + +function __wait_for_request(req, comm, my_rank, their_rank, tag, fn::String, kind::String) + time_start = time_ns() + detect = DEADLOCK_DETECT[] + warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) + timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + while true + finish, status = MPI.Test(req, MPI.Status) + if finish + if MPI.Get_error(status) != MPI.SUCCESS + error("$fn failed with error $(MPI.Get_error(status))") + end + return + end + warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, my_rank, tag, kind, their_rank) + yield() + end +end + +function bcast_send_yield(value, comm, root, tag) + @opcounter :bcast_send_yield + sz = MPI.Comm_size(comm) + rank = MPI.Comm_rank(comm) + for other_rank in 0:(sz-1) + rank == other_rank && continue + send_yield(value, comm, other_rank, tag) + end +end + +# Binomial-tree broadcast built from the tagged nonblocking P2P primitives +# (never an MPI collective). Every rank calls this at the same logical point: +# the root passes the value, the others receive from their tree parent, +# forward to their tree children, and return the received value. The tree is +# a pure function of (rank, root, comm size), hence rank-uniform; O(log P) +# latency instead of the root serially sending to P-1 peers. +function bcast_yield(comm, root::Integer, tag, value=nothing) + @opcounter :bcast_yield + sz = MPI.Comm_size(comm) + sz == 1 && return value + rank = MPI.Comm_rank(comm) + vrank = mod(rank - root, sz) + mask = 1 + while mask < sz + if vrank & mask != 0 + value = recv_yield(comm, mod(rank - mask, sz), tag) + break + end + mask <<= 1 + end + mask >>= 1 + while mask > 0 + if vrank + mask < sz + send_yield(value, comm, mod(rank + mask, sz), tag) + end + mask >>= 1 + end + return value +end + +function mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, rank, tag, kind, srcdest) + time_elapsed = (time_ns() - time_start) + if detect && time_elapsed > warn_period + @warn "[rank $rank][tag $tag] Hit probable hang on $kind (dest: $srcdest)" + return typemax(UInt64) + end + if detect && time_elapsed > timeout_period + error("[rank $rank][tag $tag] Hit hang on $kind (dest: $srcdest)") + end + return warn_period +end + +# --------------------------------------------------------------------------- +# Dispatch-decoupled binomial-tree broadcast for per-task metadata +# +# `bcast_yield` (above) is a correct O(log P) tree, but when it runs INSIDE a +# task's syncdep-gated `execute!` coroutine an intermediate rank cannot forward +# a broadcast until the scheduler dispatches that same task on it. With many +# per-task metadata broadcasts overlapping (4 ranks × many chunks) that +# "forward only once I'm dispatched" gate closes a cross-rank wait cycle — the +# hang seen on Julia 1.10/1.11. +# +# A non-owner needs none of the task's data to relay/consume the tiny status/ +# type payload, so drive the tree from an always-running per-rank relay that +# forwards ON MESSAGE ARRIVAL (independent of task dispatch) over a dedicated +# duplicated `bcast_comm`, delivering payloads to per-tag slots. The root just +# fires to its children and returns; non-roots yield-wait their slot. The wait +# graph then reduces to the acyclic datadeps DAG (receiver → owner compute → +# owner syncdeps), so it cannot deadlock, while keeping the O(log P) tree. +# Only the fixed-size `execute!` metadata uses this path; large value +# broadcasts (poolget/collect/fetch) keep the in-task tree at their sequential, +# non-overlapping points. + +# Binomial-tree children of `rank` in the tree rooted at `root` over `sz` ranks +# (identical virtual-rank math to `bcast_yield`, so relay and root agree on the +# topology). +function bcast_tree_children(rank::Int, root::Int, sz::Int) + vrank = mod(rank - root, sz) + mask = 1 + while mask < sz + (vrank & mask != 0) && break + mask <<= 1 + end + mask >>= 1 + children = Int[] + while mask > 0 + if vrank + mask < sz + push!(children, mod(rank + mask, sz)) + end + mask >>= 1 + end + return children +end + +# One-shot promise for a delivered broadcast payload (one per tag). +mutable struct BcastSlot + ev::Base.Event + value::Any + done::Bool + BcastSlot() = new(Base.Event(), nothing, false) +end + +mutable struct BcastState + bcast_comm::MPI.Comm + slots::Dict{UInt32,BcastSlot} + lock::Threads.ReentrantLock + running::Threads.Atomic{Bool} + relay::Union{Task,Nothing} +end + +# Keyed by the acceleration comm (what `execute!`/`bcast_meta_yield` see). The +# relay owns the duplicated `bcast_comm`, so it can `Improbe(ANY_SOURCE, +# ANY_TAG)` without colliding with main-comm P2P transfers or the in-task +# `bcast_yield` tree. +const BCAST_STATES = LockedObject(Dict{MPI.Comm,BcastState}()) + +bcast_state_for(comm::MPI.Comm) = + lock(BCAST_STATES) do states + states[comm] + end + +bcast_serialize(x) = (io = IOBuffer(); Serialization.serialize(io, x); take!(io)) + +function bcast_deliver!(state::BcastState, tag::UInt32, value) + slot = lock(state.lock) do + get!(BcastSlot, state.slots, tag) + end + slot.value = value + slot.done = true + notify(slot.ev) + return +end + +function bcast_slot_wait(state::BcastState, tag::UInt32) + slot = lock(state.lock) do + get!(BcastSlot, state.slots, tag) + end + # `Base.Event` latches: a `notify` that already fired makes `wait` return + # immediately, so relay-before-consumer and consumer-before-relay both work. + wait(slot.ev) + value = slot.value + lock(state.lock) do + delete!(state.slots, tag) + end + return value +end + +# Always-running per-rank progress engine: receive a broadcast on `bcast_comm`, +# forward the raw bytes to this rank's tree children, then deliver the payload. +# Forwarding is NOT gated by datadeps task dispatch — that is what breaks the +# forwarder wait-cycle. +function bcast_relay_loop(state::BcastState) + comm = state.bcast_comm + rank = MPI.Comm_rank(comm) + sz = MPI.Comm_size(comm) + while state.running[] + try + MPI.Finalized() && break + got, msg, stat = MPI.Improbe(MPI.ANY_SOURCE, MPI.ANY_TAG, comm, MPI.Status) + if !got + yield() + continue + end + src = MPI.Get_source(stat) + tag = UInt32(MPI.Get_tag(stat)) + count = MPI.Get_count(stat, UInt8) + buf = Array{UInt8}(undef, count) + req = MPI.Imrecv!(MPI.Buffer(buf), msg) + __wait_for_request(req, comm, rank, src, tag, "bcast_relay", "recv") + root, value = Serialization.deserialize(IOBuffer(buf)) + for child in bcast_tree_children(rank, Int(root), sz) + sreq = MPI.Isend(buf, comm; dest=child, tag=tag) + __wait_for_request(sreq, comm, rank, child, tag, "bcast_relay", "send") + end + bcast_deliver!(state, tag, value) + catch err + # Comm torn down (disable/finalize) races the loop: stop quietly. + # A genuine mid-session fault surfaces via errormonitor. + (state.running[] && !MPI.Finalized()) || break + rethrow(err) + end + end + return +end + +function start_bcast_relay!(accel_comm::MPI.Comm) + bcast_comm = MPI.Comm_dup(accel_comm) + state = BcastState(bcast_comm, Dict{UInt32,BcastSlot}(), + Threads.ReentrantLock(), Threads.Atomic{Bool}(true), nothing) + lock(BCAST_STATES) do states + states[accel_comm] = state + end + t = Task(() -> bcast_relay_loop(state)) + t.sticky = false + schedule(t) + Dagger.Sch.errormonitor_tracked("mpi_bcast_relay", t) + state.relay = t + # `start_bcast_relay!` runs strictly after `MPI.Init`, so this hook is + # registered after MPI.jl's Finalize hook and (atexit is LIFO) runs before + # it — the relay is stopped and joined before MPI shuts down, otherwise a + # relay mid-`Improbe` on another thread races Finalize and MPICH aborts. + atexit(() -> stop_bcast_relay!(accel_comm)) + return state +end + +function stop_bcast_relay!(accel_comm::MPI.Comm) + state = lock(BCAST_STATES) do states + s = get(states, accel_comm, nothing) + s === nothing || delete!(states, accel_comm) + s + end + state === nothing && return nothing + state.running[] = false + # Join so no in-flight probe/recv can outlive the comm; a relay that died + # to a teardown race already reported via errormonitor, so swallow it here. + relay = state.relay + if relay !== nothing && relay !== current_task() + try + wait(relay) + catch + end + end + return nothing +end + +# Dispatch-decoupled broadcast used by `execute!` for per-task metadata. The +# root fires `(root, value)` to its tree children on `bcast_comm` and returns; +# every other rank waits for the relay to deliver its copy. +function bcast_meta_yield(comm::MPI.Comm, root::Integer, tag, value=nothing) + sz = MPI.Comm_size(comm) + sz == 1 && return value + rank = MPI.Comm_rank(comm) + state = bcast_state_for(comm) + utag = UInt32(tag) + if rank == root + buf = bcast_serialize((Int(root), value)) + bcomm = state.bcast_comm + for child in bcast_tree_children(rank, Int(root), sz) + req = MPI.Isend(buf, bcomm; dest=child, tag=utag) + __wait_for_request(req, bcomm, rank, child, utag, "bcast_meta", "send") + end + return value + else + return bcast_slot_wait(state, utag) + end +end + +WeakChunk(c::Chunk{T,H}) where {T,H<:MPIRef} = WeakChunk(c.handle.rank, c.handle.id.id, WeakRef(c), T) + +function MemPool.poolget(ref::MPIRef; uniform::Bool=uniform_execution()) + @assert uniform || ref.rank == MPI.Comm_rank(ref.comm) "MPIRef rank mismatch: $(ref.rank) != $(MPI.Comm_rank(ref.comm))" + if uniform + tag = to_tag() + if ref.rank == MPI.Comm_rank(ref.comm) + return bcast_yield(ref.comm, ref.rank, tag, poolget(ref.innerRef)) + else + return bcast_yield(ref.comm, ref.rank, tag) + end + else + return poolget(ref.innerRef) + end +end +fetch_handle(ref::MPIRef; uniform::Bool=uniform_execution()) = poolget(ref; uniform) + +function move!(dep_mod, to_space::MPIMemorySpace, from_space::MPIMemorySpace, to::Chunk, from::Chunk) + @assert to.handle isa MPIRef && from.handle isa MPIRef "MPIRef expected" + @assert to.handle.comm == from.handle.comm "MPIRef comm mismatch" + @assert to.handle.rank == to_space.rank && from.handle.rank == from_space.rank "MPIRef rank mismatch" + local_rank = MPI.Comm_rank(from.handle.comm) + if to_space.rank == from_space.rank == local_rank + move!(dep_mod, to_space.innerSpace, from_space.innerSpace, to, from) + else + # Prefer the copy task's own (uniform, unique) tag; fall back to the + # source handle's tag outside of task context + tag = Dagger.in_task() ? to_tag() : to_tag(from.handle) + @dagdebug nothing :mpi "[$local_rank][$tag] Moving from $(from_space.rank) to $(to_space.rank)\n" + if dep_mod === identity + # Below the crossover, handle+ack overhead makes host staging + # competitive; each side checks its own (equal-sized) buffer + local_nbytes = local_rank == from_space.rank ? from.handle.size : + local_rank == to_space.rank ? to.handle.size : 0 + if ipc_eligible(from_space.innerSpace, to_space.innerSpace) && + local_nbytes >= IPC_MIN_BYTES[] && + same_node(current_acceleration(), from_space.rank, to_space.rank) + # Same-node device-to-device: ship only an IPC handle + if local_rank == from_space.rank + mpi_ipc_send(poolget(from.handle; uniform=false), to_space.comm, to_space.rank, tag) + elseif local_rank == to_space.rank + mpi_ipc_recv!(poolget(to.handle; uniform=false), from_space.comm, from_space.rank, tag) + end + # Full copy: receive directly into the destination buffer when possible + elseif local_rank == from_space.rank + send_yield!(poolget(from.handle; uniform=false), to_space.comm, to_space.rank, tag) + elseif local_rank == to_space.rank + to_val = poolget(to.handle; uniform=false) + val, inplace = recv_yield!(to_val, from_space.comm, from_space.rank, tag) + if !inplace + # Dense payloads may arrive host-staged; materialize in + # dest storage so inner move! sees matching spaces + val = mpi_localize(to_space.innerSpace, val) + move!(dep_mod, to_space.innerSpace, from_space.innerSpace, to_val, val) + end + end + else + # Partial copy (e.g. UpperTriangular): the destination must only + # update the dep_mod-selected region. Both sides compute the + # modifier's memory spans from their own replicas (counts and + # lengths are rank-uniform; addresses are local) and only the + # selected region travels, packed through the staging pool. + # Device arrays gather/scatter via multi_span_* (one KA launch). + pack = true + if local_rank == from_space.rank + from_val = poolget(from.handle; uniform=false) + if pack + ainfo = _with_default_acceleration() do + aliasing(from_val, dep_mod) + end + # Both sides compute this from identically-typed replicas, + # so the protocol choice stays rank-uniform + pack = !(ainfo isa NoAliasing || ainfo isa UnknownAliasing) + end + if pack + spans = memory_spans(ainfo) + len = sum(span_len, spans; init=UInt64(0)) + kind = gpu_memory_kind(from_space.innerSpace) + stage_buf = stage_acquire!(kind, len) + copies = unsafe_wrap(Array, pointer(stage_buf), Int(len)) + with_context!(from_space.innerSpace) + mpi_device_sync(from_space.innerSpace) + GC.@preserve stage_buf begin + multi_span_gather!(copies, from_val, spans) + send_yield!(copies, to_space.comm, to_space.rank, tag) + end + stage_release!(kind, stage_buf) + else + send_yield(from_val, to_space.comm, to_space.rank, tag) + end + elseif local_rank == to_space.rank + to_val = poolget(to.handle; uniform=false) + if pack + ainfo = _with_default_acceleration() do + aliasing(to_val, dep_mod) + end + pack = !(ainfo isa NoAliasing || ainfo isa UnknownAliasing) + end + if pack + spans = memory_spans(ainfo) + len = sum(span_len, spans; init=UInt64(0)) + kind = gpu_memory_kind(to_space.innerSpace) + stage_buf = stage_acquire!(kind, len) + copies = unsafe_wrap(Array, pointer(stage_buf), Int(len)) + GC.@preserve stage_buf begin + recv_yield!(copies, from_space.comm, from_space.rank, tag) + with_context!(to_space.innerSpace) + multi_span_scatter!(to_val, copies, spans) + end + mpi_device_sync(to_space.innerSpace) + stage_release!(kind, stage_buf) + else + val = recv_yield(from_space.comm, from_space.rank, tag) + # Dense payloads arrive host-staged; materialize in this + # space's storage so dep_mod views copy device-to-device + val = mpi_localize(to_space.innerSpace, val) + move!(dep_mod, to_space.innerSpace, from_space.innerSpace, to_val, val) + end + end + end + end + @dagdebug nothing :mpi "[$local_rank] Finished moving from $(from_space.rank) to $(to_space.rank) successfuly\n" +end +function move!(dep_mod::Dagger.RemainderAliasing{<:MPIMemorySpace}, to_space::MPIMemorySpace, from_space::MPIMemorySpace, to::Chunk, from::Chunk) + @assert to.handle isa MPIRef && from.handle isa MPIRef "MPIRef expected" + @assert to.handle.comm == from.handle.comm "MPIRef comm mismatch" + @assert to.handle.rank == to_space.rank && from.handle.rank == from_space.rank "MPIRef rank mismatch" + local_rank = MPI.Comm_rank(from.handle.comm) + if to_space.rank == from_space.rank == local_rank + move!(dep_mod, to_space.innerSpace, from_space.innerSpace, to, from) + else + tag = Dagger.in_task() ? to_tag() : to_tag(from.handle) + @dagdebug nothing :mpi "[$local_rank][$tag] Moving from $(from_space.rank) to $(to_space.rank)\n" + if local_rank == from_space.rank + # Gather the source spans into a host buffer (device arrays pack + # via multi_span_gather! / one KA launch, then a single DtoH) + len = sum(span_tuple->span_len(span_tuple[1]), dep_mod.spans) + kind = gpu_memory_kind(from_space.innerSpace) + stage_buf = stage_acquire!(kind, len) + copies = unsafe_wrap(Array, pointer(stage_buf), Int(len)) + from_raw = poolget(from.handle; uniform=false) + with_context!(from_space.innerSpace) + # Producers may have enqueued work on other tasks' streams + mpi_device_sync(from_space.innerSpace) + GC.@preserve stage_buf begin + multi_span_gather!(copies, from_raw, first.(dep_mod.spans)) + send_yield!(copies, to_space.comm, to_space.rank, tag) + end + stage_release!(kind, stage_buf) + elseif local_rank == to_space.rank + # Receive the spans + len = sum(span_tuple->span_len(span_tuple[1]), dep_mod.spans) + kind = gpu_memory_kind(to_space.innerSpace) + stage_buf = stage_acquire!(kind, len) + copies = unsafe_wrap(Array, pointer(stage_buf), Int(len)) + GC.@preserve stage_buf begin + recv_yield!(copies, from_space.comm, from_space.rank, tag) + + # Scatter the received bytes into the destination spans + to_raw = poolget(to.handle; uniform=false) + with_context!(to_space.innerSpace) + multi_span_scatter!(to_raw, copies, last.(dep_mod.spans)) + end + # Consumers run on other tasks/streams + mpi_device_sync(to_space.innerSpace) + stage_release!(kind, stage_buf) + + # Ensure that the data is visible + Core.Intrinsics.atomic_fence(:release) + end + end + + return +end + +# Functions route through the inner processor, so GPU backends can swap +# kernels for device equivalents (e.g. LAPACK.potrf! -> CUSOLVER.potrf!) +move(::MPIOSProc, dst::MPIProcessor, x::Union{Function,Type}) = move(OSProc(), dst.innerProc, x) +move(::MPIOSProc, ::MPIProcessor, x::Chunk{<:Union{Function,Type}}) = poolget(x.handle) + +# Out-of-place MPI move for OSProc-labeled Chunks: delegate to the +# MPIProcessor Chunk path so egress runs on the owner and ingress on dest. +function move(src::MPIOSProc, dst::MPIProcessor, x::Chunk) + @assert src.comm == dst.comm "Multi comm move not supported" + from = x.processor + if !(from isa MPIProcessor) || from.rank != x.handle.rank || from.comm != src.comm + from = MPIProcessor(OSProc(), src.comm, x.handle.rank) + end + return move(from, dst, x) +end + +function is_local(accel::MPIAcceleration, target) + return target.rank == MPI.Comm_rank(accel.comm) +end + +### SPMD-symmetric datadeps endpoints +# +# Under MPI, datadeps planning (`distribute_tasks!`) runs identically on every +# rank (SPMD). All ranks therefore enter `generate_slot!` / +# `remotecall_endpoint_toplevel` collectively and walk the exact same +# structure-driven path: the aliased-object cache is a per-rank replicated +# store updated symmetrically at the same logical points (no RPC, no +# side-channel registries), aliasing metadata is broadcast from the owner at +# these aligned points, and bulk data is routed owner->destination via P2P +# inside the endpoint. + +# A value participating in an SPMD endpoint: the owner rank holds the actual +# value, all other ranks hold only the type and the origin (owner) space. +struct MPIWireValue{T} + value::Union{Some{T},Nothing} + space::MPIMemorySpace +end +wire_type(::MPIWireValue{T}) where T = T +has_value(w::MPIWireValue) = w.value !== nothing +wire_value(w::MPIWireValue) = something(w.value) + +memory_space(w::MPIWireValue) = w.space +default_memory_space(accel::MPIAcceleration, w::MPIWireValue) = w.space + +function check_uniform(w::MPIWireValue{T}, original=w) where T + # Compare logical metadata only: the value itself is rank-local + return check_uniform(hash(T), original) && + check_uniform(w.space, original) +end + +function tochunk(w::MPIWireValue{T}, proc::P, scope::S=Dagger.AnyScope(); kwargs...) where {T,P<:Processor,S} + if has_value(w) + return tochunk(wire_value(w), proc, w.space, scope; kwargs...) + else + return tochunk(nothing, proc, w.space, scope; type=T, kwargs...) + end +end + +# Remap rank-local memory space stamps (CPURAMMemorySpace(myid()), where +# myid() == 1 on every rank) in aliasing info to the owning rank, so that +# spans from different ranks never falsely alias (SPMD heaps have very +# similar address layouts). `mpi_remap_space` defaults live in core Dagger. +mpi_remap_ptr(ptr::Dagger.RemotePtr{T}, owner::Int) where T = + RemotePtr{T}(ptr.addr, mpi_remap_space(ptr.space, owner)) +mpi_remap_span(span::MemorySpan, owner::Int) = + MemorySpan(mpi_remap_ptr(span.ptr, owner), span.len) + +mpi_remap_ainfo(a::Dagger.NoAliasing, owner::Int) = a +mpi_remap_ainfo(a::Dagger.UnknownAliasing, owner::Int) = a +mpi_remap_ainfo(a::Dagger.ContiguousAliasing, owner::Int) = + Dagger.ContiguousAliasing(mpi_remap_span(a.span, owner)) +mpi_remap_ainfo(a::Dagger.ObjectAliasing, owner::Int) = + Dagger.ObjectAliasing(mpi_remap_ptr(a.ptr, owner), a.sz) +function mpi_remap_ainfo(a::Dagger.StridedAliasing{T,N,S}, owner::Int) where {T,N,S} + base_ptr = mpi_remap_ptr(a.base_ptr, owner) + return Dagger.StridedAliasing{T,N,typeof(base_ptr.space)}(base_ptr, + mpi_remap_ptr(a.ptr, owner), + a.base_inds, a.lengths, a.strides) +end +function mpi_remap_ainfo(a::Dagger.TriangularAliasing{T,S}, owner::Int) where {T,S} + ptr = mpi_remap_ptr(a.ptr, owner) + return Dagger.TriangularAliasing{T,typeof(ptr.space)}(ptr, a.stride, a.isupper, a.diagonal) +end +function mpi_remap_ainfo(a::Dagger.DiagonalAliasing{T,S}, owner::Int) where {T,S} + ptr = mpi_remap_ptr(a.ptr, owner) + return Dagger.DiagonalAliasing{T,typeof(ptr.space)}(ptr, a.stride) +end +mpi_remap_ainfo(a::Dagger.CombinedAliasing, owner::Int) = + Dagger.CombinedAliasing(Dagger.AbstractAliasing[mpi_remap_ainfo(sub, owner) for sub in a.sub_ainfos]) +mpi_remap_ainfo(a::Dagger.AliasingWrapper, owner::Int) = + Dagger.AliasingWrapper(mpi_remap_ainfo(a.inner, owner)) +mpi_remap_ainfo(a::Dagger.AbstractAliasing, owner::Int) = a + +# Owner computes the aliasing info for its local value and broadcasts it; +# all ranks must call this collectively at the same logical point. +function Dagger.aliasing(accel::MPIAcceleration, w::MPIWireValue, dep_mod) + tag = to_tag() + check_uniform(tag) + rank = MPI.Comm_rank(accel.comm) + if w.space.rank == rank + ainfo = Dagger._with_default_acceleration() do + Dagger.aliasing(wire_value(w), dep_mod) + end + ainfo = mpi_remap_ainfo(ainfo, w.space.rank) + @opcounter :aliasing_bcast_send_yield + ainfo = bcast_yield(accel.comm, w.space.rank, tag, ainfo) + else + ainfo = bcast_yield(accel.comm, w.space.rank, tag) + end + check_uniform(ainfo) + return ainfo +end + +# All ranks enter collectively; the owner unwraps its local value, all other +# ranks construct a wire proxy carrying only type and origin space. The +# `move_rewrap` recursion below then walks the same path on every rank. +function remotecall_endpoint_toplevel(f, accel::MPIAcceleration, cache::AliasedObjectCache, from_proc, to_proc, from_space, to_space, data::Chunk) + local_rank = MPI.Comm_rank(accel.comm) + T = chunktype(data) + w = if local_rank == from_space.rank + MPIWireValue{T}(Some{T}(unwrap(data)), from_space) + else + MPIWireValue{T}(nothing, from_space) + end + return f(accel, cache, from_proc, to_proc, from_space, to_space, w)::Chunk +end + +# Materialize a host-staged received value in the native storage of `space` +# (e.g. upload a host Array into device memory for a GPU destination) +function mpi_localize(space::MemorySpace, value) + space isa CPURAMMemorySpace && return value + value isa Array || return value + return move(OSProc(), first(processors(space)), value) +end + +# Inner processor of a possibly-MPI-wrapped processor +mpi_inner_proc(proc::MPIProcessor) = proc.innerProc +mpi_inner_proc(proc::Processor) = proc + +# Ensure a processor carries MPI rank metadata for dual-sided OOP move +function as_mpi_proc(proc::MPIProcessor, space::MPIMemorySpace) + if proc.rank == space.rank && proc.comm == space.comm + return proc + end + return MPIProcessor(proc.innerProc, space.comm, space.rank) +end +as_mpi_proc(proc::Processor, space::MPIMemorySpace) = + MPIProcessor(proc, space.comm, space.rank) + +# Source-half OOP conversion: adapt off the owner's device toward a +# transportable form. Dense device buffers that send_yield can ship +# GPU-direct (or stage itself) are left as-is after a device sync. +function mpi_move_egress(src::MPIProcessor, x) + if x isa DenseArray && isbitstype(eltype(x)) && !(x isa Array) + mpi_device_sync(x) + return x + end + return move(src.innerProc, OSProc(), x) +end + +# Dest-half OOP conversion: adapt a received (often host-staged) value +# into the destination inner processor's native storage (e.g. HtoD). +mpi_move_ingress(dst::MPIProcessor, x) = move(OSProc(), dst.innerProc, x) + +# Rank-aware out-of-place move: egress on the source rank, ingress on the +# destination rank. Distributed can stay dest-pull; MPI needs both sides. +function move(src::MPIProcessor, dst::MPIProcessor, x) + @assert src.comm == dst.comm "Multi comm move not supported" + local_rank = MPI.Comm_rank(src.comm) + if src.rank == dst.rank + if local_rank == src.rank + return move(src.innerProc, dst.innerProc, x) + end + return nothing + elseif local_rank == src.rank + return mpi_move_egress(src, x) + elseif local_rank == dst.rank + return mpi_move_ingress(dst, x) + end + return nothing +end + +# Transfer the wire value from its owner to the destination rank, returning a +# uniform Chunk (real payload on the destination, placeholder elsewhere). +# Both ranks run move: owner does egress before send, dest does ingress after +# recv. Chunk type is computed uniformly on every rank via move_type. +function mpi_endpoint_transfer(accel::MPIAcceleration, from_proc, to_proc, from_space, to_space, w::MPIWireValue{T}) where T + local_rank = MPI.Comm_rank(accel.comm) + from_mpi = as_mpi_proc(from_proc, from_space) + to_mpi = as_mpi_proc(to_proc, to_space) + from_inner = mpi_inner_proc(from_mpi) + to_inner = mpi_inner_proc(to_mpi) + T_new = move_type(from_inner, to_inner, T) + if from_space.rank == to_space.rank + # No cross-rank movement; the owner (re)wraps its local value, + # converting between inner spaces (e.g. CPU -> GPU) when they differ + if local_rank == from_space.rank + value = move(from_mpi, to_mpi, wire_value(w)) + return tochunk(value, to_proc, to_space; type=T_new) + else + return tochunk(nothing, to_proc, to_space; type=T_new) + end + end + tag = to_tag() + if T <: DenseArray && isbitstype(eltype(T)) && + ipc_eligible(from_space.innerSpace, to_space.innerSpace) && + same_node(accel, from_space.rank, to_space.rank) + # Same-node device-to-device slot creation: ship only an IPC handle + if local_rank == from_space.rank + mpi_ipc_send(wire_value(w), accel.comm, to_space.rank, tag) + return tochunk(nothing, to_proc, to_space; type=T_new) + elseif local_rank == to_space.rank + value = mpi_ipc_recv_materialize(accel.comm, from_space.rank, tag) + return tochunk(value, to_proc, to_space; type=T_new) + else + return tochunk(nothing, to_proc, to_space; type=T_new) + end + end + if local_rank == from_space.rank + # Owner: egress move, then send the (possibly host-staged) value + value = move(from_mpi, to_mpi, wire_value(w)) + send_yield(value, accel.comm, to_space.rank, tag) + return tochunk(nothing, to_proc, to_space; type=T_new) + elseif local_rank == to_space.rank + value = recv_yield(accel.comm, from_space.rank, tag) + # Dest: ingress move into this space's native storage + value = move(from_mpi, to_mpi, value) + return tochunk(value, to_proc, to_space; type=T_new) + else + return tochunk(nothing, to_proc, to_space; type=T_new) + end +end + +# Generic / wrapper MPIWireValue: leaf transfer, or header+children rebuild +function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, w::MPIWireValue{T}) where T + child_types = move_rewrap_child_types(T) + 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) + end + end + + local_rank = MPI.Comm_rank(accel.comm) + child_wires = if has_value(w) + children, _ = move_rewrap_parts(wire_value(w)) + ntuple(length(child_types)) do i + CT = child_types[i] + MPIWireValue{CT}(Some{CT}(children[i]), w.space) + end + else + ntuple(length(child_types)) do i + CT = child_types[i] + MPIWireValue{CT}(nothing, w.space) + end + end + + child_chunks = map(cw -> move_rewrap(accel, cache, from_proc, to_proc, from_space, to_space, cw), child_wires) + for cc in child_chunks + check_uniform(cc.handle) + end + + mode = move_rewrap_header_mode(T) + header = if mode === :broadcast + tag = to_tag() + if local_rank == w.space.rank + _, hdr = move_rewrap_parts(wire_value(w)) + bcast_yield(accel.comm, w.space.rank, tag, hdr) + else + bcast_yield(accel.comm, w.space.rank, tag) + end + elseif mode === :none + nothing + else + error("MPIWireValue move_rewrap does not support header mode $mode for $T") + end + + child_cts = map(chunktype, child_chunks) + T_new = move_rewrap_result_type(T, child_cts, header) + if local_rank == to_space.rank + children_local = map(unwrap, child_chunks) + v_new = move_rewrap_build(T, children_local, header) + return tochunk(v_new, to_proc, to_space; type=T_new) + else + return tochunk(nothing, to_proc, to_space; type=T_new) + end +end + +# ChunkView: replicated SPMD metadata (inner chunk record + slices exist on +# every rank), so all ranks dispatch symmetrically without a wire proxy +function remotecall_endpoint_toplevel(f, accel::MPIAcceleration, cache::AliasedObjectCache, from_proc, to_proc, from_space, to_space, data::ChunkView) + return f(accel, cache, from_proc, to_proc, from_space, to_space, data)::Chunk +end + +# Share/transfer the backing chunk, then re-create the view on the destination +# (header is replicated SPMD metadata — no broadcast) +function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, slice::ChunkView) + local_rank = MPI.Comm_rank(accel.comm) + children, header = move_rewrap_parts(slice) + p_chunk = move_rewrap(accel, cache, from_proc, to_proc, from_space, to_space, children[1]) + check_uniform(p_chunk.handle) + T_view = move_rewrap_result_type(typeof(slice), (chunktype(p_chunk),), header) + if local_rank == to_space.rank + v_new = move_rewrap_build(typeof(slice), (unwrap(p_chunk),), header) + return tochunk(v_new, to_proc, to_space; type=T_view) + else + return tochunk(nothing, to_proc, to_space; type=T_view) + end +end + +# A Chunk-wrapped ChunkView cannot cross ranks: the inner chunk record is only +# resolvable on the wrapping rank. `datadeps_arg_wrap` keeps ChunkViews raw +# under MPI, so reaching this indicates a manual `tochunk(::ChunkView)`. +function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, w::MPIWireValue{T}) where {T<:ChunkView} + throw(ArgumentError("ChunkView must not be wrapped in a Chunk under MPI; pass the ChunkView directly")) +end + +# Owner computes the view's aliasing info locally and broadcasts it +function aliasing(accel::MPIAcceleration, x::ChunkView, dep_mod) + @assert dep_mod === identity "Dependency modifiers not yet supported for ChunkView: $dep_mod" + handle = x.chunk.handle::MPIRef + tag = to_tag() + check_uniform(tag) + rank = MPI.Comm_rank(accel.comm) + if handle.rank == rank + ainfo = _with_default_acceleration() do + v = view(unwrap(x.chunk), x.slices...) + aliasing(v, dep_mod) + end + ainfo = mpi_remap_ainfo(ainfo, handle.rank) + @opcounter :aliasing_bcast_send_yield + ainfo = bcast_yield(accel.comm, handle.rank, tag, ainfo) + else + ainfo = bcast_yield(accel.comm, handle.rank, tag) + end + check_uniform(ainfo) + return ainfo +end + +# The aliased-object cache is a per-rank replicated store; every rank updates +# its own copy at the same logical point during SPMD planning, so no +# cross-rank communication is required here. +function set_stored!(accel::MPIAcceleration, cache::AliasedObjectCache, value::Chunk, ainfo::AbstractAliasing) + set_stored!(unwrap(cache.chunk)::AliasedObjectCacheStore, cache.space, value, ainfo) + return +end +function set_key_stored!(accel::MPIAcceleration, cache::AliasedObjectCache, space::MemorySpace, ainfo::AbstractAliasing, value::Chunk) + set_key_stored!(unwrap(cache.chunk)::AliasedObjectCacheStore, space, ainfo, value) + return +end + +# Chunk may be MPI-backed (MPIRef) but labeled with OSProc; treat source as the owning rank +function move(src::OSProc, dst::MPIProcessor, x::Chunk) + if x.handle isa MPIRef + return move(MPIOSProc(x.handle.comm, x.handle.rank), dst, x) + end + error("MPI move not supported") +end + +move(src::Processor, dst::MPIProcessor, x::Chunk) = error("MPI move not supported") +move(to_proc::MPIProcessor, chunk::Chunk) = + move(chunk.processor, to_proc, chunk) +move(to_proc::Processor, d::MPIRef) = + move(MPIOSProc(d.rank), to_proc, d) +move(to_proc::MPIProcessor, x) = + move(MPIOSProc(), to_proc, x) + +# Non-Chunk task arguments: dual-sided OOP move via MPIProcessor ranks +move(src::MPIOSProc, dst::MPIProcessor, x) = + move(MPIProcessor(OSProc(), src.comm, src.rank), dst, x) + +# Functions route through the inner processor, so GPU backends can swap +# kernels for device equivalents (e.g. BLAS.gemm! -> CUBLAS.gemm!) +move(src::MPIProcessor, dst::MPIProcessor, x::Union{Function,Type}) = + move(OSProc(), dst.innerProc, x) +move(::MPIProcessor, ::MPIProcessor, x::Chunk{<:Union{Function,Type}}) = poolget(x.handle) + +function move(src::MPIProcessor, dst::MPIProcessor, x::Chunk) + if Dagger.Sch.SCHED_MOVE[] + # SPMD argument move at task launch: every rank enters; owner runs + # egress move before send, destination runs ingress after recv. + # When the chunk is not resident on the destination rank, the owner + # sends it P2P on the thunk's uniform tag (argument moves run + # sequentially in argument order, so FIFO matching is safe). + local_rank = MPI.Comm_rank(dst.comm) + if x.handle.rank == dst.rank + if local_rank == dst.rank + value = poolget(x.handle; uniform=false) + return move(src, dst, value) + end + return nothing + else + @assert DATADEPS_THUNK_ID[] != 0 "Cross-rank argument transfer requires the thunk's TID in scope (chunk owner $(x.handle.rank), destination rank $(dst.rank))" + tag = UInt32(DATADEPS_THUNK_ID[]) + # Source processor for egress uses the chunk's owning rank + from = src.rank == x.handle.rank ? src : + MPIProcessor(src.innerProc, src.comm, x.handle.rank) + if local_rank == x.handle.rank + value = move(from, dst, poolget(x.handle; uniform=false)) + send_yield(value, dst.comm, dst.rank, tag) + return nothing + elseif local_rank == dst.rank + value = recv_yield(dst.comm, x.handle.rank, tag) + return move(from, dst, value) + end + return nothing + end + else + uniform = uniform_execution() + # Either we're uniform (so everyone cooperates), or we're unwrapping locally + if !uniform + @assert src.rank == MPI.Comm_rank(src.comm) "Unwrapping not permitted" + @assert src.rank == x.handle.rank == dst.rank + end + return poolget(x.handle; uniform) + end +end + +gpu_kernel_backend(proc::MPIProcessor) = gpu_kernel_backend(proc.innerProc) + +# Owner-local payload that preserves the Chunk's SPMD-uniform chunktype after +# Sch unwraps a Chunk to a device value (e.g. Matrix chunktype + CuArray value). +# Without this, promote_op/chunktype diverge across ranks (CuArray vs Matrix). +struct MPILocalArg{T} + value +end +chunktype(::MPILocalArg{T}) where T = T +mpi_unwrap_arg(a::MPILocalArg) = a.value +mpi_unwrap_arg(a) = a + +function bind_moved_argument(::MPIAcceleration, original, moved) + if moved === nothing && (original isa Chunk || original isa WeakChunk) + return original + elseif (original isa Chunk || original isa WeakChunk) && + moved !== nothing && !(moved isa Chunk) + return MPILocalArg{chunktype(original)}(moved) + end + return moved +end + +# Exceptions may capture rank-local, non-serializable state; fall back to a +# string summary when they can't go over the wire +function mpi_wire_exception(err) + try + Serialization.serialize(IOBuffer(), err) + return err + catch + return ErrorException(sprint(showerror, err)) + end +end + +# Decide how much result metadata execute! must broadcast. +# - need_type_bcast: return type/space unknown → full (:ok,T,space)/(:error,ex) +# - !need_type_bcast && nothrow: concrete + proven non-throwing → zero MPI +# - !need_type_bcast && !nothrow: concrete but may throw → status-only (:ok / :error) +# +# Device processors (CUDA/ROCm/…) adapt values (e.g. Matrix→CuArray), so +# promote_op on chunktypes is not a safe SPMD stamp; keep full broadcast there. +function mpi_execute_bcast_plan(f, args, proc::MPIProcessor) + if !(proc.innerProc isa ThreadProc) + inferred = Base.promote_op(f, map(chunktype, args)...) + return (; need_type_bcast=true, nothrow=false, inferred) + end + arg_types = map(chunktype, args) + inferred = Base.promote_op(f, arg_types...) + # `Nothing` is a concrete type and is deliberately NOT forced onto the + # broadcast path: a `nothing` return (the common in-place / mutating task) + # is fully known on every rank (all ranks stamp `Chunk{Nothing}` and + # `memory_space(nothing, proc)`), so it can take the concrete fast path and + # skip metadata traffic. `Union{}` (function always throws) and `Any` + # (inference gave up) are not concrete and still need the full broadcast. + need_type_bcast = !isconcretetype(inferred) || inferred === Union{} || + inferred === Any + if need_type_bcast + return (; need_type_bcast=true, nothrow=false, inferred) + end + effects = Base.infer_effects(f, Tuple{arg_types...}) + nothrow = Core.Compiler.is_nothrow(effects) + return (; need_type_bcast=false, nothrow, inferred) +end + +function execute!(proc::MPIProcessor, f, args...; kwargs...) + local_rank = MPI.Comm_rank(proc.comm) + islocal = local_rank == proc.rank + inplace_move = f === move! + tag = to_tag() + # Plan from SPMD-uniform chunktypes (Chunk / MPILocalArg); run on unwrapped values + raw_args = map(mpi_unwrap_arg, args) + + if inplace_move + # Copy tasks are dest-scoped, but every rank runs move!. Bind the + # local endpoint's device (source vs dest), not always dest's + # mirrored innerProc — otherwise multi-GPU ranks activate the wrong + # context before DtoH/gather. Spectators use ThreadProc so no + # unrelated GPU context is set. + to_space = raw_args[2]::MPIMemorySpace + from_space = raw_args[3]::MPIMemorySpace + local_proc = if local_rank == to_space.rank + first(processors(to_space.innerSpace)) + elseif local_rank == from_space.rank + first(processors(from_space.innerSpace)) + else + ThreadProc(myid(), 1) + end + execute!(local_proc, f, raw_args...; kwargs...) + space = memory_space(nothing, proc)::MPIMemorySpace + dest_type = chunktype(args[4]) + return tochunk(nothing, proc, space; type=dest_type) + end + + plan = mpi_execute_bcast_plan(f, args, proc) + + # Concrete + nothrow: type/space known locally and the call cannot throw, + # so non-owners need no status or metadata traffic. Stamp the inferred + # type on every rank so chunktype stays SPMD-uniform. + if !plan.need_type_bcast && plan.nothrow + if islocal + result = execute!(proc.innerProc, f, raw_args...; kwargs...) + space = memory_space(result, proc)::MPIMemorySpace + return tochunk(result, proc, space; type=plan.inferred) + else + space = memory_space(nothing, proc)::MPIMemorySpace + return tochunk(nothing, proc, space; type=plan.inferred) + end + end + + # Concrete but may throw: broadcast only status (no T/space on the wire). + # Type inference still skips metadata; poison keeps SPMD aligned on error. + if !plan.need_type_bcast + if islocal + result = try + execute!(proc.innerProc, f, raw_args...; kwargs...) + catch err + @opcounter :execute_bcast_send_yield + bcast_meta_yield(proc.comm, proc.rank, tag, (:error, mpi_wire_exception(err))) + rethrow() + end + @opcounter :execute_bcast_send_yield + bcast_meta_yield(proc.comm, proc.rank, tag, :ok) + space = memory_space(result, proc)::MPIMemorySpace + return tochunk(result, proc, space; type=plan.inferred) + else + msg = bcast_meta_yield(proc.comm, proc.rank, tag) + if msg isa Tuple && first(msg) === :error + throw(msg[2]) + end + space = memory_space(nothing, proc)::MPIMemorySpace + return tochunk(nothing, proc, space; type=plan.inferred) + end + end + + # Non-concrete return type: full status + type/space (or error) broadcast + if islocal + result = try + execute!(proc.innerProc, f, raw_args...; kwargs...) + catch err + @opcounter :execute_bcast_send_yield + bcast_meta_yield(proc.comm, proc.rank, tag, (:error, mpi_wire_exception(err))) + rethrow() + end + T = typeof(result) + space = memory_space(result, proc)::MPIMemorySpace + @opcounter :execute_bcast_send_yield + bcast_meta_yield(proc.comm, proc.rank, tag, (:ok, T, space.innerSpace)) + return tochunk(result, proc, space; type=T) + else + msg = bcast_meta_yield(proc.comm, proc.rank, tag) + if msg[1] === :error + throw(msg[2]) + end + _, T, innerSpace = msg + space = MPIMemorySpace(innerSpace, proc.comm, proc.rank) + return tochunk(nothing, proc, space; type=T) + end +end + +accelerate!(::Val{:mpi}) = accelerate!(MPIAcceleration()) + +# Comms with an active MPI acceleration. Guards `initialize_acceleration!` +# idempotency. +const MPI_INITIALIZED_COMMS = LockedObject(Set{MPI.Comm}()) + +# rank -> system_uuid (node identity), populated once at initialization via +# Allgather — same node definition as Distributed CUDAExt IPC. +const MPI_NODE_UUIDS = LockedObject(Dict{MPI.Comm, Vector{UInt128}}()) + +# Whether two ranks share a node (false when the map is unavailable) +function same_node(accel::MPIAcceleration, r1::Integer, r2::Integer) + uuids = lock(MPI_NODE_UUIDS) do node_maps + get(node_maps, accel.comm, nothing) + end + uuids === nothing && return false + return uuids[r1+1] == uuids[r2+1] +end +# Compat for benches that still have a raw comm +function mpi_same_node(comm::MPI.Comm, r1::Integer, r2::Integer) + uuids = lock(MPI_NODE_UUIDS) do node_maps + get(node_maps, comm, nothing) + end + uuids === nothing && return false + return uuids[r1+1] == uuids[r2+1] +end + +function initialize_acceleration!(a::MPIAcceleration) + # Idempotent: user `accelerate!` may be called more than once, but + # initialization involves collectives (`MPI.Comm_dup`) which must only + # ever run once per installed session. `accelerate!` skips this when + # `same_acceleration` holds; after `finalize_acceleration!` the backend + # entry is gone and re-init is allowed. Per-task execution uses + # `set_task_acceleration!` in `do_task` (TLS only, no re-init). + already = lock(MPI_INITIALIZED_COMMS) do comms + a.comm in comms + end + already && return + if !MPI.Initialized() + MPI.Init(;threadlevel=:multiple) + end + lock(MPI_INITIALIZED_COMMS) do comms + push!(comms, a.comm) + end + # Node-locality map for shared-memory/IPC fast paths. Init-time + # collectives are sanctioned only here (alongside the RPC backend setup); + # the map is cached for the whole session. Uses system_uuid (same as + # Distributed CUDAExt) rather than COMM_TYPE_SHARED. + local_uuid = system_uuid().value + node_uuids = MPI.Allgather(local_uuid, a.comm) + lock(MPI_NODE_UUIDS) do node_maps + node_maps[a.comm] = node_uuids + end + # Dedicated broadcast communicator + per-rank relay for per-task metadata + # broadcasts (see `bcast_meta_yield`). `Comm_dup` is collective and, like + # `Allgather` above, is sanctioned only at init. + start_bcast_relay!(a.comm) + ctx = Dagger.Sch.eager_context() + sz = MPI.Comm_size(a.comm) + lock(ctx) do + for i in 0:(sz-1) + push!(ctx.procs, MPIOSProc(a.comm, i)) + end + unique!(ctx.procs) + end + # SPMD planning can call `rand`/`randperm`; keep default RNGs aligned. + sync_mpi_rng!(a.comm) + check_mpi_rng_coherent!(a.comm) +end + +function finalize_acceleration!(a::MPIAcceleration) + stop_bcast_relay!(a.comm) + lock(MPI_INITIALIZED_COMMS) do comms + delete!(comms, a.comm) + end + lock(MPI_NODE_UUIDS) do node_maps + delete!(node_maps, a.comm) + end + ctx = Dagger.Sch.eager_context() + lock(ctx) do + filter!(p -> !(p isa MPIOSProc && p.comm == a.comm), ctx.procs) + end + return nothing +end + +""" + mpi_propagate_chunk_types!(tasks, accel::MPIAcceleration, expected_type) + +Ensure all ranks use the same concrete type for the given tasks by setting +each task's options.return_type to expected_type when it is concrete. +This allows chunktype(task) to return the concrete type on every rank +without an MPI allgather of actual result types. +""" +function mpi_propagate_chunk_types!(tasks, accel::MPIAcceleration, expected_type) + isconcretetype(expected_type) || return + for t in tasks + if t isa Thunk + if t.options !== nothing + t.options.return_type = expected_type + else + t.options = Options(return_type=expected_type) + end + end + end + return +end + +# Deterministic round-robin assignment of array blocks to MPI processors: +# every rank computes the same grid, so block ownership is uniform by +# construction (the scheduler's measured costs are rank-local and cannot be +# used for placement decisions under SPMD). +function default_procgrid(accel::MPIAcceleration, nblocks::NTuple{N,Int}) where N + return CyclicProcGrid(uniform_mpi_processors(accel), nblocks) +end + +function post_stage_array_chunks!(accel::MPIAcceleration, tasks, eltype::Type{T}, nd::Int) where {T} + mpi_propagate_chunk_types!(tasks, accel, Array{eltype, nd}) + return tasks +end + +accel_matches_proc(accel::MPIAcceleration, proc::MPIOSProc) = true +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 diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 2b72b8523..4677fd0ac 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -52,6 +52,16 @@ function Dagger.aliasing(x::MtlArray{T}) where T return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end +# MPI (SPMD) integration: stamp owning rank on VRAM spaces +Dagger.mpi_remap_space(space::MetalVRAMMemorySpace, owner::Int) = + MetalVRAMMemorySpace(owner, space.device_id) +Dagger.value_memory_space(x::MtlArray) = Dagger.memory_space(x) + +# Staging-pool key; Metal has no CUDA-style pin of arbitrary host buffers +Dagger.gpu_memory_kind(::MtlArray) = :Metal +Dagger.gpu_memory_kind(::MetalVRAMMemorySpace) = :Metal +Dagger.pin_buffer!(::Val{:Metal}, buf::DenseArray) = nothing + function Dagger.unsafe_free!(x::MtlArray) Metal.unsafe_free!(x) return @@ -66,24 +76,32 @@ function to_device(proc::MtlArrayDeviceProc) end function to_context(proc::MtlArrayDeviceProc) @assert Dagger.root_worker_id(proc) == myid() - return Metal.global_queue() #CONTEXTS[proc.device] + return Metal.global_queue(DEVICES[proc.device_id]) end to_context(device_id::Integer) = Metal.global_queue(DEVICES[device_id]) to_context(dev::MtlDevice) = to_context(_device_id(dev)) function with_context!(device_id::Integer) + Metal.device!(DEVICES[device_id]) end function with_context!(proc::MtlArrayDeviceProc) @assert Dagger.root_worker_id(proc) == myid() + with_context!(proc.device_id) end function with_context!(space::MetalVRAMMemorySpace) @assert Dagger.root_worker_id(space) == myid() + with_context!(space.device_id) end Dagger.with_context!(proc::MtlArrayDeviceProc) = with_context!(proc) Dagger.with_context!(space::MetalVRAMMemorySpace) = with_context!(space) function with_context(f, x) + old_dev = Metal.device() with_context!(x) - return f() + try + return f() + finally + Metal.device!(old_dev) + end end function _sync_with_context(x::Union{Dagger.Processor,Dagger.MemorySpace}) @@ -100,6 +118,17 @@ function sync_with_context(x::Union{Dagger.Processor,Dagger.MemorySpace}) end end +# No Metal-aware MPI / CUDA-style IPC; always host-stage +Dagger.mpi_device_direct(::MtlArray) = false +function Dagger.mpi_device_sync(x::MtlArray) + space = Dagger.memory_space(x) + if Dagger.root_worker_id(space) == myid() + _sync_with_context(space) + end + return +end +Dagger.mpi_device_sync(space::MetalVRAMMemorySpace) = sync_with_context(space) + # Allocations Dagger.allocate_array_func(::MtlArrayDeviceProc, ::typeof(rand)) = Metal.rand Dagger.allocate_array_func(::MtlArrayDeviceProc, ::typeof(randn)) = Metal.randn @@ -129,13 +158,20 @@ end function Dagger.move!(to_space::MetalVRAMMemorySpace, from_space::MetalVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} sync_with_context(from_space) with_context!(to_space) - copyto!(to, from) + if Dagger.needs_multi_span_copy(to_space, from_space, to, from) + Dagger.multi_span_move!(to, from) + else + copyto!(to, from) + end return end # Out-of-place HtoD function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x) with_context(to_proc) do + if x isa DenseArray && isbitstype(eltype(x)) + Dagger.pin_buffer!(:Metal, x) + end arr = adapt(MtlArray, x) Metal.synchronize() return arr @@ -147,6 +183,9 @@ function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::Chunk) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) with_context(to_proc) do + if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) + Dagger.pin_buffer!(:Metal, cpu_data) + end arr = adapt(MtlArray, cpu_data) Metal.synchronize() return arr @@ -168,7 +207,8 @@ end function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::CPUProc, x) with_context(from_proc) do Metal.synchronize() - _x = adapt(Array, x) + _x = x isa DenseArray && isbitstype(eltype(x)) ? + Dagger.pinned_host_array(x) : adapt(Array, x) Metal.synchronize() return _x end @@ -185,8 +225,7 @@ end function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::CPUProc, x::MtlArray{T,N}) where {T,N} with_context(from_proc) do Metal.synchronize() - _x = Array{T,N}(undef, size(x)) - copyto!(_x, x) + _x = Dagger.pinned_host_array(x) Metal.synchronize() return _x end @@ -199,15 +238,25 @@ function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::MtlArrayDeviceProc, arr = unwrap(x) with_context(Metal.synchronize, from_proc) return arr - # FIXME: elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) + elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) + # Same process but different GPUs, use DtoD copy + from_arr = unwrap(x) + with_context(Metal.synchronize, from_proc) + return with_context(to_proc) do + to_arr = similar(from_arr) + copyto!(to_arr, from_arr) + Metal.synchronize() + return to_arr + end else - # Different node, use DtoH, serialization, HtoD + # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x return with_context(from_proc) do - Array(unwrap(x)) + Dagger.pinned_host_array(unwrap(x)) end end return with_context(to_proc) do + Dagger.pin_buffer!(:Metal, host_copy) return MtlArray(host_copy) end end @@ -218,13 +267,21 @@ function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::MtlArrayDeviceProc, # Same process and GPU, no change with_context(Metal.synchronize, from_proc) return x - # FIXME: elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) + elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) + with_context(Metal.synchronize, from_proc) + return with_context(to_proc) do + to_arr = similar(x) + copyto!(to_arr, x) + Metal.synchronize() + return to_arr + end else - # Different node, use DtoH, serialization, HtoD + # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = with_context(from_proc) do - return Array(x) + return Dagger.pinned_host_array(x) end return with_context(to_proc) do + Dagger.pin_buffer!(:Metal, host_copy) return MtlArray(host_copy) end end @@ -369,8 +426,15 @@ end Dagger.gpu_processor(::Val{:Metal}) = MtlArrayDeviceProc Dagger.gpu_can_compute(::Val{:Metal}) = Metal.functional() Dagger.gpu_kernel_backend(proc::MtlArrayDeviceProc) = MetalBackend() -# TODO: Switch devices -Dagger.gpu_with_device(f, proc::MtlArrayDeviceProc) = f() +function Dagger.gpu_with_device(f, proc::MtlArrayDeviceProc) + old = Metal.device() + try + Metal.device!(DEVICES[proc.device_id]) + return f() + finally + Metal.device!(old) + end +end function Dagger.gpu_synchronize(proc::MtlArrayDeviceProc) with_context(proc) do diff --git a/ext/OpenCLExt.jl b/ext/OpenCLExt.jl index 49deefd4d..eceb9cf6f 100644 --- a/ext/OpenCLExt.jl +++ b/ext/OpenCLExt.jl @@ -52,6 +52,16 @@ function Dagger.aliasing(x::CLArray{T}) where T return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end +# MPI (SPMD) integration: stamp owning rank on CL memory spaces +Dagger.mpi_remap_space(space::CLMemorySpace, owner::Int) = + CLMemorySpace(owner, space.device) +Dagger.value_memory_space(x::CLArray) = Dagger.memory_space(x) + +# Staging-pool key; no portable OpenCL pin of arbitrary host buffers +Dagger.gpu_memory_kind(::CLArray) = :OpenCL +Dagger.gpu_memory_kind(::CLMemorySpace) = :OpenCL +Dagger.pin_buffer!(::Val{:OpenCL}, buf::DenseArray) = nothing + function Dagger.unsafe_free!(x::CLArray) OpenCL.unsafe_free!(x) return @@ -112,6 +122,17 @@ function sync_with_context(x::Union{Dagger.Processor,Dagger.MemorySpace}) end end +# No OpenCL-aware MPI / CUDA-style IPC in Julia; always host-stage +Dagger.mpi_device_direct(::CLArray) = false +function Dagger.mpi_device_sync(x::CLArray) + space = Dagger.memory_space(x) + if Dagger.root_worker_id(space) == myid() + _sync_with_context(space) + end + return +end +Dagger.mpi_device_sync(space::CLMemorySpace) = sync_with_context(space) + # Allocations Dagger.allocate_array_func(::CLArrayDeviceProc, ::typeof(rand)) = OpenCL.rand Dagger.allocate_array_func(::CLArrayDeviceProc, ::typeof(randn)) = OpenCL.randn @@ -141,13 +162,20 @@ end function Dagger.move!(to_space::CLMemorySpace, from_space::CLMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} sync_with_context(from_space) with_context!(to_space) - copyto!(to, from) + if Dagger.needs_multi_span_copy(to_space, from_space, to, from) + Dagger.multi_span_move!(to, from) + else + copyto!(to, from) + end return end # Out-of-place HtoD function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x) with_context(to_proc) do + if x isa DenseArray && isbitstype(eltype(x)) + Dagger.pin_buffer!(:OpenCL, x) + end arr = adapt(CLArray, x) cl.finish(cl.queue()) return arr @@ -159,6 +187,9 @@ function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Chunk) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) with_context(to_proc) do + if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) + Dagger.pin_buffer!(:OpenCL, cpu_data) + end arr = adapt(CLArray, cpu_data) cl.finish(cl.queue()) return arr @@ -181,7 +212,8 @@ end function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CPUProc, x) with_context(from_proc) do cl.finish(cl.queue()) - _x = adapt(Array, x) + _x = x isa DenseArray && isbitstype(eltype(x)) ? + Dagger.pinned_host_array(x) : adapt(Array, x) cl.finish(cl.queue()) return _x end @@ -198,8 +230,7 @@ end function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CPUProc, x::CLArray{T,N}) where {T,N} with_context(from_proc) do cl.finish(cl.queue()) - _x = Array{T,N}(undef, size(x)) - copyto!(_x, x) + _x = Dagger.pinned_host_array(x) cl.finish(cl.queue()) return _x end @@ -223,18 +254,19 @@ function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x return to_arr end else - # Different node, use DtoH, serialization, HtoD + # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x return with_context(from_proc) do - Array(unwrap(x)) + Dagger.pinned_host_array(unwrap(x)) end end return with_context(to_proc) do + Dagger.pin_buffer!(:OpenCL, host_copy) return CLArray(host_copy) end end end -function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x::CLArray) where T<:CLArray +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x::CLArray) if from_proc == to_proc # Same process and GPU, no change _sync_with_context(from_proc) @@ -249,11 +281,12 @@ function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x return to_arr end else - # Different node, use DtoH, serialization, HtoD + # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = with_context(from_proc) do - return Array(x) + return Dagger.pinned_host_array(x) end return with_context(to_proc) do + Dagger.pin_buffer!(:OpenCL, host_copy) return CLArray(host_copy) end end diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index 1f65c743c..70d836369 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -47,6 +47,30 @@ function Dagger.aliasing(x::ROCArray{T}) where T return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end +# MPI (SPMD) integration: aliasing spans broadcast from an owner rank must be +# stamped with that rank so same-device addresses on different ranks never +# falsely alias (every rank has myid() == 1 under SPMD) +Dagger.mpi_remap_space(space::ROCVRAMMemorySpace, owner::Int) = + ROCVRAMMemorySpace(owner, space.device_id) +Dagger.value_memory_space(x::ROCArray) = Dagger.memory_space(x) + +# Page-lock host staging buffers (~2x DtoH/HtoD bandwidth); AMDGPU.Mem.pin +# does not register a finalizer, so we unpin on GC +Dagger.gpu_memory_kind(::ROCArray) = :ROC +Dagger.gpu_memory_kind(::ROCVRAMMemorySpace) = :ROC +function Dagger.pin_buffer!(::Val{:ROC}, buf::DenseArray) + ptr = Ptr{Cvoid}(pointer(buf)) + sz = sizeof(buf) + # Avoid double-registering a finalizer if already page-locked + if !AMDGPU.Mem.is_pinned(ptr) + AMDGPU.Mem.pin(ptr, sz) + finalizer(buf) do b + AMDGPU.Mem.unpin(Ptr{Cvoid}(pointer(b))) + end + end + return nothing +end + function Dagger.unsafe_free!(x::ROCArray) AMDGPU.unsafe_free!(x) return @@ -142,13 +166,20 @@ end function Dagger.move!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} sync_with_context(from_space) with_context!(to_space) - copyto!(to, from) + if Dagger.needs_multi_span_copy(to_space, from_space, to, from) + Dagger.multi_span_move!(to, from) + else + copyto!(to, from) + end return end # Out-of-place HtoD function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x) with_context(to_proc) do + if x isa DenseArray && isbitstype(eltype(x)) + Dagger.pin_buffer!(:ROC, x) + end arr = adapt(ROCArray, x) AMDGPU.synchronize() return arr @@ -160,6 +191,9 @@ function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Chunk) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) with_context(to_proc) do + if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) + Dagger.pin_buffer!(:ROC, cpu_data) + end arr = adapt(ROCArray, cpu_data) AMDGPU.synchronize() return arr @@ -181,7 +215,8 @@ end function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, x) with_context(from_proc) do AMDGPU.synchronize() - _x = adapt(Array, x) + _x = x isa DenseArray && isbitstype(eltype(x)) ? + Dagger.pinned_host_array(x) : adapt(Array, x) AMDGPU.synchronize() return _x end @@ -198,8 +233,7 @@ end function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, x::ROCArray{T,N}) where {T,N} with_context(AMDGPU.device(x).device_id) do AMDGPU.synchronize() - _x = Array{T,N}(undef, size(x)) - copyto!(_x, x) + _x = Dagger.pinned_host_array(x) AMDGPU.synchronize() return _x end @@ -224,13 +258,14 @@ function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, to_arr end else - # Different node, use DtoH, serialization, HtoD + # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x return with_context(from_proc) do - Array(unwrap(x)) + Dagger.pinned_host_array(unwrap(x)) end end return with_context(to_proc) do + Dagger.pin_buffer!(:ROC, host_copy) return ROCArray(host_copy) end end @@ -252,10 +287,11 @@ function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, else host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x return with_context(from_proc) do - Array(unwrap(x)) + Dagger.pinned_host_array(x) end end return with_context(to_proc) do + Dagger.pin_buffer!(:ROC, host_copy) return ROCArray(host_copy) end end @@ -402,6 +438,89 @@ end Dagger.scope_key_precedence(::Val{:rocm_gpu}) = 2 Dagger.scope_key_precedence(::Val{:rocm_gpus}) = 1 +# MPI data plane: pass ROCArrays to MPI directly when the library is ROCm-aware +# (DAGGER_MPI_GPU_DIRECT=0/1 forces the decision; detection is best-effort) +const MPI_GPU_DIRECT = Ref{Union{Nothing,Bool}}(nothing) +function mpi_gpu_direct_enabled() + v = MPI_GPU_DIRECT[] + v !== nothing && return v + env = get(ENV, "DAGGER_MPI_GPU_DIRECT", "") + v = if !isempty(env) + something(tryparse(Bool, env), false) + else + Dagger.mpi_library_gpu_aware(:ROC) + end + MPI_GPU_DIRECT[] = v + return v +end +Dagger.mpi_device_direct(x::AMDGPU.StridedROCArray) = mpi_gpu_direct_enabled() +Dagger.mpi_device_sync(x::ROCArray) = AMDGPU.HIP.device_synchronize() +Dagger.mpi_device_sync(::ROCVRAMMemorySpace) = AMDGPU.HIP.device_synchronize() + +# Same-node device IPC via hipIpc*: pool-backed ROCArrays may not be +# exportable, so the sender stages into a hipMalloc allocation and ships its +# 64-byte handle; the receiver maps it and copies device-to-device. +# DAGGER_IPC=0 disables the path. +const GPU_IPC = Ref{Union{Nothing,Bool}}(nothing) +function ipc_enabled() + v = GPU_IPC[] + v !== nothing && return v + v = something(tryparse(Bool, get(ENV, "DAGGER_IPC", "true")), true) + GPU_IPC[] = v + return v +end +Dagger.ipc_eligible(::ROCVRAMMemorySpace, ::ROCVRAMMemorySpace) = ipc_enabled() + +struct HIPIpcInfo{T,N} + handle::AMDGPU.HIP.hipIpcMemHandle_t + shape::Dims{N} +end +# Token keeps a hipMalloc'd staging allocation alive until ipc_release! +struct HIPIpcToken + ptr::Ptr{Cvoid} + bytesize::Int +end +function Dagger.ipc_export(value::AMDGPU.StridedROCArray{T,N}) where {T,N} + with_context!(Dagger.memory_space(value)) + nbytes = sizeof(value) + ptr_ref = Ref{Ptr{Cvoid}}() + AMDGPU.HIP.hipMalloc(ptr_ref, nbytes) + staged = unsafe_wrap(ROCArray, Ptr{T}(ptr_ref[]), size(value); own=false) + copyto!(staged, value) + AMDGPU.HIP.device_synchronize() + handle_ref = Ref{AMDGPU.HIP.hipIpcMemHandle_t}() + AMDGPU.HIP.hipIpcGetMemHandle(handle_ref, ptr_ref[]) + return HIPIpcInfo{T,N}(handle_ref[], size(value)), HIPIpcToken(ptr_ref[], nbytes) +end +function Dagger.ipc_release!(token::HIPIpcToken) + token.ptr == C_NULL && return + AMDGPU.HIP.hipFree(token.ptr) + return +end +function _ipc_open(info::HIPIpcInfo{T,N}) where {T,N} + ptr_ref = Ref{Ptr{Cvoid}}() + AMDGPU.HIP.hipIpcOpenMemHandle(ptr_ref, info.handle, + AMDGPU.HIP.hipIpcMemLazyEnablePeerAccess) + src = unsafe_wrap(ROCArray, Ptr{T}(ptr_ref[]), info.shape; own=false) + return src, ptr_ref[] +end +function Dagger.ipc_copyto!(dest::AMDGPU.StridedROCArray{T,N}, info::HIPIpcInfo{T,N}) where {T,N} + @assert size(dest) == info.shape "IPC shape mismatch: $(size(dest)) != $(info.shape)" + with_context!(Dagger.memory_space(dest)) + src, raw = _ipc_open(info) + try + copyto!(dest, src) + AMDGPU.HIP.device_synchronize() + finally + AMDGPU.HIP.hipIpcCloseMemHandle(raw) + end + return dest +end +function Dagger.ipc_materialize(info::HIPIpcInfo{T,N}) where {T,N} + dest = ROCArray{T,N}(undef, info.shape) + return Dagger.ipc_copyto!(dest, info) +end + const DEVICES = Dict{Int, HIPDevice}() const CONTEXTS = Dict{Int, HIPContext}() const STREAMS = Dict{Int, HIPStream}() diff --git a/src/Dagger.jl b/src/Dagger.jl index 593eee4aa..1a4099bee 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -53,6 +53,13 @@ import Adapt include("lib/util.jl") include("utils/dagdebug.jl") +# Type definitions (for MPI/acceleration) +include("types/processor.jl") +include("types/scope.jl") +include("types/memory-space.jl") +include("types/chunk.jl") +include("types/acceleration.jl") + # Distributed data include("utils/locked-object.jl") include("utils/tasks.jl") @@ -77,12 +84,14 @@ include("queue.jl") include("thunk.jl") include("utils/fetch.jl") include("utils/chunks.jl") +include("weakchunk.jl") include("utils/logging.jl") include("submission.jl") abstract type MemorySpace end include("utils/memory-span.jl") include("utils/interval_tree.jl") include("memory-spaces.jl") +include("acceleration.jl") # Task scheduling include("compute.jl") @@ -90,6 +99,7 @@ include("utils/clock.jl") include("utils/system_uuid.jl") include("utils/caching.jl") include("sch/Sch.jl"); using .Sch +include("tochunk.jl") # Data dependency task queue include("datadeps/aliasing.jl") @@ -111,6 +121,7 @@ include("stream-transfer.jl") include("file-io.jl") # Array computations +include("procgrid.jl") include("array/darray.jl") include("array/alloc.jl") include("array/map-reduce.jl") @@ -133,6 +144,7 @@ include("array/svd.jl") # GPU include("gpu.jl") +include("utils/span_copy.jl") # Logging include("utils/logging-events.jl") diff --git a/src/acceleration.jl b/src/acceleration.jl new file mode 100644 index 000000000..99a71e1e2 --- /dev/null +++ b/src/acceleration.jl @@ -0,0 +1,158 @@ +const ACCELERATION = TaskLocalValue{Acceleration}(() -> DistributedAcceleration()) + +current_acceleration() = ACCELERATION[] + +default_processor(::DistributedAcceleration) = OSProc(myid()) +default_processor(accel::DistributedAcceleration, x) = default_processor(accel) +default_processor() = default_processor(current_acceleration()) + +accelerate!(accel::Symbol) = accelerate!(Val{accel}()) +accelerate!(::Val{:distributed}) = accelerate!(DistributedAcceleration()) + +function _with_default_acceleration(f) + old_accel = ACCELERATION[] + ACCELERATION[] = DistributedAcceleration() + result = try + f() + finally + ACCELERATION[] = old_accel + end + return result +end + +# `do_task` runs on pooled/reusable tasks, so `ACCELERATION` (a TaskLocalValue) +# can carry over a previous task's value. Rather than wrapping the whole task +# body in a try/finally closure to restore it, each `do_task` sets it once up +# front: the value is always overwritten before anything reads it, and no other +# code runs on those pooled tasks, so a restore is unnecessary. `nothing` (an +# unstamped task) falls back to the Distributed default. +set_task_acceleration!(accel::Acceleration) = (ACCELERATION[] = accel; nothing) +set_task_acceleration!(::Nothing) = (ACCELERATION[] = DistributedAcceleration(); nothing) + +initialize_acceleration!(a::DistributedAcceleration) = nothing +finalize_acceleration!(::Acceleration) = nothing + +"Whether two accelerations represent the same installed backend (type + fields)." +same_acceleration(::Acceleration, ::Acceleration) = false +same_acceleration(a::T, b::T) where {T<:Acceleration} = a == b + +function accelerate!(accel::Acceleration) + old = ACCELERATION[] + if !same_acceleration(old, accel) + finalize_acceleration!(old) + initialize_acceleration!(accel) + end + ACCELERATION[] = accel +end + +accel_matches_proc(accel::DistributedAcceleration, proc::OSProc) = true +accel_matches_proc(accel::DistributedAcceleration, proc) = true + +function compatible_processors(accel::Union{Acceleration,Nothing}, scope::AbstractScope, procs::Vector{<:Processor}) + comp = compatible_processors(scope, procs) + accel === nothing && return comp + return Set(p for p in comp if accel_matches_proc(accel, p)) +end +compatible_processors(accel::Union{Acceleration,Nothing}, scope::AbstractScope=get_compute_scope(), ctx::Context=Sch.eager_context()) = + compatible_processors(accel, scope, procs(ctx)) + +uniform_execution(::DistributedAcceleration) = false +uniform_execution() = uniform_execution(current_acceleration()) + +# SPMD uniformity checking. Under a uniform-execution acceleration (e.g. MPI), +# planning-time values are hashed and compared across ranks to catch divergence +# early. Core datadeps calls `check_uniform` unconditionally; the real checks +# live in MPIExt, so core owns the toggle plus a no-op fallback that always +# reports uniformity when no such extension is loaded. +const CHECK_UNIFORMITY = Ref{Bool}(false) +check_uniformity!(check::Bool=true) = (CHECK_UNIFORMITY[] = check) +# Generic fallback: when checking is off or execution isn't uniform (e.g. no +# MPIExt loaded), this is a no-op. Otherwise route through `hash` so MPIExt's +# `check_uniform(::Integer, ...)` performs the cross-rank comparison. MPIExt +# adds type-specific methods (more specific than this one) rather than +# overwriting it, which precompilation forbids. +check_uniform(value, original=value) = + CHECK_UNIFORMITY[] && uniform_execution() ? check_uniform(hash(value), original) : true + +# Scheduler hooks: Sch.jl calls these instead of MPI-specific symbols. +# Uniform-execution accelerations (e.g. MPI) inherit capacity/occupancy +# behavior via `uniform_execution`; move/cleanup hooks default to no-ops / +# async spawn and are overridden per backend where needed. +scheduling_ignore_capacity(accel::Acceleration) = uniform_execution(accel) + +scheduling_task_occupancy(accel::Acceleration, est_occupancy::UInt32) = + uniform_execution(accel) ? UInt32(0) : est_occupancy + +# Region-end cleanup (wait_all): reclaim accel state for finished tasks in one +# batch instead of locking on every Sch finish_task!. +cleanup_task_accel!(::Acceleration, task) = nothing +cleanup_tasks_accel!(accel::Acceleration, tasks) = + foreach(t -> cleanup_task_accel!(accel, t), tasks) + +schedule_argument_move(::Acceleration, ::Integer, f::Function) = Threads.@spawn f() + +""" + bind_moved_argument(accel, original, moved) -> bound + +After a scheduler argument move, choose the value stored on the Argument. The +default simply stores `moved` — the actual moved value — so a genuine `nothing` +(e.g. a `Chunk{Nothing}` from an ordinary task) reaches the kernel as `nothing`, +not as a `Chunk` the user never asked for. + +Accelerations that use non-resident placeholders override this: under MPI a +`move` returning `nothing` on a non-owning rank means "data not local", so the +`Chunk` is kept to preserve SPMD-uniform type metadata (and `execute!` later +materializes it), and an owner-unwrapped `Chunk` is rewrapped to keep its +`chunktype` uniform across ranks. +""" +bind_moved_argument(::Acceleration, original, moved) = moved + +"Logical execution rank of `proc` used for uniform task placement preference." +processor_owner_rank(accel::Acceleration, proc::Processor) = first(fire_order_key(proc)) + +"Deterministic ordering key for `proc` under uniform execution." +processor_order_key(accel::Acceleration, proc::Processor) = + (fire_order_key(proc)..., short_name(proc)) + +"Preferred execution rank for `task`, or `nothing` if no preference." +task_processor_preference(accel::Acceleration, task, state) = nothing + +""" + select_processors_uniform!(procs, accel, task=nothing, state=nothing) + +Re-sort `procs` with a deterministic, acceleration-dispatched key when +`uniform_execution(accel)` so every rank gets the same order (and with a +`task` the acceleration may prefer the rank owning task argument data). +When not under uniform execution, leaves `procs` unchanged. +""" +function select_processors_uniform!(procs, accel::Acceleration, + task=nothing, state=nothing) + uniform_execution(accel) || return procs + pref = task_processor_preference(accel, task, state) + sort!(procs, by=p -> ( + pref !== nothing && processor_owner_rank(accel, p) != pref, + processor_order_key(accel, p)..., + )) + return procs +end + +# Default processor grid for block assignment of new distributed arrays. +# Returns `nothing` to let the scheduler choose freely; uniform-execution +# accelerations (MPI) override this to produce a deterministic assignment +# that is identical on every rank. +default_procgrid(accel, nblocks::NTuple{N,Int}) where N = nothing + +"Acceleration-specific finalization after array chunk tasks are spawned." +post_stage_array_chunks!(accel::Acceleration, tasks, eltype::Type, nd::Int) = tasks + +default_processor(space::CPURAMMemorySpace) = OSProc(space.owner) +default_memory_space(accel::DistributedAcceleration) = CPURAMMemorySpace(myid()) +# Chunk records must be labeled with the space the value actually resides in +# (e.g. GPU memory for device arrays) +default_memory_space(accel::DistributedAcceleration, x) = value_memory_space(x) +default_memory_space(x) = default_memory_space(current_acceleration(), x) +default_memory_space() = default_memory_space(current_acceleration()) + +# Whether two logical owners share a physical node (for device IPC fast paths). +# Defaults false; Distributed and MPI provide concrete methods. +same_node(::Acceleration, ::Integer, ::Integer) = false diff --git a/src/array/alloc.jl b/src/array/alloc.jl index aa1050210..93862d545 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -9,49 +9,11 @@ mutable struct AllocateArray{T,N} <: ArrayOp{T,N} domain::ArrayDomain{N} domainchunks partitioning::AbstractBlocks{N} - procgrid::Union{AbstractArray{<:Processor, N}, Nothing} + procgrid::Union{AbstractProcGrid{N}, Nothing} function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, domainchunks, p::AbstractBlocks{N}, assignment::Union{AssignmentType{N},Nothing} = nothing) where {T,N} sizeA = map(length, d.indexes) - procgrid = nothing - availprocs = collect(Dagger.compatible_processors()) - if !(assignment isa AbstractArray{<:Processor, N}) - filter!(p -> p isa ThreadProc, availprocs) - sort!(availprocs, by = x -> (x.owner, x.tid)) - end - np = length(availprocs) - if assignment isa Symbol - if assignment == :arbitrary - procgrid = nothing - elseif assignment == :blockrow - q = ntuple(i -> i == 1 ? Int(ceil(sizeA[1] / p.blocksize[1])) : 1, N) - rows_per_proc, extra = divrem(Int(ceil(sizeA[1] / p.blocksize[1])), np) - counts = [rows_per_proc + (i <= extra ? 1 : 0) for i in 1:np] - procgrid = reshape(vcat(fill.(availprocs, counts)...), q) - elseif assignment == :blockcol - q = ntuple(i -> i == N ? Int(ceil(sizeA[N] / p.blocksize[N])) : 1, N) - cols_per_proc, extra = divrem(Int(ceil(sizeA[N] / p.blocksize[N])), np) - counts = [cols_per_proc + (i <= extra ? 1 : 0) for i in 1:np] - procgrid = reshape(vcat(fill.(availprocs, counts)...), q) - elseif assignment == :cyclicrow - q = ntuple(i -> i == 1 ? np : 1, N) - procgrid = reshape(availprocs, q) - elseif assignment == :cycliccol - q = ntuple(i -> i == N ? np : 1, N) - procgrid = reshape(availprocs, q) - else - error("Unsupported assignment symbol: $assignment, use :arbitrary, :blockrow, :blockcol, :cyclicrow or :cycliccol") - end - elseif assignment isa AbstractArray{<:Int, N} - missingprocs = filter(q -> q ∉ procs(), assignment) - isempty(missingprocs) || error("Specified workers are not available: $missingprocs") - procgrid = [ThreadProc(proc, 1) for proc in assignment] - elseif assignment isa AbstractArray{<:Processor, N} - missingprocs = filter(q -> q ∉ availprocs, assignment) - isempty(missingprocs) || error("Specified processors are not available: $missingprocs") - procgrid = assignment - end - + procgrid = build_procgrid(something(assignment, :arbitrary), Tuple(sizeA), p.blocksize, current_acceleration()) return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid) end @@ -80,27 +42,17 @@ function allocate_array(f, T, sz) end allocate_array_func(::Processor, f) = f function stage(ctx, A::AllocateArray) - tasks = Array{DTask,ndims(A.domainchunks)}(undef, size(A.domainchunks)...) - Dagger.spawn_datadeps() do - default_scope = get_compute_scope() - for I in CartesianIndices(A.domainchunks) - x = A.domainchunks[I] - i = LinearIndices(A.domainchunks)[I] - - if isnothing(A.procgrid) - scope = default_scope - else - scope = ExactScope(A.procgrid[CartesianIndex(mod1.(Tuple(I), size(A.procgrid))...)]) - end - - if A.want_index - task = Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, i, size(x)) - else - task = Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, size(x)) - end - tasks[i] = task + tasks = emit_chunk_tasks!(A.domainchunks, A.procgrid, A.eltype, + (scope, I, i) -> begin + x = A.domainchunks[I] + N = ndims(A.domainchunks) + ret_type = Array{A.eltype, N} + if A.want_index + Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, i, size(x)) + else + Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, size(x)) end - end + end) return DArray(A.eltype, A.domain, A.domainchunks, tasks, A.partitioning) end diff --git a/src/array/darray.jl b/src/array/darray.jl index fb27adb7a..bd0e5efc5 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -1,4 +1,4 @@ -import Base: ==, fetch +import Base: ==, fetch, length, isempty, size export DArray, DVector, DMatrix, DVecOrMat, Blocks, AutoBlocks export distribute @@ -36,7 +36,7 @@ Base.getindex(arr::AbstractArray{T,0} where T, d::ArrayDomain{0}) = arr Base.getindex(arr::GPUArraysCore.AbstractGPUArray, d::ArrayDomain) = arr[indexes(d)...] Base.getindex(arr::GPUArraysCore.AbstractGPUArray{T,0} where T, d::ArrayDomain{0}) = arr -function intersect(a::ArrayDomain, b::ArrayDomain) +function Base.intersect(a::ArrayDomain, b::ArrayDomain) if a === b return a end @@ -73,7 +73,7 @@ function size(a::ArrayDomain, dim) end size(a::ArrayDomain) = map(length, indexes(a)) length(a::ArrayDomain) = prod(size(a)) -ndims(a::ArrayDomain) = length(size(a)) +Base.ndims(a::ArrayDomain) = length(size(a)) isempty(a::ArrayDomain) = length(a) == 0 @@ -84,7 +84,6 @@ The domain of an array is an ArrayDomain. """ domain(x::AbstractArray) = ArrayDomain([1:l for l in size(x)]) - abstract type ArrayOp{T, N} <: AbstractArray{T, N} end Base.IndexStyle(::Type{<:ArrayOp}) = IndexCartesian() @@ -174,9 +173,14 @@ domain(d::DArray) = d.domain chunks(d::DArray) = d.chunks domainchunks(d::DArray) = d.subdomains size(x::DArray) = size(domain(x)) +Base.ndims(d::DArray{T,N}) where {T,N} = N stage(ctx, c::DArray) = c -function Base.collect(d::DArray{T,N}; tree=false, copyto=false) where {T,N} +# Named so the spawned function is rank-uniform under MPI (anonymous +# closures get non-uniform ArgumentWrapper hashes). +_collect_cat(concat, dims::Int, xs...) = concat(xs...; dims) + +function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} a = fetch(d) if isempty(d.chunks) return Array{eltype(d)}(undef, size(d)...) @@ -193,12 +197,24 @@ function Base.collect(d::DArray{T,N}; tree=false, copyto=false) where {T,N} return C end - dimcatfuncs = [(x...) -> d.concat(x..., dims=i) for i in 1:ndims(d)] - if tree - collect(fetch(treereduce_nd(map(x -> ((args...,) -> Dagger.@spawn x(args...)) , dimcatfuncs), a.chunks))) - else - collect(treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) + concat = d.concat + if uniform_execution() + # Under SPMD/MPI, gather chunks via a datadeps concat tree so chunk + # movement is scheduled uniformly across ranks (rather than a direct + # per-chunk `fetch`, which isn't well-defined for rank-owned data). + spawn_catfuncs = ntuple(Val(N)) do i + (args...) -> Dagger.@spawn _collect_cat(concat, i, map(In, args)...) + end + result = Dagger.spawn_datadeps() do + treereduce_nd(spawn_catfuncs, a.chunks) + end + return collect(fetch(result)) end + # Distributed: fetch chunks directly and concat in-process. This avoids + # routing chunk data through datadeps aliasing, which requires an + # aliasing-resolvable (e.g. isbits) element type. + dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:N] + return collect(treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) end Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A)) @@ -401,6 +417,18 @@ If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk. """ function Base.fetch(c::DArray{T}) where T if any(istask, chunks(c)) + if uniform_execution() + # SPMD (MPI): every rank holds every task locally, and chunk + # records are rank-local views (placeholders for non-owned data). + # Resolve them locally instead of assembling on one rank and + # broadcasting (which would clobber the local records with the + # assembling rank's placeholders). + thunks = chunks(c) + new_chunks = Any[istask(t) ? fetch(t; raw=true) : t for t in thunks] + return DArray(T, domain(c), domainchunks(c), + reshape(new_chunks, size(thunks)), + c.partitioning, c.concat) + end thunks = chunks(c) sz = size(thunks) dmn = domain(c) @@ -422,7 +450,7 @@ struct Distribute{T,N,B<:AbstractBlocks} <: ArrayOp{T, N} domainchunks partitioning::B data::AbstractArray{T,N} - procgrid::Union{AbstractArray{<:Processor, N}, Nothing} + procgrid::Union{AbstractProcGrid{N}, Nothing} end size(x::Distribute) = size(domain(x.data)) @@ -430,18 +458,18 @@ size(x::Distribute) = size(domain(x.data)) Base.@deprecate BlockPartition Blocks -Distribute(p::Blocks, data::AbstractArray, procgrid::Union{AbstractArray{<:Processor},Nothing} = nothing) = - Distribute(partition(p, domain(data)), p, data, procgrid) +Distribute(p::Blocks, data::AbstractArray, procgrid::Union{AbstractProcGrid,AbstractArray{<:Processor},Nothing} = nothing) = + Distribute(partition(p, domain(data)), p, data, normalize_procgrid(procgrid)) -function Distribute(domainchunks::DomainBlocks{N}, data::AbstractArray{T,N}, procgrid::Union{AbstractArray{<:Processor, N},Nothing} = nothing) where {T,N} +function Distribute(domainchunks::DomainBlocks{N}, data::AbstractArray{T,N}, procgrid::Union{AbstractProcGrid{N},AbstractArray{<:Processor, N},Nothing} = nothing) where {T,N} p = Blocks(ntuple(i->first(domainchunks.cumlength[i]), N)) - Distribute(domainchunks, p, data, procgrid) + Distribute(domainchunks, p, data, normalize_procgrid(procgrid)) end -function Distribute(data::AbstractArray{T,N}, procgrid::Union{AbstractArray{<:Processor, N},Nothing} = nothing) where {T,N} +function Distribute(data::AbstractArray{T,N}, procgrid::Union{AbstractProcGrid{N},AbstractArray{<:Processor, N},Nothing} = nothing) where {T,N} nprocs = sum(w->length(get_processors(OSProc(w))),procs()) p = Blocks(ntuple(i->max(cld(size(data, i), nprocs), 1), N)) - return Distribute(partition(p, domain(data)), p, data, procgrid) + return Distribute(partition(p, domain(data)), p, data, normalize_procgrid(procgrid)) end function stage(ctx::Context, d::Distribute) @@ -458,13 +486,7 @@ function stage(ctx::Context, d::Distribute) idx = d.domainchunks[I] chunks = stage(ctx, x[idx]).chunks shape = size(chunks) - # TODO: fix hashing - #hash = uhash(idx, Base.hash(Distribute, Base.hash(d.data))) - if isnothing(d.procgrid) - scope = get_compute_scope() - else - scope = ExactScope(d.procgrid[CartesianIndex(mod1.(Tuple(I), size(d.procgrid))...)]) - end + scope = procgrid_scope(d.procgrid, I) options = Options(compute_scope=scope) Dagger.spawn(options, shape, chunks...) do shape, parts... if prod(shape) == 0 @@ -475,18 +497,14 @@ function stage(ctx::Context, d::Distribute) collect(treereduce_nd(dimcatfuncs, ps)) end end + cs = post_stage_array_chunks!(current_acceleration(), cs, T, Nd) else - cs = map(CartesianIndices(d.domainchunks)) do I - # TODO: fix hashing - #hash = uhash(c, Base.hash(Distribute, Base.hash(d.data))) + T = eltype(d.data) + cs = emit_chunk_tasks!(d.domainchunks, d.procgrid, T, + (scope, I, i) -> begin c = d.domainchunks[I] - if isnothing(d.procgrid) - scope = get_compute_scope() - else - scope = ExactScope(d.procgrid[CartesianIndex(mod1.(Tuple(I), size(d.procgrid))...)]) - end Dagger.@spawn compute_scope=scope identity(d.data[c]) - end + end) end return DArray(eltype(d.data), domain(d.data), @@ -510,49 +528,9 @@ function auto_blocks(dims::Dims{N}) where N end auto_blocks(A::AbstractArray{T,N}) where {T,N} = auto_blocks(size(A)) -const AssignmentType{N} = Union{Symbol, AbstractArray{<:Int, N}, AbstractArray{<:Processor, N}} - distribute(A::AbstractArray, assignment::AssignmentType = :arbitrary) = distribute(A, AutoBlocks(), assignment) function distribute(A::AbstractArray{T,N}, dist::Blocks{N}, assignment::AssignmentType{N} = :arbitrary) where {T,N} - procgrid = nothing - availprocs = collect(Dagger.compatible_processors()) - if !(assignment isa AbstractArray{<:Processor, N}) - filter!(p -> p isa ThreadProc, availprocs) - sort!(availprocs, by = x -> (x.owner, x.tid)) - end - np = length(availprocs) - if assignment isa Symbol - if assignment == :arbitrary - procgrid = nothing - elseif assignment == :blockrow - p = ntuple(i -> i == 1 ? Int(ceil(size(A,1) / dist.blocksize[1])) : 1, N) - rows_per_proc, extra = divrem(Int(ceil(size(A,1) / dist.blocksize[1])), np) - counts = [rows_per_proc + (i <= extra ? 1 : 0) for i in 1:np] - procgrid = reshape(vcat(fill.(availprocs, counts)...), p) - elseif assignment == :blockcol - p = ntuple(i -> i == N ? Int(ceil(size(A,N) / dist.blocksize[N])) : 1, N) - cols_per_proc, extra = divrem(Int(ceil(size(A,N) / dist.blocksize[N])), np) - counts = [cols_per_proc + (i <= extra ? 1 : 0) for i in 1:np] - procgrid = reshape(vcat(fill.(availprocs, counts)...), p) - elseif assignment == :cyclicrow - p = ntuple(i -> i == 1 ? np : 1, N) - procgrid = reshape(availprocs, p) - elseif assignment == :cycliccol - p = ntuple(i -> i == N ? np : 1, N) - procgrid = reshape(availprocs, p) - else - error("Unsupported assignment symbol: $assignment, use :arbitrary, :blockrow, :blockcol, :cyclicrow or :cycliccol") - end - elseif assignment isa AbstractArray{<:Int, N} - missingprocs = filter(p -> p ∉ procs(), assignment) - isempty(missingprocs) || error("Specified workers are not available: $missingprocs") - procgrid = [ThreadProc(proc, 1) for proc in assignment] - elseif assignment isa AbstractArray{<:Processor, N} - missingprocs = filter(p -> p ∉ availprocs, assignment) - isempty(missingprocs) || error("Specified processors are not available: $missingprocs") - procgrid = assignment - end - + procgrid = build_procgrid(assignment, size(A), dist.blocksize, current_acceleration()) return _to_darray(Distribute(dist, A, procgrid)) end @@ -625,7 +603,7 @@ end mapchunk(f, chunk) = tochunk(f(poolget(chunk.handle))) function mapchunks(f, d::DArray{T,N,F}) where {T,N,F} chunks = map(d.chunks) do chunk - owner = get_parent(chunk.processor).pid + owner = root_worker_id(chunk.processor) remotecall_fetch(mapchunk, owner, f, chunk) end DArray{T,N,F}(d.domain, d.subdomains, chunks, d.concat) diff --git a/src/array/linalg.jl b/src/array/linalg.jl index 6d6e32552..2fcaa42fa 100644 --- a/src/array/linalg.jl +++ b/src/array/linalg.jl @@ -71,10 +71,17 @@ function LinearAlgebra.issymmetric(A::DArray{T,2}) where T Ac = A.chunks mt = size(Ac, 1) - to_check = [Dagger.@spawn issymmetric(Ac[i, i]) for i in 1:mt] - for i in 2:mt - for j in 1:i-1 - push!(to_check, Dagger.@spawn is_cross_symmetric(Ac[i, j], Ac[j, i])) + # The cross-block checks read pairs of blocks that may live in different + # memory spaces (e.g. different MPI ranks); datadeps handles the transfers + to_check = DTask[] + Dagger.spawn_datadeps() do + for i in 1:mt + push!(to_check, Dagger.@spawn issymmetric(In(Ac[i, i]))) + end + for i in 2:mt + for j in 1:i-1 + push!(to_check, Dagger.@spawn is_cross_symmetric(In(Ac[i, j]), In(Ac[j, i]))) + end end end @@ -99,10 +106,17 @@ function LinearAlgebra.ishermitian(A::DArray{T,2}) where T Ac = A.chunks mt = size(Ac, 1) - to_check = [Dagger.@spawn ishermitian(Ac[i, i]) for i in 1:mt] - for i in 2:mt - for j in 1:i-1 - push!(to_check, Dagger.@spawn is_cross_hermitian(Ac[i, j], Ac[j, i])) + # The cross-block checks read pairs of blocks that may live in different + # memory spaces (e.g. different MPI ranks); datadeps handles the transfers + to_check = DTask[] + Dagger.spawn_datadeps() do + for i in 1:mt + push!(to_check, Dagger.@spawn ishermitian(In(Ac[i, i]))) + end + for i in 2:mt + for j in 1:i-1 + push!(to_check, Dagger.@spawn is_cross_hermitian(In(Ac[i, j]), In(Ac[j, i]))) + end end end diff --git a/src/array/lu.jl b/src/array/lu.jl index 2f63e8bf6..45d0c2c38 100644 --- a/src/array/lu.jl +++ b/src/array/lu.jl @@ -107,6 +107,7 @@ function swaprows_panel!(A::AbstractMatrix{T}, M::AbstractMatrix{T}, ipiv_chunk: A[p,:], M[r,:] = M[r,:], A[p,:] end end + return A end @kernel function _geru_kernel!(alpha, x, y, A) @@ -173,6 +174,7 @@ end end end end + return A end # Same as above, but for a vector (e.g. a right-hand-side `DVector` in diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index fa0510b5a..6e16a912b 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -16,9 +16,26 @@ function stage(ctx::Context, node::Map) partitioning = primary.partitioning concat = primary.concat f = node.f - for i=eachindex(domains) - inps = map(x->chunks(x)[i], inputs) - thunks[i] = Dagger.@spawn map(f, inps...) + if uniform_execution() + # Under uniform (MPI) execution, spawn per-chunk datadeps tasks: wrap + # each chunk in In(...), not the whole chunks(::DArray) array. The plain + # `@spawn` path below issues a concurrent `execute!` status-broadcast wave + # that deadlocks under MPI. Datadeps serializes those moves safely. + Dagger.spawn_datadeps() do + for i in eachindex(domains) + inps = ntuple(j -> In(chunks(inputs[j])[i]), length(inputs)) + thunks[i] = Dagger.spawn(map, f, inps...) + end + end + else + # Non-uniform (e.g. Distributed) execution: spawn directly, without + # datadeps. This avoids forcing datadeps aliasing analysis on the chunk + # element type, which is unsupported for non-isbits eltypes (e.g. the + # `OnlineStats` accumulators produced by `mean`/`var`/`std`). + for i in eachindex(domains) + inps = map(x->chunks(x)[i], inputs) + thunks[i] = Dagger.@spawn map(f, inps...) + end end RT = Base.promote_op(node.f, map(eltype, node.inputs)...) return DArray(RT, domain(primary), domainchunks(primary), thunks, partitioning, concat) diff --git a/src/array/mul.jl b/src/array/mul.jl index 3a2637b76..14ed9f0ba 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -41,7 +41,7 @@ function LinearAlgebra.generic_matmatmul!( return gemm_dagger!(C, transA, transB, A, B, alpha, beta) end end -function _repartition_matmatmul(C, A, B, transA::Char, transB::Char) +function _repartition_matmatmul(C, A, B, transA::Char, transB::Char)::Tuple{Blocks{2}, Blocks{2}, Blocks{2}} partA = A.partitioning.blocksize partB = B.partitioning.blocksize istransA = transA == 'T' || transA == 'C' @@ -398,16 +398,17 @@ end return A end -@inline function copytile!(A, B) +@inline function copytile!(A::AbstractMatrix{T}, B::AbstractMatrix{T})::Nothing where {T} m, n = size(A) C = B' for i = 1:m, j = 1:n A[i, j] = C[i, j] end + return nothing end -@inline function copydiagtile!(A, uplo) +@inline function copydiagtile!(A::AbstractMatrix{T}, uplo::AbstractChar)::Nothing where {T} m, n = size(A) Acpy = copy(A) @@ -422,6 +423,7 @@ end for i = 1:m, j = 1:n A[i, j] = C[i, j] end + return nothing end function LinearAlgebra.generic_matvecmul!( C::DVector{T}, @@ -445,7 +447,7 @@ function LinearAlgebra.generic_matvecmul!( return gemv_dagger!(C, transA, A, B, _alpha, _beta) end end -function _repartition_matvecmul(C, A, B, transA::Char) +function _repartition_matvecmul(C, A, B, transA::Char)::Tuple{Blocks{1}, Blocks{2}, Blocks{1}} partA = A.partitioning.blocksize partB = B.partitioning.blocksize istransA = transA == 'T' || transA == 'C' diff --git a/src/array/sort.jl b/src/array/sort.jl index 4eb4015d9..7abbe373f 100644 --- a/src/array/sort.jl +++ b/src/array/sort.jl @@ -290,26 +290,6 @@ function dsort_chunks(cs, nchunks=length(cs), nsamples=2000; splitters, batchsize; merge=merge, by=by, sub=sub, order=order) - #= - for (w, c) in zip(Iterators.cycle(affinities), cs) - propagate_affinity!(c, Dagger.OSProc(w) => 1) - end - =# -end - -function propagate_affinity!(c, aff) - if !isa(c, DTask) - return - end - if c.affinity !== nothing - push!(c.affinity, aff) - else - c.affinity = [aff] - end - - for t in c.inputs - propagate_affinity!(t, aff) - end end function Base.sort(v::ArrayOp; diff --git a/src/array/trsm.jl b/src/array/trsm.jl index 65e87c5d5..c0c025468 100644 --- a/src/array/trsm.jl +++ b/src/array/trsm.jl @@ -189,4 +189,4 @@ function trsm!(side::Char, uplo::Char, trans::Char, diag::Char, alpha::T, A::DMa end end -end \ No newline at end of file +end diff --git a/src/chunks.jl b/src/chunks.jl index 03bdfb65d..9961bd890 100644 --- a/src/chunks.jl +++ b/src/chunks.jl @@ -24,38 +24,20 @@ we use the `UnitDomain` on it. A `UnitDomain` is indivisible. """ domain(x::Any) = UnitDomain() -###### Chunk ###### - -""" - Chunk - -A reference to a piece of data located on a remote worker. `Chunk`s are -typically created with `Dagger.tochunk(data)`, and the data can then be -accessed from any worker with `collect(::Chunk)`. `Chunk`s are -serialization-safe, and use distributed refcounting (provided by -`MemPool.DRef`) to ensure that the data referenced by a `Chunk` won't be GC'd, -as long as a reference exists on some worker. - -Each `Chunk` is associated with a given `Dagger.Processor`, which is (in a -sense) the processor that "owns" or contains the data. Calling -`collect(::Chunk)` will perform data movement and conversions defined by that -processor to safely serialize the data to the calling worker. - -## Constructors -See [`tochunk`](@ref). -""" -mutable struct Chunk{T, H, P<:Processor, S<:AbstractScope} - chunktype::Type{T} - domain - handle::H - processor::P - scope::S -end +###### Chunk Methods ###### domain(c::Chunk) = c.domain chunktype(c::Chunk) = c.chunktype processor(c::Chunk) = c.processor -affinity(c::Chunk) = affinity(c.handle) + +""" + datasize(x) + +Returns the estimated memory size of `x`'s data, used for transfer-cost estimation. +""" +datasize(c::Chunk) = datasize(c.handle) +datasize(r::DRef) = r.size +datasize(r::FileRef) = r.size is_task_or_chunk(c::Chunk) = true @@ -72,20 +54,29 @@ function collect(ctx::Context, chunk::Chunk; options=nothing) elseif chunk.handle isa FileRef return poolget(chunk.handle) else - return move(chunk.processor, OSProc(), chunk.handle) + return move(chunk.processor, default_processor(), chunk.handle) end end collect(ctx::Context, ref::DRef; options=nothing) = move(OSProc(ref.owner), OSProc(), ref) collect(ctx::Context, ref::FileRef; options=nothing) = poolget(ref) # FIXME: Do move call -function Base.fetch(chunk::Chunk; raw=false) - if raw - poolget(chunk.handle) - else - collect(chunk) +function Base.fetch(chunk::Chunk{T}; unwrap::Bool=false, uniform::Bool=uniform_execution(), kwargs...) where T + # N.B. Do not assert `::T`: the chunktype is not always the restored value + # type. File-backed chunks (`tochunk(FileRef(path); device=...)`) carry + # chunktype `FileRef` but restore to the file's deserialized contents. + value = fetch_handle(chunk.handle; uniform) + if unwrap && unwrappable(value) + return fetch(value; unwrap, uniform, kwargs...) end + return value end +fetch_handle(ref::DRef; uniform::Bool) = poolget(ref) +fetch_handle(ref::FileRef; uniform::Bool) = poolget(ref) +unwrappable(x::Chunk) = true +unwrappable(x::DRef) = true +unwrappable(x::FileRef) = true +unwrappable(x) = false # Unwrap Chunk, DRef, and FileRef by default move(from_proc::Processor, to_proc::Processor, x::Chunk) = @@ -100,32 +91,3 @@ move(to_proc::Processor, d::DRef) = move(OSProc(d.owner), to_proc, d) move(to_proc::Processor, x) = move(OSProc(), to_proc, x) - -### ChunkIO -affinity(r::DRef) = OSProc(r.owner)=>r.size -# this previously returned a vector with all machines that had the file cached -# but now only returns the owner and size, for consistency with affinity(::DRef), -# see #295 -affinity(r::FileRef) = OSProc(1)=>r.size - -struct WeakChunk - wid::Int - id::Int - x::WeakRef - function WeakChunk(c::Chunk) - return new(c.handle.owner, c.handle.id, WeakRef(c)) - end -end -unwrap_weak(c::WeakChunk) = c.x.value -function unwrap_weak_checked(c::WeakChunk) - cw = unwrap_weak(c) - @assert cw !== nothing "WeakChunk expired: ($(c.wid), $(c.id))" - return cw -end -wrap_weak(c::Chunk) = WeakChunk(c) -isweak(c::WeakChunk) = true -isweak(c::Chunk) = false -is_task_or_chunk(c::WeakChunk) = true -Serialization.serialize(io::AbstractSerializer, wc::WeakChunk) = - error("Cannot serialize a WeakChunk") -chunktype(c::WeakChunk) = chunktype(unwrap_weak_checked(c)) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 13432748a..725e1094e 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -226,6 +226,7 @@ struct ArgumentWrapper function ArgumentWrapper(arg, dep_mod) h = hash(dep_mod) h = _identity_hash(arg, h) + check_uniform(h, arg) return new(arg, dep_mod, h) end end @@ -242,6 +243,7 @@ struct HistoryEntry end struct AliasedObjectCacheStore + accel::Acceleration keys::Vector{AbstractAliasing} derived::Dict{AbstractAliasing,AbstractAliasing} stored::Dict{MemorySpace,Set{AbstractAliasing}} @@ -253,8 +255,9 @@ struct AliasedObjectCacheStore # space holding that key is a copy we allocated and may free. originals::Set{Tuple{MemorySpace,AbstractAliasing}} end -AliasedObjectCacheStore() = - AliasedObjectCacheStore(Vector{AbstractAliasing}(), +AliasedObjectCacheStore(accel::Acceleration) = + AliasedObjectCacheStore(accel, + Vector{AbstractAliasing}(), Dict{AbstractAliasing,AbstractAliasing}(), Dict{MemorySpace,Set{AbstractAliasing}}(), Dict{MemorySpace,Dict{AbstractAliasing,Chunk}}(), @@ -290,8 +293,9 @@ function get_stored(cache::AliasedObjectCacheStore, space::MemorySpace, ainfo::A end function set_stored!(cache::AliasedObjectCacheStore, dest_space::MemorySpace, value::Chunk, ainfo::AbstractAliasing) @assert !is_stored(cache, dest_space, ainfo) "Cache already has derived ainfo $ainfo" + check_uniform(value) key = cache.derived[ainfo] - value_ainfo = aliasing(value, identity) + value_ainfo = aliasing(cache.accel, value, identity) cache.derived[value_ainfo] = key push!(get!(Set{AbstractAliasing}, cache.stored, dest_space), key) values_dict = get!(Dict{AbstractAliasing,Chunk}, cache.values, dest_space) @@ -299,6 +303,7 @@ function set_stored!(cache::AliasedObjectCacheStore, dest_space::MemorySpace, va return end function set_key_stored!(cache::AliasedObjectCacheStore, space::MemorySpace, ainfo::AbstractAliasing, value::Chunk) + check_uniform(value) push!(cache.keys, ainfo) cache.derived[ainfo] = ainfo # A key is first registered at the space where the original object lives, so @@ -311,6 +316,7 @@ function set_key_stored!(cache::AliasedObjectCacheStore, space::MemorySpace, ain end struct AliasedObjectCache + accel::Acceleration space::MemorySpace chunk::Chunk end @@ -338,31 +344,34 @@ function get_stored(cache::AliasedObjectCache, ainfo::AbstractAliasing) cache_raw = unwrap(cache.chunk)::AliasedObjectCacheStore return get_stored(cache_raw, cache.space, ainfo) end -function set_stored!(cache::AliasedObjectCache, value::Chunk, ainfo::AbstractAliasing) + +function set_stored!(accel::DistributedAcceleration, cache::AliasedObjectCache, value::Chunk, ainfo::AbstractAliasing) wid = root_worker_id(cache.chunk) if wid != myid() - return remotecall_fetch(set_stored!, wid, cache, value, ainfo) + return remotecall_fetch(set_stored!, wid, accel, cache, value, ainfo) end cache_raw = unwrap(cache.chunk)::AliasedObjectCacheStore set_stored!(cache_raw, cache.space, value, ainfo) return end -function set_key_stored!(cache::AliasedObjectCache, space::MemorySpace, ainfo::AbstractAliasing, value::Chunk) + +function set_key_stored!(accel::DistributedAcceleration, cache::AliasedObjectCache, space::MemorySpace, ainfo::AbstractAliasing, value::Chunk) wid = root_worker_id(cache.chunk) if wid != myid() - return remotecall_fetch(set_key_stored!, wid, cache, space, ainfo, value) + return remotecall_fetch(set_key_stored!, wid, accel, cache, space, ainfo, value) end cache_raw = unwrap(cache.chunk)::AliasedObjectCacheStore set_key_stored!(cache_raw, space, ainfo, value) end -function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(x, identity)) + +function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(cache.accel, x, identity)) x_space = memory_space(x) if !is_key_present(cache, x_space, ainfo) # Preserve the object's memory-space/processor pairing when inserting # the source key. Using bare `tochunk(x)` defaults to OSProc, which can # incorrectly wrap GPU-backed objects as CPU chunks. x_chunk = x isa Chunk ? x : tochunk(x, first(processors(x_space))) - set_key_stored!(cache, x_space, ainfo, x_chunk) + set_key_stored!(cache.accel, cache, x_space, ainfo, x_chunk) end if is_stored(cache, ainfo) return get_stored(cache, ainfo) @@ -371,16 +380,17 @@ function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(x, iden @assert y isa Chunk "Didn't get a Chunk from functor" @assert memory_space(y) == cache.space "Space mismatch! $(memory_space(y)) != $(cache.space)" if memory_space(x) != cache.space - @assert ainfo != aliasing(y, identity) "Aliasing mismatch! $ainfo == $(aliasing(y, identity))" + @assert ainfo != aliasing(cache.accel, y, identity) "Aliasing mismatch! $ainfo == $(aliasing(cache.accel, y, identity))" end - set_stored!(cache, y, ainfo) + set_stored!(cache.accel, cache, y, ainfo) return y end end struct DataDepsState # The mapping of original raw argument to its Chunk - raw_arg_to_chunk::IdDict{Any,Chunk} + # N.B. Values are Chunks, or raw remote handles (e.g. ChunkView under MPI) + raw_arg_to_chunk::IdDict{Any,Any} # The origin memory space of each argument # Used to track the original location of an argument, for final copy-from @@ -410,6 +420,12 @@ struct DataDepsState # Used by remainder copies to lookup the "backstop" if any portion of the target ainfo is not updated by the remainder arg_owner::Dict{ArgumentWrapper,MemorySpace} + # The set of memory spaces holding a fully-current replica of each argument + # and dep_mod: a task write resets it to the written space, while a copy + # extends it (the copy's source remains current) + # Used to elide copies, notably the final copy-from of read-only arguments + arg_current::Dict{ArgumentWrapper,Set{MemorySpace}} + # The overlap of each argument with every other argument, based on the ainfo overlaps # Incrementally updated as new ainfos are created # Used for fast history updates @@ -442,7 +458,7 @@ struct DataDepsState ainfos_readers::Dict{AliasingWrapper,Vector{Pair{DTask,Int}}} function DataDepsState() - arg_to_chunk = IdDict{Any,Chunk}() + arg_to_chunk = IdDict{Any,Any}() arg_origin = IdDict{Any,MemorySpace}() remote_args = Dict{MemorySpace,IdDict{Any,Any}}() remote_arg_to_original = IdDict{Any,Any}() @@ -450,8 +466,12 @@ struct DataDepsState ainfo_arg = Dict{AliasingWrapper,Set{ArgumentWrapper}}() arg_history = Dict{ArgumentWrapper,Vector{HistoryEntry}}() arg_owner = Dict{ArgumentWrapper,MemorySpace}() + arg_current = Dict{ArgumentWrapper,Set{MemorySpace}}() arg_overlaps = Dict{ArgumentWrapper,Set{ArgumentWrapper}}() - ainfo_backing_chunk = tochunk(AliasedObjectCacheStore()) + accel = current_acceleration() + ainfo_backing_chunk = _with_default_acceleration() do + tochunk(AliasedObjectCacheStore(accel)) + end supports_inplace_cache = IdDict{Any,Bool}() ainfo_cache = Dict{ArgumentWrapper,AliasingWrapper}() @@ -462,7 +482,7 @@ struct DataDepsState ainfos_owner = Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}}() ainfos_readers = Dict{AliasingWrapper,Vector{Pair{DTask,Int}}}() - return new(arg_to_chunk, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_overlaps, ainfo_backing_chunk, + return new(arg_to_chunk, 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) end end @@ -479,6 +499,11 @@ function is_writedep(arg, deps, task::DTask) return any(dep->dep[3], deps) end +# Wrap a non-Chunk argument for datadeps tracking. Most values become local +# Chunks; remote handles that are already rank-replicated metadata under SPMD +# (e.g. ChunkView under MPI) override this to stay raw. +datadeps_arg_wrap(arg) = tochunk(arg) + # Aliasing state setup function populate_task_info!(state::DataDepsState, task_args, spec::DTaskSpec, task::DTask) # Track the task's arguments and access patterns @@ -512,7 +537,9 @@ function populate_task_info!(state::DataDepsState, task_args, spec::DTaskSpec, t arg_chunk = state.raw_arg_to_chunk[arg] else if !(arg isa Chunk) - arg_chunk = tochunk(arg) + arg_chunk = with(DATADEPS_THUNK_ID=>task.uid) do + datadeps_arg_wrap(arg) + end state.raw_arg_to_chunk[arg] = arg_chunk else state.raw_arg_to_chunk[arg] = arg @@ -522,6 +549,7 @@ function populate_task_info!(state::DataDepsState, task_args, spec::DTaskSpec, t # Track the origin space of the argument origin_space = memory_space(arg_chunk) + check_uniform(origin_space) state.arg_origin[arg_chunk] = origin_space state.remote_arg_to_original[arg_chunk] = arg_chunk @@ -583,7 +611,7 @@ function aliasing!(state::DataDepsState, target_space::MemorySpace, arg_w::Argum end # Calculate the ainfo - ainfo = AliasingWrapper(aliasing(remote_arg, arg_w.dep_mod)) + ainfo = AliasingWrapper(aliasing(current_acceleration(), remote_arg, arg_w.dep_mod)) # Cache the result state.ainfo_cache[remote_arg_w] = ainfo @@ -683,19 +711,25 @@ implemented `move!` and want to skip in-place copies. When this returns `false`, datadeps will instead perform out-of-place copies for each non-local use of `x`, and the data in `x` will not be updated when the `spawn_datadeps` region returns. + +This is deliberately a *type-level* query: it excludes values that +`type_may_alias` reports as aliasing but which cannot actually be mutated in +place (a `Function` closure that captures arrays is the canonical example — +it aliases its captures, but there is no `move!` that writes into a closure). +Deciding from the type alone (rather than inspecting the value) is what makes +this usable under SPMD/MPI execution, where a `Chunk`'s payload is not +materialized on every rank but its `chunktype` is known everywhere. + +Note this is distinct from whether a backend can perform a *zero-copy* +in-place transfer of the payload (e.g. MPIExt's `supports_inplace_mpi`, which +additionally requires a bits-eltype `DenseArray`); that is a transport-layer +optimization applied only once an in-place move has been scheduled. """ -supports_inplace_move(x) = true -supports_inplace_move(t::DTask) = supports_inplace_move(fetch(t; raw=true)) -function supports_inplace_move(c::Chunk) - # FIXME: Use MemPool.access_ref - pid = root_worker_id(c.processor) - if pid == myid() - return supports_inplace_move(poolget(c.handle)) - else - return remotecall_fetch(supports_inplace_move, pid, c) - end -end -supports_inplace_move(::Function) = false +supports_inplace_move(x) = supports_inplace_move(typeof(x)) +supports_inplace_move(t::DTask) = supports_inplace_move(chunktype(t)) +supports_inplace_move(c::Chunk) = supports_inplace_move(chunktype(c)) +supports_inplace_move(::Type) = true +supports_inplace_move(::Type{<:Function}) = false # Read/write dependency management function get_write_deps!(state::DataDepsState, dest_space::MemorySpace, ainfo::AbstractAliasing, write_num, syncdeps) @@ -732,8 +766,14 @@ function _get_read_deps!(state::DataDepsState, dest_space::MemorySpace, ainfo::A end end end +# Whether the raw data backing `x` can be inspected on the current process to +# compute its aliasing. Distributed `Chunk`s are always inspectable (locally or +# via `remotecall`); the MPI extension overrides this for refs owned by a +# different rank, where the data is not present on this rank. +aliasing_available(@nospecialize(x)) = true + """ - gather_free_syncdeps!(state, space, remote_arg, write_num, chunk_to_ainfos, syncdeps) + gather_free_syncdeps!(state, space, key_ainfo, remote_arg, write_num, chunk_to_ainfos, syncdeps) Collect into `syncdeps` every task that must complete before the backing buffer `remote_arg` (a Datadeps-allocated copy in `space`) can be freed. @@ -743,9 +783,11 @@ arguments), its ainfos are in `chunk_to_ainfos` and we reuse their precomputed overlap sets. Otherwise the buffer only underlies wrapper arguments (e.g. it is the parent array shared by several `view`s, whose tracked slots are the views rather than this buffer); in that case we compute the buffer's own aliasing and -sync with every tracked ainfo that overlaps its memory. +sync with every tracked ainfo that overlaps its memory. When the buffer's data +is not inspectable here (`aliasing_available` is `false`, e.g. an `MPIRef` owned +by another rank), we fall back to the rank-uniform cache key ainfo `key_ainfo`. """ -function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, remote_arg, write_num::Int, chunk_to_ainfos, syncdeps) +function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, key_ainfo, remote_arg, write_num::Int, chunk_to_ainfos, syncdeps) ainfos = get(chunk_to_ainfos, remote_arg, nothing) if ainfos !== nothing for ainfo in ainfos @@ -754,6 +796,18 @@ function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, remote_ return end + # If the buffer's raw data isn't inspectable here (e.g. an `MPIRef` owned by + # another rank under uniform execution), we cannot compute its aliasing. + # Fall back to the rank-uniform cache key ainfo, which is metadata available + # identically on every rank. The cache stores raw ainfos, so wrap it to match + # the `AliasingWrapper` keys used by the overlap tracking. + if !aliasing_available(remote_arg) + wrapped = key_ainfo isa AliasingWrapper ? key_ainfo : AliasingWrapper(key_ainfo) + haskey(state.ainfos_overlaps, wrapped) && + get_write_deps!(state, space, wrapped, write_num, syncdeps) + return + end + # Buffer underlies wrapper arguments: find all tracked ainfos overlapping # it. We reuse the lookup's interval-tree overlap search (which prunes most # `will_alias` comparisons via bounding spans) rather than scanning every @@ -777,7 +831,7 @@ function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, remote_ end return end -function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::MemorySpace, ainfo::AbstractAliasing, task, write_num) +function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::MemorySpace, ainfo::AbstractAliasing, task, write_num; copy_src::Union{MemorySpace,Nothing}=nothing) state.ainfos_owner[ainfo] = task=>write_num empty!(state.ainfos_readers[ainfo]) @@ -793,6 +847,33 @@ function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::M push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num)) end + # Track which spaces hold a fully-current replica of this region + if copy_src === nothing + # Task write: only the written space is current, and other spaces' + # replicas of overlapping regions become stale + state.arg_current[arg_w] = Set{MemorySpace}((dest_space,)) + for other_arg_w in state.arg_overlaps[arg_w] + other_arg_w == arg_w && continue + other_current = get(state.arg_current, other_arg_w, nothing) + if other_current !== nothing + intersect!(other_current, (dest_space,)) + else + # Lazily-tracked args are current only at their origin + other_origin = state.arg_origin[other_arg_w.arg] + state.arg_current[other_arg_w] = other_origin == dest_space ? + Set{MemorySpace}((dest_space,)) : Set{MemorySpace}() + end + end + else + # Copy: the destination joins the current set; the source (and any + # other current replica) stays current, and since a copy only moves + # already-canonical bytes, overlapping arguments are unaffected + current = get!(state.arg_current, arg_w) do + Set{MemorySpace}((state.arg_origin[arg_w.arg],)) + end + push!(current, dest_space) + end + # Record the last place we were fully written to state.arg_owner[arg_w] = dest_space @@ -813,19 +894,30 @@ function generate_slot!(state::DataDepsState, dest_space, data) # because all we want here is to make a copy of some version of the data, # even if the data is not up to date. orig_space = memory_space(data) + check_uniform(orig_space) to_proc = first(processors(dest_space)) + check_uniform(to_proc) from_proc = first(processors(orig_space)) + check_uniform(from_proc) + check_uniform(typeof(data)) dest_space_args = get!(IdDict{Any,Any}, state.remote_args, dest_space) - aliased_object_cache = AliasedObjectCache(dest_space, state.ainfo_backing_chunk) + aliased_object_cache = AliasedObjectCache(current_acceleration(), dest_space, state.ainfo_backing_chunk) ctx = Sch.eager_context() id = rand(Int) @maybelog ctx timespan_start(ctx, :move, (;thunk_id=0, id, position=ArgPosition(), processor=to_proc), (;f=nothing, data)) - data_chunk = move_rewrap(aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + tid = something(DATADEPS_CURRENT_TASK[], (;uid=0)).uid + data_chunk = 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 @maybelog ctx timespan_finish(ctx, :move, (;thunk_id=0, id, position=ArgPosition(), processor=to_proc), (;f=nothing, data=data_chunk)) @assert memory_space(data_chunk) == dest_space "space mismatch! $dest_space (dest) != $(memory_space(data_chunk)) (actual) ($(typeof(data)) (data) vs. $(typeof(data_chunk)) (chunk)), spaces ($orig_space -> $dest_space)" dest_space_args[data] = data_chunk state.remote_arg_to_original[data_chunk] = data + check_uniform(memory_space(dest_space_args[data])) + check_uniform(processor(dest_space_args[data])) + check_uniform(dest_space_args[data].handle) + return dest_space_args[data] end function get_or_generate_slot!(state, dest_space, data) @@ -838,82 +930,93 @@ function get_or_generate_slot!(state, dest_space, data) end return state.remote_args[dest_space][data] end -function remotecall_endpoint(f, from_proc, to_proc, from_space, to_space, data) - to_w = root_worker_id(to_proc) - if to_w == myid() - data_converted = f(move(from_proc, to_proc, data)) - return tochunk(data_converted, to_proc) - end - return remotecall_fetch(to_w, from_proc, to_proc, to_space, data) do from_proc, to_proc, to_space, data - data_converted = f(move(from_proc, to_proc, data)) - return tochunk(data_converted, to_proc) - end + +function is_local(accel::DistributedAcceleration, target) + return root_worker_id(target) == myid() end -function rewrap_aliased_object!(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, x) - return aliased_object!(cache, x) do x - return remotecall_endpoint(identity, from_proc, to_proc, from_space, to_space, x) + +function remotecall_endpoint_toplevel(f, accel::DistributedAcceleration, cache::AliasedObjectCache, from_proc, to_proc, from_space, to_space, data::Chunk) + wid = root_worker_id(from_proc) + if wid == myid() + return f(accel, cache, from_proc, to_proc, from_space, to_space, unwrap(data))::Chunk end -end -function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, data::Chunk) - # Unwrap so that we hit the right dispatch - wid = root_worker_id(data) - if wid != myid() - return remotecall_fetch(move_rewrap, wid, cache, from_proc, to_proc, from_space, to_space, data) + return remotecall_fetch(wid, f, accel, cache, from_proc, to_proc, from_space, to_space, data) do f, accel, cache, from_proc, to_proc, from_space, to_space, data + return f(accel, cache, from_proc, to_proc, from_space, to_space, unwrap(data))::Chunk end - data_raw = unwrap(data) - return move_rewrap(cache, from_proc, to_proc, from_space, to_space, data_raw) end -function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, data) - # For generic data - return aliased_object!(cache, data) do data - return remotecall_endpoint(identity, from_proc, to_proc, from_space, to_space, data) +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() + return f(accel, from_proc, to_proc, from_space, to_space, data) end -end -function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, v::SubArray) - to_w = root_worker_id(to_proc) - p_chunk = rewrap_aliased_object!(cache, from_proc, to_proc, from_space, to_space, parent(v)) - inds = parentindices(v) - return remotecall_fetch(to_w, from_proc, to_proc, from_space, to_space, p_chunk, inds) do from_proc, to_proc, from_space, to_space, p_chunk, inds - p_new = move(from_proc, to_proc, p_chunk) - v_new = view(p_new, inds...) - return tochunk(v_new, to_proc) + return remotecall_fetch(wid, f, accel, from_proc, to_proc, from_space, to_space, data) do f, accel, from_proc, to_proc, from_space, to_space, data + return f(accel, from_proc, to_proc, from_space, to_space, data) end end -# FIXME: Do this programmatically via recursive dispatch + +#============================================================================== + move_rewrap header + children protocol + + Wrappers that share underlying storage (SubArray, triangular, ChunkView, + HaloArray, ...) register: + move_rewrap_parts(x) -> (children::Tuple, header) | nothing (leaf) + move_rewrap_build(T, children, header) -> reconstructed value + move_rewrap_child_types(T) -> Tuple of child types | nothing (leaf) + move_rewrap_header_mode(T) -> :none | :broadcast | :replicated + move_rewrap_result_type(T, child_chunktypes, header) -> result DataType + + Distributed and MPI accelerations share one generic move_rewrap that + transfers/shares children, carries or broadcasts the header, then rebuilds. +==============================================================================# + +# Leaf defaults +move_rewrap_parts(x) = nothing +move_rewrap_child_types(::Type) = nothing +move_rewrap_header_mode(::Type) = :none + +# SubArray: parent payload + indices header +move_rewrap_parts(x::SubArray) = ((parent(x),), parentindices(x)) +move_rewrap_build(::Type{<:SubArray}, (p,), inds) = view(p, inds...) +move_rewrap_child_types(::Type{T}) where {T<:SubArray} = (T.parameters[3],) +move_rewrap_header_mode(::Type{<:SubArray}) = :broadcast +move_rewrap_result_type(::Type{<:SubArray}, (PT,), inds) = + Base.promote_op(view, PT, typeof.(inds)...) + +# Triangular wrappers: parent payload, empty header for wrapper in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) - @eval function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, v::$(wrapper)) - to_w = root_worker_id(to_proc) - p_chunk = rewrap_aliased_object!(cache, from_proc, to_proc, from_space, to_space, parent(v)) - return remotecall_fetch(to_w, from_proc, to_proc, from_space, to_space, p_chunk) do from_proc, to_proc, from_space, to_space, p_chunk - p_new = move(from_proc, to_proc, p_chunk) - v_new = $(wrapper)(p_new) - return tochunk(v_new, to_proc) + @eval begin + move_rewrap_parts(x::$wrapper) = ((parent(x),), nothing) + move_rewrap_build(::Type{<:$wrapper}, (p,), ::Nothing) = $wrapper(p) + move_rewrap_child_types(::Type{T}) where {T<:$wrapper} = (T.parameters[2],) + move_rewrap_result_type(::Type{<:$wrapper}, (PT,), ::Nothing) = + $wrapper{eltype(PT),PT} + end +end + +# Nested Chunk: unwrap on the owner and recurse +move_rewrap(accel, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, data::Chunk) = + remotecall_endpoint_toplevel(move_rewrap, accel, cache, from_proc, to_proc, from_space, to_space, data) + +function move_rewrap(accel, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, data) + parts = move_rewrap_parts(data) + if parts === nothing + # 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) + end end end -end -function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, v::Base.RefValue) - return aliased_object!(cache, v) do v - return remotecall_endpoint(identity, from_proc, to_proc, from_space, to_space, v) + # Wrapper: recurse on children, then rebuild with header on the destination + children, header = parts + T = typeof(data) + child_chunks = map(c -> move_rewrap(accel, cache, from_proc, to_proc, from_space, to_space, c), children) + for cc in child_chunks + check_uniform(cc.handle) end -end -#= FIXME: Make this work so we can automatically move-rewrap recursive objects -function move_rewrap_recursive(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, x::T) where T - if isstructtype(T) - # Check all object fields (recursive) - for field in fieldnames(T) - value = getfield(x, field) - new_value = aliased_object!(cache, value) do value - return move_rewrap_recursive(cache, from_proc, to_proc, from_space, to_space, value) - end - setfield!(x, field, new_value) - end - return x - else - @warn "Cannot move-rewrap object of type $T" - return x + return remotecall_endpoint_transfer(accel, from_proc, to_proc, from_space, to_space, child_chunks) do accel, from_proc, to_proc, from_space, to_space, child_chunks + children_local = map(c -> move(from_proc, to_proc, c), child_chunks) + v_new = move_rewrap_build(T, children_local, header) + return tochunk(v_new, to_proc, to_space) end end -move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, x::String) = x # FIXME: Not necessarily true -move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, x::Symbol) = x -move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, x::Type) = x -=# diff --git a/src/datadeps/chunkview.jl b/src/datadeps/chunkview.jl index 1c2aa600f..903785f93 100644 --- a/src/datadeps/chunkview.jl +++ b/src/datadeps/chunkview.jl @@ -8,7 +8,9 @@ function _identity_hash(arg::ChunkView, h::UInt=UInt(0)) end function Base.view(c::Chunk, slices...) - if c.domain isa ArrayDomain + # N.B. Placeholder chunk records (e.g. MPI non-owner views of remote data) + # carry an empty domain; skip slice validation for them + if c.domain isa ArrayDomain && !(ndims(c.domain) == 0 && length(slices) > 0) nd, sz = ndims(c.domain), size(c.domain) nd == length(slices) || throw(DimensionMismatch("Expected $nd slices, got $(length(slices))")) @@ -29,32 +31,73 @@ function Base.view(c::Chunk, slices...) return ChunkView(c, slices) end +# Compose nested ChunkView slices onto the underlying Chunk (flatten; do not +# nest ChunkView structs). Mirrors Base.reindex for Int / AbstractRange / Colon. +_compose_chunkview_slices(::Tuple{}, ::Tuple{}) = () +_compose_chunkview_slices(::Tuple{}, ::Tuple{Any, Vararg}) = + throw(DimensionMismatch("Too many indices for ChunkView")) +_compose_chunkview_slices(::Tuple{Colon, Vararg}, ::Tuple{}) = + throw(DimensionMismatch("Too few indices for ChunkView")) +_compose_chunkview_slices(::Tuple{AbstractRange, Vararg}, ::Tuple{}) = + throw(DimensionMismatch("Too few indices for ChunkView")) +# Parent Int: dimension already dropped; keep and do not consume a sub-slice +_compose_chunkview_slices(ps::Tuple{Int, Vararg}, ss::Tuple) = + (ps[1], _compose_chunkview_slices(Base.tail(ps), ss)...) +# Parent Colon: pass the sub-slice through +_compose_chunkview_slices(ps::Tuple{Colon, Vararg}, ss::Tuple{Any, Vararg}) = + (ss[1], _compose_chunkview_slices(Base.tail(ps), Base.tail(ss))...) +# Parent AbstractRange: re-index into it +function _compose_chunkview_slices(ps::Tuple{AbstractRange, Vararg}, ss::Tuple{Any, Vararg}) + composed = try + ps[1][ss[1]] + catch e + e isa BoundsError || rethrow() + throw(ArgumentError("Slice $(ss[1]) out of bounds for parent slice $(ps[1])")) + end + return (composed, _compose_chunkview_slices(Base.tail(ps), Base.tail(ss))...) +end + +function Base.view(cv::ChunkView, slices...) + composed = _compose_chunkview_slices(cv.slices, slices) + return view(cv.chunk, composed...) +end + Base.view(c::DTask, slices...) = view(fetch(c; raw=true), slices...) -function aliasing(x::ChunkView{N}) where N +function aliasing(accel::Acceleration, x::ChunkView{N}, dep_mod) where N + @assert dep_mod === identity "Dependency modifiers not yet supported for ChunkView: $dep_mod" return remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices x = unwrap(x) v = view(x, slices...) - return aliasing(v) + return aliasing(accel, v, dep_mod) end end +aliasing(x::ChunkView) = aliasing(current_acceleration(), x, identity) +aliasing(x::ChunkView, dep_mod) = aliasing(current_acceleration(), x, dep_mod) memory_space(x::ChunkView) = memory_space(x.chunk) isremotehandle(x::ChunkView) = true -function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, slice::ChunkView) - to_w = root_worker_id(to_proc) - # N.B. We use move_rewrap (not rewrap_aliased_object!) so that if the inner - # chunk is a SubArray, it goes through the SubArray-aware path which shares - # the parent array via the aliased object cache. Using rewrap_aliased_object! - # would simply serialize the entire SubArray, creating a new parent copy on - # the destination, breaking aliasing with other views of the same parent. - p_chunk = move_rewrap(cache, from_proc, to_proc, from_space, to_space, slice.chunk) - return remotecall_fetch(to_w, from_proc, to_proc, from_space, to_space, p_chunk, slice.slices) do from_proc, to_proc, from_space, to_space, p_chunk, inds - p_new = move(from_proc, to_proc, p_chunk) - v_new = view(p_new, inds...) - return tochunk(v_new, to_proc) - end +# Under uniform (SPMD) execution, a ChunkView is replicated metadata on every +# rank (inner chunk record + slices); wrapping it in a rank-0-owned Chunk +# would strand the other ranks' inner chunk records, so keep it raw +datadeps_arg_wrap(arg::ChunkView) = uniform_execution() ? arg : tochunk(arg) + +# ChunkView is a remote handle: its move_rewrap accesses the underlying chunk +# remotely itself, so dispatch directly from the caller +function remotecall_endpoint_toplevel(f, accel::DistributedAcceleration, cache::AliasedObjectCache, from_proc, to_proc, from_space, to_space, data::ChunkView) + return f(accel, cache, from_proc, to_proc, from_space, to_space, data)::Chunk end + +# Header+children registration: share the backing chunk, rebuild as a view. +# N.B. We use move_rewrap (not rewrap_aliased_object!) so that if the inner +# chunk is a SubArray, it goes through the SubArray-aware path which shares +# the parent array via the aliased object cache. +move_rewrap_parts(slice::ChunkView) = ((slice.chunk,), slice.slices) +move_rewrap_build(::Type{<:ChunkView}, (p,), inds) = view(p, inds...) +move_rewrap_header_mode(::Type{<:ChunkView}) = :replicated +move_rewrap_result_type(::Type{<:ChunkView}, (PT,), inds) = + Base.promote_op(view, PT, typeof.(inds)...) + function move(from_proc::Processor, to_proc::Processor, slice::ChunkView) to_w = root_worker_id(to_proc) return remotecall_fetch(to_w, from_proc, to_proc, slice.chunk, slice.slices) do from_proc, to_proc, chunk, slices @@ -64,4 +107,4 @@ function move(from_proc::Processor, to_proc::Processor, slice::ChunkView) end end -Base.fetch(slice::ChunkView) = view(fetch(slice.chunk), slice.slices...) \ No newline at end of file +Base.fetch(slice::ChunkView) = view(fetch(slice.chunk), slice.slices...) diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 13d0c6b63..9441f3bbe 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -25,6 +25,13 @@ function enqueue!(queue::DataDepsTaskQueue, pairs::Vector{DTaskPair}) append!(queue.seen_tasks, pairs) end +const DATADEPS_CURRENT_TASK = TaskLocalValue{Union{DTask,Nothing}}(Returns(nothing)) + +# Tag for datadeps-internal tasks (copies, frees) launched outside the user +# task queue. Under uniform (MPI) execution every task needs a unique, +# rank-uniform tag for its P2P transfers; under Distributed this is unused. +datadeps_task_tag() = uniform_execution() ? UInt32(to_tag()) : nothing + """ spawn_datadeps(f::Base.Callable) @@ -88,6 +95,15 @@ end const DATADEPS_SCHEDULER = ScopedValue{Union{DataDepsScheduler,Nothing}}(nothing) const DATADEPS_LAUNCH_WAIT = ScopedValue{Union{Bool,Nothing}}(nothing) +# Current task uid, propagated into `tochunk` so uniform-execution backends +# (MPIExt) can derive deterministic, rank-agreed handle IDs. Core datadeps sets +# this during planning/execution; MPIExt reads it. 0 means "no task in scope". +const DATADEPS_THUNK_ID = ScopedValue{Int64}(0) + +# Deterministic, rank-agreed task tag for uniform execution. Only invoked when +# `uniform_execution()` holds (i.e. under MPIExt), which provides the method. +function to_tag end + function distribute_tasks!(queue::DataDepsTaskQueue) #= TODO: Improvements to be made: # - Support for copying non-AbstractArray arguments @@ -98,20 +114,24 @@ function distribute_tasks!(queue::DataDepsTaskQueue) =# # Get the set of all processors to be scheduled on - all_procs = Processor[] - scope = get_compute_scope() - for w in procs() - append!(all_procs, get_processors(OSProc(w))) + accel = current_acceleration() + accel_procs = filter(procs(Dagger.Sch.eager_context())) do proc + Dagger.accel_matches_proc(accel, proc) end + all_procs = unique(vcat([collect(Dagger.get_processors(gp)) for gp in accel_procs]...)) + select_processors_uniform!(all_procs, accel) + scope = get_compute_scope() filter!(proc->proc_in_scope(proc, scope), all_procs) if isempty(all_procs) throw(Sch.SchedulingException("No processors available, try widening scope")) end + if uniform_execution(accel) + for proc in all_procs + check_uniform(proc) + end + end all_scope = UnionScope(map(ExactScope, all_procs)) exec_spaces = unique(vcat(map(proc->collect(memory_spaces(proc)), all_procs)...)) - if !all(space->space isa CPURAMMemorySpace, exec_spaces) && !all(space->root_worker_id(space) == myid(), exec_spaces) - @warn "Datadeps support for multi-GPU, multi-worker is currently broken\nPlease be prepared for incorrect results or errors" maxlog=1 - end # Round-robin assign tasks to processors upper_queue = get_options(:task_queue) @@ -128,10 +148,21 @@ function distribute_tasks!(queue::DataDepsTaskQueue) # Copy args from remote to local # N.B. We sort the keys to ensure a deterministic order for uniformity + check_uniform(length(state.arg_owner)) for arg_w in sort(collect(keys(state.arg_owner)); by=arg_w->arg_w.hash) + check_uniform(arg_w) arg = arg_w.arg origin_space = state.arg_origin[arg] - remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) + # When the origin still holds a fully-current replica (the argument was + # only read, or copies merely propagated it), the write-back is elided. + # This is only safe here at region end: mid-region, the copy tasks also + # serialize readers against later writers, so they must not be skipped. + current = get(state.arg_current, arg_w, nothing) + if current !== nothing && origin_space in current + remainder = NoAliasing() + else + remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) + end if remainder isa MultiRemainderAliasing origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) enqueue_remainder_copy_from!(state, origin_space, arg_w, remainder, origin_scope, write_num) @@ -158,6 +189,13 @@ function distribute_tasks!(queue::DataDepsTaskQueue) # regardless of array wrappers. A buffer is a Datadeps-allocated copy (safe # to free) at every space except the `key`'s source space, where it is the # user's original data (`is_original`). + # + # N.B. Do not mpi_cleanup_tid the planning-time uid keys here. Eager uid == + # Sch thunk id, so planning (DATADEPS_THUNK_ID=>uid) and later in-task tochunk share + # the same next_ref_sub_id! counter. Resetting it before execution would + # reissue (tid, sub_id) pairs while planning-time MPIRefs are still live, + # colliding ArgumentWrapper / Chunk identity hashes. Counters are reclaimed + # in wait_all after the region finishes (thunks generate no further IDs). obj_cache = unwrap(state.ainfo_backing_chunk) # Map each tracked slot chunk to its ainfos, to compute free syncdeps. chunk_to_ainfos = IdDict{Any,Vector{AliasingWrapper}}() @@ -176,8 +214,9 @@ function distribute_tasks!(queue::DataDepsTaskQueue) haskey(freed, remote_arg) && continue freed[remote_arg] = nothing free_syncdeps = Set{ThunkSyncdep}() - gather_free_syncdeps!(state, remote_space, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) - Dagger.@spawn scope=free_scope syncdeps=free_syncdeps Dagger.unsafe_free!(remote_arg) + gather_free_syncdeps!(state, remote_space, ainfo, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) + # `tag` keeps the free task rank-uniform under MPI/uniform execution. + Dagger.@spawn scope=free_scope syncdeps=free_syncdeps tag=datadeps_task_tag() Dagger.unsafe_free!(remote_arg) end end end @@ -215,11 +254,15 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr fargs::Vector{Argument} end + DATADEPS_CURRENT_TASK[] = task + task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) scheduler = queue.scheduler our_proc = datadeps_schedule_task(scheduler, state, all_procs, all_scope, task_scope, spec, task) @assert our_proc in all_procs our_space = only(memory_spaces(our_proc)) + check_uniform(our_proc) + check_uniform(our_space) # Find the scope for this task (and its copies) task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) @@ -238,15 +281,16 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr throw(Sch.SchedulingException("Scopes are not compatible: $(our_scope.x), $(our_scope.y)")) end - f = spec.fargs[1] tid = task.uid - # FIXME: May not be correct to move this under uniformity - #f.value = move(default_processor(), our_proc, value(f)) + # N.B. `with_value` returns a fresh argument rather than mutating; this is + # required for typed specs, whose `TypedArgument`s disallow `setproperty!`. + f = with_value(spec.fargs[1], move(default_processor(), our_proc, value(spec.fargs[1]))) @dagdebug tid :spawn_datadeps "($(repr(value(f)))) Scheduling: $our_proc ($our_space)" # Copy raw task arguments for analysis - # N.B. Used later for checking dependencies - task_args = map_or_ntuple(idx->copy(spec.fargs[idx]), spec.fargs) + # N.B. Used later for checking dependencies. The moved function argument + # (`f`) replaces the original at position 1. + task_args = map_or_ntuple(idx->copy(idx == 1 ? f : spec.fargs[idx]), spec.fargs) # Populate all task dependencies task_arg_ws = populate_task_info!(state, task_args, spec, task) @@ -324,6 +368,9 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr if spec.options.syncdeps === nothing spec.options.syncdeps = Set{ThunkSyncdep}() end + if spec.options.tag === nothing && uniform_execution() + spec.options.tag = to_tag() + end syncdeps = spec.options.syncdeps map_or_ntuple(task_arg_ws) do idx arg_ws = task_arg_ws[idx] @@ -358,7 +405,9 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr new_spec = DTaskSpec(new_fargs, spec.options) new_spec.options.scope = our_scope new_spec.options.exec_scope = our_scope - new_spec.options.occupancy = Dict(Any=>0) + if uniform_execution() + new_spec.options.occupancy = Dict(Any=>0) + end ctx = Sch.eager_context() @maybelog ctx timespan_start(ctx, :datadeps_execute, (;thunk_id=task.uid), (;)) enqueue!(queue.upper_queue, DTaskPair(new_spec, task)) @@ -386,5 +435,7 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr write_num += 1 + DATADEPS_CURRENT_TASK[] = nothing + return write_num end diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 33ca8c419..1000471b2 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -98,7 +98,7 @@ function compute_remainder_for_arg!(state::DataDepsState, for entry in state.arg_history[arg_w] push!(spaces_set, entry.space) end - spaces = collect(spaces_set) + spaces = sort(collect(spaces_set), by=short_name) N = length(spaces) # Lookup all memory spans for arg_w in these spaces @@ -118,6 +118,8 @@ function compute_remainder_for_arg!(state::DataDepsState, @goto restart end end + check_uniform(spaces) + check_uniform(target_ainfos) # We may only need to schedule a full copy from the origin space to the # target space if this is the first time we've written to `arg_w` @@ -159,6 +161,8 @@ function compute_remainder_for_arg!(state::DataDepsState, other_ainfo = aliasing!(state, owner_space, arg_w) other_space = owner_space end + check_uniform(other_ainfo) + check_uniform(other_space) # Lookup all memory spans for arg_w in these spaces other_remote_arg_w = first(collect(state.ainfo_arg[other_ainfo])) @@ -174,6 +178,7 @@ function compute_remainder_for_arg!(state::DataDepsState, foreach(other_many_spans) do span verify_span(span) end + check_uniform(other_many_spans) if other_space == target_space # Only subtract, this data is already up-to-date in target_space @@ -250,7 +255,9 @@ Enqueues a copy operation to update the remainder regions of an object before a function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::MultiRemainderAliasing, f, idx, dest_scope, task, write_num::Int) for remainder in remainder_aliasing.remainders + check_uniform(remainder.space) @assert !isempty(remainder.spans) + check_uniform(remainder.spans) enqueue_remainder_copy_to!(state, dest_space, arg_w, remainder, f, idx, dest_scope, task, write_num) end end @@ -286,14 +293,14 @@ function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpac ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true Dagger.move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true tag=datadeps_task_tag() Dagger.move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the sources and writes to the target for ainfo in source_ainfos add_reader!(state, arg_w, source_space, ainfo, copy_task, write_num) end - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) end """ enqueue_remainder_copy_from!(state::DataDepsState, target_ainfo::AliasingWrapper, arg, remainder_aliasing, @@ -304,7 +311,9 @@ Enqueues a copy operation to update the remainder regions of an object back to t function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::MultiRemainderAliasing, dest_scope, write_num::Int) for remainder in remainder_aliasing.remainders + check_uniform(remainder.space) @assert !isempty(remainder.spans) + check_uniform(remainder.spans) enqueue_remainder_copy_from!(state, dest_space, arg_w, remainder, dest_scope, write_num) end end @@ -340,14 +349,14 @@ function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySp ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true Dagger.move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true tag=datadeps_task_tag() Dagger.move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the sources and writes to the target for ainfo in source_ainfos add_reader!(state, arg_w, source_space, ainfo, copy_task, write_num) end - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) end # FIXME: Document me @@ -376,12 +385,12 @@ function enqueue_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w:: ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true Dagger.move!(dep_mod, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true tag=datadeps_task_tag() Dagger.move!(dep_mod, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the source and writes to the target add_reader!(state, arg_w, source_space, source_ainfo, copy_task, write_num) - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) end function enqueue_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, dest_scope, write_num::Int) @@ -408,44 +417,50 @@ function enqueue_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true Dagger.move!(dep_mod, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true tag=datadeps_task_tag() Dagger.move!(dep_mod, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the source and writes to the target add_reader!(state, arg_w, source_space, source_ainfo, copy_task, write_num) - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) end # Main copy function for RemainderAliasing function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space::MemorySpace, to::Chunk, from::Chunk) where S - # TODO: Support direct copy between GPU memory spaces + # Same-worker device-to-device: copy spans directly with one KA launch + if root_worker_id(from_space) == myid() && root_worker_id(to_space) == myid() && + is_device_space(from_space) && is_device_space(to_space) + from_s = storage_array(unwrap(from)) + to_s = storage_array(unwrap(to)) + with_context!(from_space) + with_context!(to_space) + multi_span_copy!(to_s, from_s, dep_mod.spans) + return + end - # Copy the data from the source object + # Gather source spans into a host buffer (device arrays pack via KA first). + # Pin when the source is device memory so DtoH hits page-locked host RAM. copies = remotecall_fetch(root_worker_id(from_space), from_space, dep_mod, from) do from_space, dep_mod, from len = sum(span_tuple->span_len(span_tuple[1]), dep_mod.spans) copies = Vector{UInt8}(undef, len) + pin_buffer!(gpu_memory_kind(from_space), copies) from_raw = unwrap(from) - offset = UInt64(1) with_context!(from_space) GC.@preserve copies begin - for (from_span, _) in dep_mod.spans - read_remainder!(copies, offset, from_raw, from_span.ptr, from_span.len) - offset += from_span.len - end + # Lazily view the source (first) span of each pair — avoid + # materializing a whole second span vector for large span sets + multi_span_gather!(copies, from_raw, Iterators.map(first, dep_mod.spans)) end - @assert offset == len+UInt64(1) return copies end - # Copy the data into the destination object - offset = UInt64(1) + # Scatter into the destination spans (pin again after remotecall deserialize) to_raw = unwrap(to) + with_context!(to_space) + pin_buffer!(gpu_memory_kind(to_space), copies) GC.@preserve copies begin - for (_, to_span) in dep_mod.spans - write_remainder!(copies, offset, to_raw, to_span.ptr, to_span.len) - offset += to_span.len - end - @assert offset == length(copies)+UInt64(1) + # Lazily view the dest (last) span of each pair (see gather above) + multi_span_scatter!(to_raw, copies, Iterators.map(last, dep_mod.spans)) end # Ensure that the data is visible @@ -469,6 +484,8 @@ function read_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, from::Arr unsafe_copyto!(Ptr{eltype(from)}(pointer(copies, copies_offset)), pointer(from_vec, from_offset_n), n) end function read_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, from::DenseArray, from_ptr::UInt64, len::UInt64) + # Device arrays (e.g. CuArray) need their VRAM context for pointer/copyto! + with_context!(memory_space(from)) elsize = sizeof(eltype(from)) @assert len / elsize == round(UInt64, len / elsize) "Span length is not an integer multiple of the element size: $(len) / $(elsize) = $(len / elsize) (elsize: $elsize)" n = UInt64(len / elsize) @@ -492,6 +509,8 @@ function write_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, to::Arra unsafe_copyto!(pointer(to_vec, to_offset_n), Ptr{eltype(to)}(pointer(copies, copies_offset)), n) end function write_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, to::DenseArray, to_ptr::UInt64, len::UInt64) + # Device arrays (e.g. CuArray) need their VRAM context for pointer/copyto! + with_context!(memory_space(to)) elsize = sizeof(eltype(to)) @assert len / elsize == round(UInt64, len / elsize) "Span length is not an integer multiple of the element size: $(len) / $(elsize) = $(len / elsize) (elsize: $elsize)" n = UInt64(len / elsize) @@ -542,4 +561,4 @@ function find_object_holding_ptr(A::SparseMatrixCSC, ptr::UInt64) span = LocalMemorySpan(pointer(A.rowval), length(A.rowval)*sizeof(eltype(A.rowval))) @assert span_start(span) <= ptr <= span_end(span) "Pointer $ptr not found in SparseMatrixCSC" return A.rowval -end \ No newline at end of file +end diff --git a/src/dtask.jl b/src/dtask.jl index fc78ee67b..3bdefb543 100644 --- a/src/dtask.jl +++ b/src/dtask.jl @@ -24,14 +24,14 @@ Base.wait(t::ThunkFuture) = Dagger.Sch.thunk_yield() do wait(t.future) return end -function Base.fetch(t::ThunkFuture; proc=OSProc(), raw=false) +function Base.fetch(t::ThunkFuture; proc=OSProc(), raw=false, move_value=!raw, unwrap=!raw, uniform=uniform_execution()) error, value = Dagger.Sch.thunk_yield() do fetch(t.future) end if error throw(value) end - if raw + if !move_value return value else return move(proc, value) @@ -78,11 +78,11 @@ function Base.wait(t::DTask) wait(t.future) return end -function Base.fetch(t::DTask; raw=false) +function Base.fetch(t::DTask; raw=false, move_value=!raw, unwrap=!raw, uniform=false) if !istaskstarted(t) throw(ConcurrencyViolationError("Cannot `fetch` an unlaunched `DTask`")) end - return fetch(t.future; raw) + return fetch(t.future; move_value, unwrap, uniform) end function waitany(tasks::Vector{DTask}) if isempty(tasks) diff --git a/src/gpu.jl b/src/gpu.jl index fac44893d..0100bfbd2 100644 --- a/src/gpu.jl +++ b/src/gpu.jl @@ -105,6 +105,118 @@ gpu_synchronize(::Val{:CPU}) = nothing with_context!(proc::Processor) = nothing with_context!(space::MemorySpace) = nothing +# Backend kind for an array or memory space (`:CPU`, `:CUDA`, `:ROC`, …). +# GPU extensions override for their array / VRAM space types. +gpu_memory_kind(x) = :CPU +gpu_memory_kind(::MemorySpace) = :CPU + +# Page-lock a host buffer for faster DtoH/HtoD (GPU extensions override per Val). +# Unpinning happens via the buffer's GC finalizer. Used by MPI and Distributed. +pin_buffer!(kind::Symbol, buf) = pin_buffer!(Val(kind), buf) +pin_buffer!(::Val, buf) = nothing + +### Host staging buffer pools (per GPU backend kind) +# Staging device payloads through fresh pageable allocations is slow and +# GC-heavy; transfers reuse pooled host buffers that GPU extensions +# page-lock (pinned memory roughly doubles DtoH/HtoD bandwidth). +# Pools are keyed by `gpu_memory_kind` (`:CPU`/`:CUDA`/`:ROC`, …) so +# differently pinned buffers never share a free-list. +const STAGE_POOLS = LockedObject(Dict{Symbol, Dict{Int,Vector{Vector{UInt8}}}}()) +const STAGE_POOL_MAX_BYTES = Ref(256 * 1024 * 1024) +const STAGE_POOL_BYTES = Threads.Atomic{Int}(0) + +function stage_acquire!(kind::Symbol, nbytes::Integer) + cls = nextpow(2, max(Int(nbytes), 4096)) + buf = lock(STAGE_POOLS) do pools + pool = get!(Dict{Int,Vector{Vector{UInt8}}}, pools, kind) + bufs = get(pool, cls, nothing) + (bufs === nothing || isempty(bufs)) ? nothing : pop!(bufs) + end + if buf === nothing + buf = Vector{UInt8}(undef, cls) + pin_buffer!(kind, buf) + else + Threads.atomic_sub!(STAGE_POOL_BYTES, cls) + end + return buf +end +function stage_release!(kind::Symbol, buf::Vector{UInt8}) + cls = length(buf) + if STAGE_POOL_BYTES[] + cls > STAGE_POOL_MAX_BYTES[] + return # drop it; GC unpins via the registration finalizer + end + Threads.atomic_add!(STAGE_POOL_BYTES, cls) + lock(STAGE_POOLS) do pools + pool = get!(Dict{Int,Vector{Vector{UInt8}}}, pools, kind) + push!(get!(Vector{Vector{UInt8}}, pool, cls), buf) + end + return +end + +# Copy a dense device payload into a pooled host buffer; returns +# (host_view, buf, kind) — keep `buf` rooted while using the view, then release +function stage_to_host!(value::DenseArray{T}) where T + kind = gpu_memory_kind(value) + buf = stage_acquire!(kind, sizeof(value)) + host = unsafe_wrap(Array, Ptr{T}(pointer(buf)), size(value)) + with_context!(memory_space(value)) + copyto!(host, value) + return host, buf, kind +end + +# Owned host Array for DtoH (Distributed remotecall / return paths): pin then copy +function pinned_host_array(value::DenseArray{T}) where T + kind = gpu_memory_kind(value) + host = Array{T}(undef, size(value)) + pin_buffer!(kind, host) + with_context!(memory_space(value)) + copyto!(host, value) + return host +end + +### Same-node device IPC +# When both endpoints are device memory on the same node, the payload can stay +# on-device: the sender stages into an IPC-exportable allocation and ships only +# a small handle; the receiver maps it and copies device-to-device. +# GPU extensions override eligibility and export/import. Transport of the +# handle (MPI P2P vs Distributed remotecall) is acceleration-specific. +# DAGGER_IPC=0 disables the path. +ipc_eligible(from_inner::MemorySpace, to_inner::MemorySpace) = false +# Transfers smaller than this stay on the staged path (handle exchange and +# ack latency dominate below the crossover, ~128KiB on PCIe-attached GPUs) +const IPC_MIN_BYTES = Ref{Int}(128 * 1024) +# Export `value`: returns `(info, token)` where `info` is a small serializable +# description and `token` keeps the staging allocation alive until release +ipc_export(value) = error("No IPC export implementation for $(typeof(value))") +ipc_release!(token) = nothing +# Copy the exported data into `dest` (receiver side) +ipc_copyto!(dest, info) = error("No IPC import implementation for $(typeof(dest))") +# Materialize a fresh device array from the exported data (receiver side) +ipc_materialize(info) = error("No IPC materialize implementation for $(typeof(info))") + +# MPI/GPU interop hooks. GPU package extensions (which cannot depend on MPI) +# specialize these per array/space type, and MPIExt refines `mpi_library_gpu_aware`. +# Declared in core so both extension families extend the same Dagger generics +# (and so GPU extensions load even when MPIExt isn't present). +# +# Can `value` be handed to MPI directly as a device buffer? (needs GPU-aware MPI) +mpi_device_direct(value) = false +# Make pending device work on `value`/`space` visible to the host/MPI +mpi_device_sync(value) = nothing +mpi_device_sync(space::MemorySpace) = nothing +# Remap rank-local memory space stamps to the owning rank so spans from +# different ranks never falsely alias. GPU extensions stamp their VRAM spaces. +mpi_remap_space(space::CPURAMMemorySpace, owner::Int) = CPURAMMemorySpace(owner) +mpi_remap_space(space::MemorySpace, owner::Int) = + throw(ArgumentError("mpi_remap_space not defined for $(typeof(space)); backends must stamp the owning MPI rank")) +# Whether the loaded MPI library supports device buffers of `kind`. Only invoked +# under MPI acceleration; MPIExt provides the method (queries the MPI library). +function mpi_library_gpu_aware end + +# Distributed: same node iff workers share a system_uuid (CUDAExt IPC gate) +same_node(::DistributedAcceleration, w1::Integer, w2::Integer) = + system_uuid(Int(w1)) == system_uuid(Int(w2)) + # Adapt RefValue mutable struct GPURef{T,S<:MemorySpace} <: Ref{T} value::T diff --git a/src/lib/domain-blocks.jl b/src/lib/domain-blocks.jl index 2a0854e3b..ae79797cf 100644 --- a/src/lib/domain-blocks.jl +++ b/src/lib/domain-blocks.jl @@ -6,6 +6,8 @@ struct DomainBlocks{N} <: AbstractArray{ArrayDomain{N, NTuple{N, UnitRange{Int}} end Base.@deprecate_binding BlockedDomains DomainBlocks +Base.ndims(::DomainBlocks{N}) where N = N + size(x::DomainBlocks) = map(length, x.cumlength) function _getindex(x::DomainBlocks{N}, idx::Tuple) where N starts = map((vec, i) -> i == 0 ? 0 : getindex(vec,i), x.cumlength, map(x->x-1, idx)) diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index eb4f7ad5b..c78b96d83 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -4,23 +4,26 @@ end CPURAMMemorySpace() = CPURAMMemorySpace(myid()) root_worker_id(space::CPURAMMemorySpace) = space.owner -memory_space(x) = CPURAMMemorySpace(myid()) -function memory_space(x::Chunk) - proc = processor(x) - if proc isa OSProc - # TODO: This should probably be programmable - return CPURAMMemorySpace(proc.pid) - else - return only(memory_spaces(proc)) - end -end -memory_space(x::DTask) = - memory_space(fetch(x; raw=true)) +# Stable textual key for deterministic space ordering +short_name(space::MemorySpace) = string(space) +short_name(space::CPURAMMemorySpace) = "CPU: $(space.owner)" + +memory_space(x, proc::Processor=default_processor()) = first(memory_spaces(proc)) + +# Acceleration-free memory space of a raw value, used to label chunk records; +# GPU package extensions add methods for device array types (e.g. +# CuArray -> CUDAVRAMMemorySpace) +value_memory_space(x) = CPURAMMemorySpace(myid()) +memory_space(x::Processor) = first(memory_spaces(x)) +memory_space(x::Chunk) = x.space +memory_space(x::DTask) = memory_space(fetch(x; move_value=false, unwrap=false)) memory_spaces(::P) where {P<:Processor} = throw(ArgumentError("Must define `memory_spaces` for `$P`")) memory_spaces(proc::ThreadProc) = Set([CPURAMMemorySpace(proc.owner)]) +memory_spaces(proc::OSProc) = + Set([CPURAMMemorySpace(proc.pid)]) processors(::S) where {S<:MemorySpace} = throw(ArgumentError("Must define `processors` for `$S`")) processors(space::CPURAMMemorySpace) = @@ -28,9 +31,10 @@ processors(space::CPURAMMemorySpace) = ### In-place Data Movement -function unwrap(x::Chunk) - @assert x.handle.owner == myid() - MemPool.poolget(x.handle) +unwrap(x::Chunk) = unwrap(x.handle) +function unwrap(handle::DRef) + @assert root_worker_id(handle) == myid() "DRef $handle is not owned by this process: $(root_worker_id(handle)) != $(myid())" + return MemPool.poolget(handle) end move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::T, from::F) where {T,F} = throw(ArgumentError("No `move!` implementation defined for $F -> $T")) @@ -49,18 +53,24 @@ function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::Base to[] = from[] return end -function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} - move!(to_space, from_space, dep_mod(to), dep_mod(from)) -end -function move!(to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} - copyto!(to, from) +function move!(dep_mod::Symbol, to_space::MemorySpace, from_space::MemorySpace, to::Base.RefValue{T}, from::Base.RefValue{T}) where {T} + to_f = getfield(to, dep_mod) + from_f = getfield(from, dep_mod) + if to_f isa AbstractArray && from_f isa AbstractArray && axes(to_f) == axes(from_f) + # Update the field's storage in place: aliasing spans computed when the + # destination slot was created point into that storage, and rebinding + # the field would silently disconnect them + move!(to_space, from_space, to_f, from_f) + else + setfield!(to, dep_mod, from_f) + end return end - -function move!(::Type{<:Diagonal}, to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} - copyto!(view(to, diagind(to)), view(from, diagind(from))) - return +function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} + move!(to_space, from_space, dep_mod(to), dep_mod(from)) end +# AbstractArray and Diagonal move! are defined in utils/span_copy.jl (after +# multi-span helpers) so GPU views can share one KA path without method overwrite. # FIXME: Bidiagonal (need direction specified in type) function move!(::Type{<:Tridiagonal}, to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} copyto!(view(to, diagind(to, -1)), view(from, diagind(from, -1))) @@ -69,6 +79,16 @@ function move!(::Type{<:Tridiagonal}, to_space::MemorySpace, from_space::MemoryS return end +# FIXME: Take MemorySpace instead +function move_type(from_proc::Processor, to_proc::Processor, ::Type{T}) where T + if from_proc == to_proc + return T + end + return Base._return_type(move, Tuple{typeof(from_proc), typeof(to_proc), T}) +end +move_type(from_proc::Processor, to_proc::Processor, ::Type{<:Chunk{T}}) where T = + move_type(from_proc, to_proc, T) + ### Aliasing and Memory Spans type_may_alias(::Type{String}) = false @@ -355,6 +375,7 @@ function memory_spans(oa::ObjectAliasing{S}) where S return [span] end +aliasing(accel::Acceleration, x, T) = aliasing(x, T) function aliasing(x, dep_mod) if dep_mod isa Symbol return aliasing(getfield(x, dep_mod)) @@ -391,19 +412,25 @@ aliasing(::String) = NoAliasing() # FIXME: Not necessarily true aliasing(::Symbol) = NoAliasing() aliasing(::Type) = NoAliasing() function aliasing(x::Chunk, T) - @assert x.handle isa DRef if root_worker_id(x.processor) == myid() return aliasing(unwrap(x), T) end + @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T aliasing(unwrap(x), T) end end -aliasing(x::Chunk) = remotecall_fetch(root_worker_id(x.processor), x) do x - aliasing(unwrap(x)) +function aliasing(x::Chunk) + if root_worker_id(x.processor) == myid() + return aliasing(unwrap(x)) + end + @assert x.handle isa DRef + return remotecall_fetch(root_worker_id(x.processor), x) do x + aliasing(unwrap(x)) + end end -aliasing(x::DTask, T) = aliasing(fetch(x; raw=true), T) -aliasing(x::DTask) = aliasing(fetch(x; raw=true)) +aliasing(x::DTask, T) = aliasing(fetch(x; move_value=false, unwrap=false), T) +aliasing(x::DTask) = aliasing(fetch(x; move_value=false, unwrap=false)) function aliasing(x::Base.RefValue{T}) where T addr = UInt(Base.pointer_from_objref(x) + fieldoffset(typeof(x), 1)) @@ -539,14 +566,26 @@ function memory_spans(a::TriangularAliasing{T,S}) where {T,S} end return spans end -aliasing(x::UpperTriangular{T}) where T = - TriangularAliasing{T,CPURAMMemorySpace}(pointer(parent(x)), size(parent(x), 1), true, true) -aliasing(x::LowerTriangular{T}) where T = - TriangularAliasing{T,CPURAMMemorySpace}(pointer(parent(x)), size(parent(x), 1), false, true) -aliasing(x::UnitUpperTriangular{T}) where T = - TriangularAliasing{T,CPURAMMemorySpace}(pointer(parent(x)), size(parent(x), 1), true, false) -aliasing(x::UnitLowerTriangular{T}) where T = - TriangularAliasing{T,CPURAMMemorySpace}(pointer(parent(x)), size(parent(x), 1), false, false) +function aliasing(x::UpperTriangular{T}) where T + p = parent(x) + space = memory_space(p) + TriangularAliasing{T,typeof(space)}(RemotePtr{Cvoid}(UInt(pointer(p)), space), size(p, 1), true, true) +end +function aliasing(x::LowerTriangular{T}) where T + p = parent(x) + space = memory_space(p) + TriangularAliasing{T,typeof(space)}(RemotePtr{Cvoid}(UInt(pointer(p)), space), size(p, 1), false, true) +end +function aliasing(x::UnitUpperTriangular{T}) where T + p = parent(x) + space = memory_space(p) + TriangularAliasing{T,typeof(space)}(RemotePtr{Cvoid}(UInt(pointer(p)), space), size(p, 1), true, false) +end +function aliasing(x::UnitLowerTriangular{T}) where T + p = parent(x) + space = memory_space(p) + TriangularAliasing{T,typeof(space)}(RemotePtr{Cvoid}(UInt(pointer(p)), space), size(p, 1), false, false) +end struct DiagonalAliasing{T,S} <: AbstractAliasing ptr::RemotePtr{Cvoid,S} @@ -611,5 +650,5 @@ unsafe_free!(x::Chunk) = remotecall_fetch(root_worker_id(x), x) do x unsafe_free!(unwrap(x)) return end -unsafe_free!(x::DTask) = unsafe_free!(fetch(x; raw=true)) +unsafe_free!(x::DTask) = unsafe_free!(fetch(x; move_value=false, unwrap=false)) unsafe_free!(x) = nothing # Do nothing by default diff --git a/src/mutable.jl b/src/mutable.jl new file mode 100644 index 000000000..1f48ead53 --- /dev/null +++ b/src/mutable.jl @@ -0,0 +1,41 @@ +function _mutable_inner(@nospecialize(f), proc, scope) + result = f() + return Ref(Dagger.tochunk(result, proc, scope)) +end + +""" + mutable(f::Base.Callable; worker, processor, scope) -> Chunk + +Calls `f()` on the specified worker or processor, returning a `Chunk` +referencing the result with the specified scope `scope`. +""" +function mutable(@nospecialize(f); worker=nothing, processor=nothing, scope=nothing) + if processor === nothing + if worker === nothing + processor = OSProc() + else + processor = OSProc(worker) + end + else + @assert worker === nothing "mutable: Can't mix worker and processor" + end + if scope === nothing + scope = processor isa OSProc ? ProcessScope(processor) : ExactScope(processor) + end + return fetch(Dagger.@spawn scope=scope _mutable_inner(f, processor, scope))[] +end + +""" + @mutable [worker=1] [processor=OSProc()] [scope=ProcessorScope()] f() + +Helper macro for [`mutable()`](@ref). +""" +macro mutable(exs...) + opts = esc.(exs[1:end-1]) + ex = exs[end] + quote + let f = @noinline ()->$(esc(ex)) + $mutable(f; $(opts...)) + end + end +end diff --git a/src/options.jl b/src/options.jl index 89a27a214..aabd20002 100644 --- a/src/options.jl +++ b/src/options.jl @@ -26,10 +26,13 @@ Stores per-task options to be passed to the scheduler. - `storage_leaf_tag::Union{MemPool.Tag,Nothing}=nothing`: If not `nothing`, specifies the MemPool storage leaf tag to associate with the task's result. This tag can be used by MemPool's storage devices to manipulate their behavior, such as the file name used to store data on disk." - `storage_retain::Union{Bool,Nothing}=nothing`: The value of `retain` to pass to `MemPool.poolset` when constructing the result `Chunk`. `nothing` defaults to `false`. - `name::Union{String,Nothing}=nothing`: If not `nothing`, annotates the task with a name for logging purposes. +- `tag::Union{UInt32,Nothing}=nothing`: (Datadeps/MPI) MPI message tag for this task; assigned automatically if `nothing`. - `stream_input_buffer_amount::Union{Int,Nothing}=nothing`: (Streaming only) Specifies the amount of slots to allocate for the input buffer of the task. Defaults to 1. - `stream_output_buffer_amount::Union{Int,Nothing}=nothing`: (Streaming only) Specifies the amount of slots to allocate for the output buffer of the task. Defaults to 1. - `stream_buffer_type::Union{Type,Nothing}=nothing`: (Streaming only) Specifies the type of buffer to use for the input and output buffers of the task. Defaults to `Dagger.ProcessRingBuffer`. - `stream_max_evals::Union{Int,Nothing}=nothing`: (Streaming only) Specifies the maximum number of times the task will be evaluated before returning a result. Defaults to infinite evaluations. +- `acceleration::Union{Acceleration,Nothing}=nothing`: The acceleration backend used to plan and execute this task (e.g. `DistributedAcceleration`, `MPIAcceleration`). When `nothing`, the current acceleration (`Dagger.current_acceleration()`) is used. +- `return_type::Union{Type,Nothing}=nothing`: The expected return type of the task's function. When set to a concrete type, it is used as the task's `chunktype` before the task has run (e.g. so downstream metadata and, under MPI, cross-rank type uniformity are known ahead of execution). When `nothing`, the type is left unknown until the result is available. """ Base.@kwdef mutable struct Options propagates::Union{Vector{Symbol},Nothing} = nothing @@ -61,10 +64,16 @@ Base.@kwdef mutable struct Options name::Union{String,Nothing} = nothing + tag::Union{UInt32,Nothing} = nothing + stream_input_buffer_amount::Union{Int,Nothing} = nothing stream_output_buffer_amount::Union{Int,Nothing} = nothing stream_buffer_type::Union{Type, Nothing} = nothing stream_max_evals::Union{Int,Nothing} = nothing + + acceleration::Union{Acceleration,Nothing} = nothing + + return_type::Union{Type,Nothing} = nothing end Options(::Nothing) = Options() function Options(old_options::NamedTuple) diff --git a/src/processor.jl b/src/processor.jl index d21a05548..156c6cd5e 100644 --- a/src/processor.jl +++ b/src/processor.jl @@ -155,3 +155,9 @@ iscompatible_arg(proc::OSProc, opts, args...) = "Returns a very brief `String` representation of `proc`." short_name(proc::Processor) = string(proc) short_name(p::OSProc) = "W: $(p.pid)" + +"Returns true if the processor is on the local worker (for MPI/ordering)." +is_local_processor(proc::Processor) = (root_worker_id(proc) == myid()) + +"Ordering key for task firing (used by MPI to avoid deadlock)." +fire_order_key(proc::Processor) = (root_worker_id(proc), 0) diff --git a/src/procgrid.jl b/src/procgrid.jl new file mode 100644 index 000000000..7208a4ebf --- /dev/null +++ b/src/procgrid.jl @@ -0,0 +1,150 @@ +# Compact processor grids for distributed array block assignment. + +abstract type AbstractProcGrid{N} end + +struct CyclicProcGrid{N} <: AbstractProcGrid{N} + procs::Vector{Processor} + shape::NTuple{N,Int} +end + +struct BlockedProcGrid{N} <: AbstractProcGrid{N} + procs::Vector{Processor} + counts::Vector{Int} + shape::NTuple{N,Int} + dim::Int +end + +struct DenseProcGrid{N} <: AbstractProcGrid{N} + grid::Array{Processor,N} +end + +DenseProcGrid(grid::AbstractArray{<:Processor,N}) where {N} = + DenseProcGrid{N}(Array{Processor,N}(grid)) + +const AssignmentType{N} = Union{Symbol, AbstractArray{<:Int, N}, AbstractArray{<:Processor, N}} + +procgrid_shape(pg::AbstractProcGrid{N}) where {N} = pg.shape +procgrid_shape(pg::DenseProcGrid{N}) where {N} = size(pg.grid) + +function _tile_index(pg::AbstractProcGrid{N}, I::CartesianIndex{N}) where {N} + shape = procgrid_shape(pg) + return CartesianIndex(mod1.(Tuple(I), shape)) +end + +function procgrid_processor(pg::CyclicProcGrid{N}, I::CartesianIndex{N}) where {N} + I_tiled = _tile_index(pg, I) + lin = LinearIndices(pg.shape)[I_tiled] + return pg.procs[mod1(lin, length(pg.procs))] +end + +function procgrid_processor(pg::BlockedProcGrid{N}, I::CartesianIndex{N}) where {N} + I_tiled = _tile_index(pg, I) + idx = I_tiled[pg.dim] + acc = 0 + for (p, c) in enumerate(pg.counts) + acc += c + if idx <= acc + return pg.procs[p] + end + end + return pg.procs[end] +end + +procgrid_processor(pg::DenseProcGrid{N}, I::CartesianIndex{N}) where {N} = + pg.grid[_tile_index(pg, I)] + +function procgrid_scope(procgrid, I::CartesianIndex) + if procgrid === nothing + return get_compute_scope() + end + return ExactScope(procgrid_processor(procgrid, I)) +end + +normalize_procgrid(procgrid::Union{AbstractProcGrid, Nothing}) = procgrid +normalize_procgrid(procgrid::AbstractArray{<:Processor,N}) where {N} = DenseProcGrid(procgrid) + +function build_procgrid(assignment::AssignmentType{N}, sizeA::NTuple{N,Int}, + blocksize::NTuple{N,Int}, accel::Acceleration) where {N} + # N.B. Collect with the abstract `Processor` eltype: the grid structs store + # `Vector{Processor}` (invariant), so a concrete `Vector{ThreadProc}` from + # `compatible_processors` would not match their constructors. + availprocs = collect(Processor, compatible_processors(accel)) + if !(assignment isa AbstractArray{<:Processor, N}) + filter!(p -> p isa ThreadProc, availprocs) + sort!(availprocs, by=x -> (x.owner, x.tid)) + end + np = length(availprocs) + nblocks = ntuple(i -> cld(sizeA[i], blocksize[i]), N) + + if assignment isa Symbol + if assignment == :arbitrary + return default_procgrid(accel, nblocks) + elseif assignment == :blockrow + shape = ntuple(i -> i == 1 ? Int(ceil(sizeA[1] / blocksize[1])) : 1, N) + nrows = shape[1] + rows_per_proc, extra = divrem(nrows, np) + counts = [rows_per_proc + (i <= extra ? 1 : 0) for i in 1:np] + return BlockedProcGrid(availprocs, counts, shape, 1) + elseif assignment == :blockcol + shape = ntuple(i -> i == N ? Int(ceil(sizeA[N] / blocksize[N])) : 1, N) + ncols = shape[N] + cols_per_proc, extra = divrem(ncols, np) + counts = [cols_per_proc + (i <= extra ? 1 : 0) for i in 1:np] + return BlockedProcGrid(availprocs, counts, shape, N) + elseif assignment == :cyclicrow + shape = ntuple(i -> i == 1 ? np : 1, N) + return CyclicProcGrid(availprocs, shape) + elseif assignment == :cycliccol + shape = ntuple(i -> i == N ? np : 1, N) + return CyclicProcGrid(availprocs, shape) + else + error("Unsupported assignment symbol: $assignment, use :arbitrary, :blockrow, :blockcol, :cyclicrow or :cycliccol") + end + elseif assignment isa AbstractArray{<:Int, N} + missingprocs = filter(p -> p ∉ procs(), assignment) + isempty(missingprocs) || error("Specified workers are not available: $missingprocs") + return DenseProcGrid([ThreadProc(proc, 1) for proc in assignment]) + elseif assignment isa AbstractArray{<:Processor, N} + missingprocs = filter(p -> p ∉ availprocs, assignment) + isempty(missingprocs) || error("Specified processors are not available: $missingprocs") + return DenseProcGrid(assignment) + end + error("Invalid assignment type: $(typeof(assignment))") +end + +function emit_chunk_tasks!(domainchunks, procgrid, eltype::Type{T}, spawn_chunk::Function) where {T} + N = ndims(domainchunks) + tasks = Array{Any,N}(undef, size(domainchunks)...) + default_scope = get_compute_scope() + function emit!() + for I in CartesianIndices(domainchunks) + scope = procgrid === nothing ? default_scope : procgrid_scope(procgrid, I) + i = LinearIndices(domainchunks)[I] + tasks[i] = spawn_chunk(scope, I, i) + end + end + # Under SPMD/MPI (uniform execution) the chunk-placement tasks must run + # inside a datadeps region so every rank agrees on scheduling and data + # movement. Under Distributed the placements are independent `@spawn`s with + # explicit scopes; wrapping them in datadeps would needlessly require the + # element type to be aliasing-resolvable (breaking e.g. non-isbits eltypes). + if uniform_execution() + spawn_datadeps(emit!) + else + emit!() + end + return post_stage_array_chunks!(current_acceleration(), tasks, eltype, N) +end + +"Materialize a dense processor array (for tests and parity checks)." +function materialize_procgrid(procgrid::Union{AbstractProcGrid{N},Nothing}, nblocks::NTuple{N,Int}) where {N} + procgrid === nothing && return nothing + shape = ntuple(i -> max(nblocks[i], procgrid_shape(procgrid)[i]), N) + grid = Array{Processor,N}(undef, shape) + for I in CartesianIndices(grid) + grid[I] = procgrid_processor(procgrid, I) + end + return grid +end + +materialize_procgrid(procgrid::DenseProcGrid{N}, ::NTuple{N,Int}) where {N} = procgrid.grid diff --git a/src/queue.jl b/src/queue.jl index ee57c1d4e..2c8fd1f51 100644 --- a/src/queue.jl +++ b/src/queue.jl @@ -164,10 +164,11 @@ function wait_all(f; check_errors::Bool=false) result = with_options(f; task_queue=queue) for task in queue.tasks if check_errors - fetch(task; raw=true) + fetch(task; move_value=false, unwrap=false) else wait(task) end end + cleanup_tasks_accel!(current_acceleration(), queue.tasks) return result end diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 958bb30c1..6a353e28c 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -17,7 +17,8 @@ import ..Dagger import ..Dagger: Context, Processor, SchedulerOptions, Options, Thunk, WeakThunk, ThunkFuture, ThunkID, DTaskFailedException, Chunk, WeakChunk, OSProc, AnyScope, DefaultScope, InvalidScope, LockedObject, Argument, Signature import ..Dagger: Sealed, SEALED, FutureNode, futures_push!, futures_seal! import ..Dagger: DepNode, deps_push!, deps_seal! -import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak, unwrap_weak_checked, wrap_weak, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, default_enabled, processor, get_processors, get_parent, execute!, rmprocs!, task_processor, constrain, cputhreadtime, maybe_take_or_alloc! +import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak, unwrap_weak_checked, wrap_weak, tochunk, timespan_start, timespan_finish, procs, move, chunktype, default_enabled, processor, get_processors, get_parent, execute!, rmprocs!, task_processor, constrain, cputhreadtime, maybe_take_or_alloc! +import ..Dagger: datasize, root_worker_id, is_local_processor, fire_order_key, short_name, select_processors_uniform!, processor_order_key, current_acceleration, set_task_acceleration!, scheduling_ignore_capacity, scheduling_task_occupancy, schedule_argument_move, bind_moved_argument import ..Dagger: @dagdebug, @safe_lock_spin1, @maybelog, @take_or_alloc! import DataStructures: PriorityQueue @@ -27,7 +28,7 @@ import ..Dagger: @reusable, @reusable_dict, @reusable_vector, @reusable_tasks, @ import TimespanLogging import TaskLocalValues: TaskLocalValue -import ScopedValues: @with +import ScopedValues: ScopedValue, @with, with include("util.jl") @@ -211,27 +212,30 @@ const WORKER_MONITOR_TASKS = Dict{Int,Task}() const WORKER_MONITOR_CHANS = Dict{Int,Dict{UInt64,RemoteChannel}}() function init_proc(state, p, log_sink) ctx = Context(Int[]; log_sink) - @maybelog ctx timespan_start(ctx, :init_proc, (;uid=state.uid, worker=p.pid), nothing) - # Initialize pressure and capacity - gproc = OSProc(p.pid) + pid = Dagger.root_worker_id(p) + @maybelog ctx timespan_start(ctx, :init_proc, (;uid=state.uid, worker=pid), nothing) + # Initialize pressure and capacity. + # N.B. Key by `root_worker_id(p)` (== `pid`) rather than `p.pid`; MPI + # processors carry no `pid` field but report a root worker id. + gproc = OSProc(pid) lock(state.worker_time_pressure) do wtp - wtp[p.pid] = Dict{Processor,Threads.Atomic{UInt64}}() + wtp[pid] = Dict{Processor,Threads.Atomic{UInt64}}() end lock(state.worker_transfer_rate) do wtr - wtr[p.pid] = Dict{Processor,UInt64}() + wtr[pid] = Dict{Processor,UInt64}() end lock(state.worker_storage_pressure) do wsp - wsp[p.pid] = Dict{Union{StorageResource,Nothing},UInt64}() + wsp[pid] = Dict{Union{StorageResource,Nothing},UInt64}() end lock(state.worker_storage_capacity) do wsc - wsc[p.pid] = Dict{Union{StorageResource,Nothing},UInt64}() + wsc[pid] = Dict{Union{StorageResource,Nothing},UInt64}() end lock(state.worker_loadavg) do wla - wla[p.pid] = (0.0, 0.0, 0.0) + wla[pid] = (0.0, 0.0, 0.0) end - if p.pid != 1 + if pid != 1 lock(WORKER_MONITOR_LOCK) do - wid = p.pid + wid = pid if !haskey(WORKER_MONITOR_TASKS, wid) t = Threads.@spawn begin try @@ -265,16 +269,16 @@ function init_proc(state, p, log_sink) end # Setup worker-to-scheduler channels - inp_chan = RemoteChannel(p.pid) - out_chan = RemoteChannel(p.pid) + inp_chan = RemoteChannel(pid) + out_chan = RemoteChannel(pid) lock(state.lock) do - state.worker_chans[p.pid] = (inp_chan, out_chan) + state.worker_chans[pid] = (inp_chan, out_chan) end # Setup dynamic listener - dynamic_listener!(ctx, state, p.pid) + dynamic_listener!(ctx, state, pid) - @maybelog ctx timespan_finish(ctx, :init_proc, (;uid=state.uid, worker=p.pid), nothing) + @maybelog ctx timespan_finish(ctx, :init_proc, (;uid=state.uid, worker=pid), nothing) end function _cleanup_proc(uid, log_sink) empty!(CHUNK_CACHE) # FIXME: Should be keyed on uid! @@ -292,7 +296,7 @@ function _cleanup_proc(uid, log_sink) end function cleanup_proc(state, p, log_sink) ctx = Context(Int[]; log_sink) - wid = p.pid + wid = root_worker_id(p) @maybelog ctx timespan_start(ctx, :cleanup_proc, (;uid=state.uid, worker=wid), nothing) lock(WORKER_MONITOR_LOCK) do if haskey(WORKER_MONITOR_CHANS, wid) @@ -355,7 +359,7 @@ function compute_dag(ctx::Context, d::Thunk, options=SchedulerOptions()) node_order = x -> -get(ord, x, 0) state = start_state(deps, node_order, chan, ctx, options) - master = OSProc(myid()) + master = Dagger.default_processor() # Register so in-process completions can find this state register_scheduler_state!(state) @@ -516,10 +520,10 @@ function handle_result!(ctx, state::ComputeState, pid, proc, thunk_id, res, meta end end end - if res isa Chunk + if res isa Chunk && res.handle isa DRef lock(state.equiv_chunks) do ec - if !haskey(ec, res.handle::DRef) - ec[res.handle::DRef] = res + if !haskey(ec, res.handle) + ec[res.handle] = res end end end @@ -712,7 +716,7 @@ end const CHUNK_CACHE = Dict{Chunk,Dict{Processor,Any}}() struct ScheduleTaskLocation - gproc::OSProc + gproc::Processor proc::Processor end struct ScheduleTaskSpec @@ -889,9 +893,15 @@ concurrently across threads. return end + # `accel` drives uniform (e.g. MPI) scheduling: it filters compatible + # processors, imposes a deterministic processor order, and (under uniform + # execution) makes selection ignore rank-local capacity/occupancy so every + # rank picks the same processor for this task. + accel = something(options.acceleration, current_acceleration()) + input_procs = @reusable_vector :schedule_one!_input_procs Processor OSProc() 32 input_procs_cleanup = @reuse_defer_cleanup empty!(input_procs) - compat = Dagger.compatible_processors(scope, procs_filt) + compat = Dagger.compatible_processors(accel, scope, procs_filt) for proc in compat if !(proc in input_procs) push!(input_procs, proc) @@ -906,6 +916,10 @@ concurrently across threads. estimate_task_costs!(sorted_procs, costs, state, input_procs, task; sig) input_procs_cleanup() + # Under uniform execution, measured costs are rank-local, so re-order by a + # deterministic, acceleration-dispatched key instead (no-op otherwise). + select_processors_uniform!(sorted_procs, accel, task, state) + # Select the best available processor and reserve time pressure optimistically. # These operations use their own fine-grained locks — no state.lock needed. scheduled = false @@ -915,11 +929,13 @@ concurrently across threads. can_use, scope = can_use_proc(state, task, gproc, proc, options, scope) if can_use has_cap, est_time_util, est_alloc_util, est_occupancy = - has_capacity(state, proc, gproc.pid, options.time_util, options.alloc_util, options.occupancy, sig) - if has_cap + has_capacity(state, proc, root_worker_id(gproc), options.time_util, options.alloc_util, options.occupancy, sig) + # Under uniform execution capacity is rank-local; every rank must + # take the first usable processor in the deterministic order. + if has_cap || scheduling_ignore_capacity(accel) # Optimistic reservation: capture-the-ref, atomic_add! counter_ref = lock(state.worker_time_pressure) do wtp - proc_map = get!(wtp, gproc.pid) do + proc_map = get!(wtp, root_worker_id(gproc)) do Dict{Processor,Threads.Atomic{UInt64}}() end get!(proc_map, proc) do @@ -931,6 +947,9 @@ concurrently across threads. rpr[task.id] = (counter_ref, est_time_util) end @dagdebug task :schedule "Scheduling to $gproc -> $proc (cost: $(costs[proc]), pressure: $(counter_ref[]))" + # Uniform execution overlaps communicating tasks on a processor + # (rank-local shares); avoid serializing them into wait cycles. + est_occupancy = scheduling_task_occupancy(accel, est_occupancy) best_loc = ScheduleTaskLocation(gproc, proc) best_spec = ScheduleTaskSpec(task, scope, est_time_util, est_alloc_util, est_occupancy) scheduled = true @@ -1003,17 +1022,19 @@ function monitor_procs_changed!(ctx, state, options) end function remove_dead_proc!(ctx, state, proc, options) - @assert options.single !== proc.pid "Single worker failed, cannot continue." - rmprocs!(ctx, [proc]) + pid = root_worker_id(proc) + @assert options.single !== pid "Single worker failed, cannot continue." + rmprocs!(ctx, [pid]) # COW-style membership removal: lock each map and delete the worker's entry. # Any in-flight tasks that captured a counter ref from this worker will still # release via atomic_sub! on the (now-orphaned) Atomic object — harmless. - lock(state.worker_time_pressure) do wtp; delete!(wtp, proc.pid); end - lock(state.worker_transfer_rate) do wtr; delete!(wtr, proc.pid); end - lock(state.worker_storage_pressure) do wsp; delete!(wsp, proc.pid); end - lock(state.worker_storage_capacity) do wsc; delete!(wsc, proc.pid); end - lock(state.worker_loadavg) do wla; delete!(wla, proc.pid); end - delete!(state.worker_chans, proc.pid) + # N.B. Key by `root_worker_id` (MPI processors have no `pid` field). + lock(state.worker_time_pressure) do wtp; delete!(wtp, pid); end + lock(state.worker_transfer_rate) do wtr; delete!(wtr, pid); end + lock(state.worker_storage_pressure) do wsp; delete!(wsp, pid); end + lock(state.worker_storage_capacity) do wsc; delete!(wsc, pid); end + lock(state.worker_loadavg) do wla; delete!(wla, pid); end + delete!(state.worker_chans, pid) end function finish_task!(ctx, state, node, thunk_failed, ready::Vector{Thunk}) @@ -1073,7 +1094,7 @@ end function evict_all_chunks!(ctx, options, to_evict) if !isempty(to_evict) - @sync for w in map(p->p.pid, procs_to_use(ctx, options)) + @sync for w in map(p->root_worker_id(p), procs_to_use(ctx, options)) Threads.@spawn remote_do(evict_chunks!, w, ctx.log_sink, to_evict) end end @@ -1163,9 +1184,10 @@ Base.hash(task::TaskSpec, h::UInt) = hash(task.thunk_id, hash(TaskSpec, h)) end Tf = chunktype(first(args)) - @assert (options.single === nothing) || (gproc.pid == options.single) + pid = root_worker_id(gproc) + @assert (options.single === nothing) || (pid == options.single) # TODO: Set `sch_handle.tid.ref` to the right `DRef` - sch_handle = SchedulerHandle(ThunkID(thunk.id, nothing), state.worker_chans[gproc.pid]...) + sch_handle = SchedulerHandle(ThunkID(thunk.id, nothing), state.worker_chans[pid]...) # TODO: De-dup common fields (log_sink, uid, etc.) push!(to_send, TaskSpec( @@ -1177,7 +1199,7 @@ Base.hash(task::TaskSpec, h::UInt) = hash(task.thunk_id, hash(TaskSpec, h)) end if !isempty(to_send) - if Dagger.root_worker_id(gproc) == myid() + if root_worker_id(gproc) == myid() @reusable_tasks :fire_tasks!_task_cache 32 _->nothing "fire_tasks!" FireTaskSpec(proc, state.chan, to_send) else # N.B. We don't batch these because we might get a deserialization @@ -1392,7 +1414,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re proc_occupancy = istate.proc_occupancy time_pressure = istate.time_pressure - wid = get_parent(to_proc).pid + wid = root_worker_id(to_proc) work_to_do = false while isopen(return_queue) # Wait for new tasks @@ -1450,7 +1472,6 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re # Try to steal from local queues randomly # TODO: Prioritize stealing from busiest processors states = proc_states_values(uid) - # TODO: Try to pre-allocate this P = randperm(length(states)) for state in getindex.(Ref(states), P) other_istate = state.state @@ -1467,7 +1488,8 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re end task, occupancy = first(queue) scope = task.scope - if Dagger.proc_in_scope(to_proc, scope) + accel = something(task.options.acceleration, Dagger.DistributedAcceleration()) + if Dagger.proc_in_scope(to_proc, scope) && Dagger.accel_matches_proc(accel, to_proc) && typemax(UInt32) - proc_occupancy_cached >= occupancy # Compatible, steal this task return popfirst!(queue) @@ -1750,6 +1772,8 @@ function do_tasks(to_proc, return_queue, tasks) @dagdebug nothing :processor "Kicked processors" end +const SCHED_MOVE = ScopedValue{Bool}(false) + """ do_task(to_proc, task::TaskSpec) -> Any @@ -1761,13 +1785,15 @@ Executes a single task specified by `task` on `to_proc`. ctx_vars = task.ctx_vars ctx = Context(Processor[]; log_sink=ctx_vars.log_sink, profile=ctx_vars.profile) - from_proc = OSProc() + options = task.options + Dagger.set_task_acceleration!(options.acceleration) + + from_proc = Dagger.default_processor() data = task.data Tf = task.Tf f = isdefined(Tf, :instance) ? Tf.instance : nothing # Wait for required resources to become available - options = task.options propagated = get_propagated_options(options) to_storage = options.storage !== nothing ? fetch(options.storage) : MemPool.GLOBAL_DEVICE[] #to_storage_name = nameof(typeof(to_storage)) @@ -1835,7 +1861,7 @@ Executes a single task specified by `task` on `to_proc`. @maybelog ctx timespan_finish(ctx, :storage_wait, (;thunk_id, processor=to_proc), (;f, device=typeof(to_storage))) =# - @dagdebug thunk_id :execute "Moving data" + @dagdebug thunk_id :execute "Moving data for $Tf" # Initiate data transfers for function and arguments transfer_time = Threads.Atomic{UInt64}(0) @@ -1845,9 +1871,13 @@ Executes a single task specified by `task` on `to_proc`. else data end + # Under uniform execution, argument moves may communicate between ranks on + # rank-uniform tags derived from the thunk ID; they must run sequentially + # in argument order (FIFO tag matching), with the thunk's TID in scope + accel = something(options.acceleration, current_acceleration()) fetch_tasks = map(_data) do arg #=FIXME:REALLOC_TASKS=# - Threads.@spawn begin + move_one = () -> begin value = Dagger.value(arg) position = arg.pos @maybelog ctx timespan_start(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=value)) @@ -1858,7 +1888,9 @@ Executes a single task specified by `task` on `to_proc`. Some{Any}(get!(CHUNK_CACHE[x], to_proc) do # Convert from cached value # TODO: Choose "closest" processor of same type first - some_proc = first(keys(CHUNK_CACHE[x])) + cache_procs = keys(CHUNK_CACHE[x]) + accel = something(options.acceleration, current_acceleration()) + some_proc = minimum(cache_procs, by=p -> processor_order_key(accel, p)) some_x = CHUNK_CACHE[x][some_proc] @dagdebug thunk_id :move "Cache hit for argument $id at $some_proc: $some_x" @invokelatest move(some_proc, to_proc, some_x) @@ -1893,17 +1925,26 @@ Executes a single task specified by `task` on `to_proc`. end else =# - new_value = @invokelatest move(to_proc, value) + new_value = with(SCHED_MOVE=>true) do + @invokelatest move(to_proc, value) + end #end - if new_value !== value - @dagdebug thunk_id :move "Moved argument @ $position to $to_proc: $(typeof(value)) -> $(typeof(new_value))" + # Acceleration decides how to bind the moved value (e.g. keep a + # Chunk placeholder, or wrap an owner unwrap so chunktype stays + # SPMD-uniform). Materializing for the kernel may still happen in + # execute! when this rank only holds a placeholder. + bound = bind_moved_argument(accel, value, new_value) + if bound !== value + @dagdebug thunk_id :move "Moved argument @ $position to $to_proc: $(typeof(value)) -> $(typeof(bound))" end - @maybelog ctx timespan_finish(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=new_value); tasks=[Base.current_task()]) - arg.value = new_value + arg.value = bound + @maybelog ctx timespan_finish(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=Dagger.value(arg)); tasks=[Base.current_task()]) return end + return schedule_argument_move(accel, thunk_id, move_one) end for task in fetch_tasks + task === nothing && continue fetch_report(task) end @@ -1938,7 +1979,7 @@ Executes a single task specified by `task` on `to_proc`. # FIXME #gcnum_start = Base.gc_num() - @dagdebug thunk_id :execute "Executing $(typeof(f))" + @dagdebug thunk_id :execute "Executing $Tf" logging_enabled = !(ctx.log_sink isa TimespanLogging.NoOpLog) @@ -2001,7 +2042,7 @@ Executes a single task specified by `task` on `to_proc`. notify(TASK_SYNC) end - @dagdebug thunk_id :execute "Returning" + @dagdebug thunk_id :execute "Returning $Tf with $(typeof(result_meta))" # TODO: debug_storage("Releasing $to_storage_name") metadata = ( diff --git a/src/sch/util.jl b/src/sch/util.jl index 6f7f21a5f..9057f5740 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -428,7 +428,7 @@ function signature(f, args) value = Dagger.value(arg) if value isa Dagger.DTask # Only occurs via manual usage of signature - value = fetch(value; raw=true) + value = fetch(value; move_value=false, unwrap=false) end if istask(value) throw(ConcurrencyViolationError("Must call `collect_task_inputs!(state, task)` before calling `signature`")) @@ -485,8 +485,8 @@ function can_use_proc(state, task, gproc, proc, opts, scope) # Check against single if opts.single !== nothing @warn "The `single` option is deprecated, please use scopes instead\nSee https://juliaparallel.org/Dagger.jl/stable/scopes/ for details" maxlog=1 - if gproc.pid != opts.single - @dagdebug task :scope "Rejected $proc: gproc.pid ($(gproc.pid)) != single ($(opts.single))" + if root_worker_id(gproc) != opts.single + @dagdebug task :scope "Rejected $proc: gproc root_worker_id ($(root_worker_id(gproc))) != single ($(opts.single))" return false, scope end scope = constrain(scope, Dagger.ProcessScope(opts.single)) @@ -660,18 +660,18 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) chunks_filt = Iterators.filter(c->get_parent(processor(c)) != gproc, chunks) # Estimate network transfer costs based on data size - # N.B. `affinity(x)` really means "data size of `x`" # N.B. We treat same-worker transfers as having zero transfer cost # TODO: For non-Chunk, model cost from scheduler to worker # TODO: Measure and model processor move overhead - tx_cost = impute_sum(affinity(chunk)[2] for chunk in chunks_filt) + tx_cost = impute_sum(datasize(chunk) for chunk in chunks_filt) # Add fixed cost for cross-worker task transfer (esimated at 1ms) # TODO: Actually estimate/benchmark this - task_xfer_cost = gproc.pid != myid() ? 1_000_000 : 0 # 1ms + task_xfer_cost = root_worker_id(gproc) != myid() ? 1_000_000 : 0 # 1ms + pid = Dagger.root_worker_id(gproc) tx_rate = lock(state.worker_transfer_rate) do wtr - get(get(wtr, gproc.pid, Dict{Processor,UInt64}()), proc, DEFAULT_TRANSFER_RATE) + get(get(wtr, pid, Dict{Processor,UInt64}()), proc, DEFAULT_TRANSFER_RATE) end costs[proc] = est_time_util + (tx_cost/tx_rate) + task_xfer_cost end diff --git a/src/scopes.jl b/src/scopes.jl index 79190c292..28aa8fa00 100644 --- a/src/scopes.jl +++ b/src/scopes.jl @@ -101,7 +101,7 @@ struct ExactScope <: AbstractScope parent::ProcessScope processor::Processor end -ExactScope(proc) = ExactScope(ProcessScope(get_parent(proc).pid), proc) +ExactScope(proc) = ExactScope(ProcessScope(root_worker_id(get_parent(proc))), proc) proc_in_scope(proc::Processor, scope::ExactScope) = proc == scope.processor "Indicates that the applied scopes `x` and `y` are incompatible." diff --git a/src/shard.jl b/src/shard.jl new file mode 100644 index 000000000..ecd0ee570 --- /dev/null +++ b/src/shard.jl @@ -0,0 +1,89 @@ +""" +Maps a value to one of multiple distributed "mirror" values automatically when +used as a thunk argument. Construct using `@shard` or `shard`. +""" +struct Shard + chunks::Dict{Processor,Chunk} +end + +""" + shard(f; kwargs...) -> Chunk{Shard} + +Executes `f` on all workers in `workers`, wrapping the result in a +process-scoped `Chunk`, and constructs a `Chunk{Shard}` containing all of these +`Chunk`s on the current worker. + +Keyword arguments: +- `procs` -- The list of processors to create pieces on. May be any iterable container of `Processor`s. +- `workers` -- The list of workers to create pieces on. May be any iterable container of `Integer`s. +- `per_thread::Bool=false` -- If `true`, creates a piece per each thread, rather than a piece per each worker. +""" +function shard(@nospecialize(f); procs=nothing, workers=nothing, per_thread=false) + if procs === nothing + if workers !== nothing + procs = [OSProc(w) for w in workers] + else + procs = lock(Sch.eager_context()) do + copy(Sch.eager_context().procs) + end + end + if per_thread + _procs = ThreadProc[] + for p in procs + append!(_procs, filter(p->p isa ThreadProc, get_processors(p))) + end + procs = _procs + end + else + if workers !== nothing + throw(ArgumentError("Cannot combine `procs` and `workers`")) + elseif per_thread + throw(ArgumentError("Cannot combine `procs` and `per_thread=true`")) + end + end + isempty(procs) && throw(ArgumentError("Cannot create empty Shard")) + shard_running_dict = Dict{Processor,DTask}() + for proc in procs + scope = proc isa OSProc ? ProcessScope(proc) : ExactScope(proc) + thunk = Dagger.@spawn scope=scope _mutable_inner(f, proc, scope) + shard_running_dict[proc] = thunk + end + shard_dict = Dict{Processor,Chunk}() + for proc in procs + shard_dict[proc] = fetch(shard_running_dict[proc])[] + end + return Shard(shard_dict) +end + +"Creates a `Shard`. See [`Dagger.shard`](@ref) for details." +macro shard(exs...) + opts = esc.(exs[1:end-1]) + ex = exs[end] + quote + let f = @noinline ()->$(esc(ex)) + $shard(f; $(opts...)) + end + end +end + +function move(from_proc::Processor, to_proc::Processor, shard::Shard) + # Match either this proc or some ancestor + # N.B. This behavior may bypass the piece's scope restriction + proc = to_proc + if haskey(shard.chunks, proc) + return move(from_proc, to_proc, shard.chunks[proc]) + end + parent = Dagger.get_parent(proc) + while parent != proc + proc = parent + parent = Dagger.get_parent(proc) + if haskey(shard.chunks, proc) + return move(from_proc, to_proc, shard.chunks[proc]) + end + end + + throw(KeyError(to_proc)) +end +Base.iterate(s::Shard) = iterate(values(s.chunks)) +Base.iterate(s::Shard, state) = iterate(values(s.chunks), state) +Base.length(s::Shard) = length(s.chunks) diff --git a/src/submission.jl b/src/submission.jl index 221a83e79..8d82c5034 100644 --- a/src/submission.jl +++ b/src/submission.jl @@ -112,6 +112,11 @@ const UID_TO_TID_CACHE = TaskLocalValue{ReusableCache{Dict{UInt64,Int},Nothing}} # so we just pick an equivalent Chunk as our upstream chunk = value(arg)::Chunk function find_equivalent_chunk(state, chunk::C) where {C<:Chunk} + # `equiv_chunks` is a `WeakKeyDict{DRef,Chunk}`; only + # DRef-backed chunks participate. Other handles (e.g. + # `MPIRef` under MPI) are not valid keys and manage their + # own identity, so pass them through unchanged. + chunk.handle isa DRef || return chunk lock(state.equiv_chunks) do ec if haskey(ec, chunk.handle) return ec[chunk.handle]::C @@ -123,7 +128,19 @@ const UID_TO_TID_CACHE = TaskLocalValue{ReusableCache{Dict{UInt64,Int},Nothing}} end chunk = find_equivalent_chunk(state, chunk) #=FIXME:UNIQUE=# - @inbounds fargs[idx] = Argument(arg.pos, WeakChunk(chunk)) + if chunk.handle isa DRef + @inbounds fargs[idx] = Argument(arg.pos, WeakChunk(chunk)) + else + # Non-DRef chunks (e.g. `MPIRef` under MPI) are not kept + # alive by `equiv_chunks` (a `WeakKeyDict{DRef,Chunk}`, so + # only DRef-backed wrappers get a strong keeper there). + # Weakening such a chunk here would let its `Chunk` wrapper + # be GC'd before the consuming task runs, expiring the + # `WeakChunk` (observed on Julia 1.10, whose GC is more + # eager). Keep a strong reference; it is released when the + # task's `Thunk` is cleaned up. + @inbounds fargs[idx] = Argument(arg.pos, chunk) + end end end # TODO: Iteration protocol would be faster @@ -332,7 +349,13 @@ function cached_return_type(@nospecialize(f), @nospecialize(arg_types::Tuple)) end end -DTaskMetadata(spec::DTaskSpec) = DTaskMetadata(eager_metadata(spec.fargs)) +function DTaskMetadata(spec::DTaskSpec) + rt = spec.options.return_type + if rt !== nothing && isconcretetype(rt) && rt !== Any + return DTaskMetadata(rt) + end + return DTaskMetadata(eager_metadata(spec.fargs)) +end function eager_metadata(fargs) f = value(fargs[1]) f = f isa StreamingFunction ? f.f : f @@ -345,6 +368,10 @@ function eager_spawn(spec::DTaskSpec) uid = eager_next_id() future = ThunkFuture() metadata = DTaskMetadata(spec) + # Propagate inferred return type to options + if isconcretetype(metadata.return_type) + spec.options.return_type = metadata.return_type + end return DTask(uid, future, metadata) end @@ -367,10 +394,16 @@ function eager_launch!(pair::DTaskPair) end end + # Propagate DTask return_type into options so the created Thunk has chunktype for downstream inference + options = spec.options + if isconcretetype(task.metadata.return_type) + options = copy(options) + options.return_type = task.metadata.return_type + end # Submit the task #=FIXME:REALLOC=# thunk_id = eager_submit!(PayloadOne(task.uid, task.future, - fargs, spec.options, true)) + fargs, options, true)) task.thunk_ref = thunk_id.ref end # FIXME: Don't convert Tuple to Vector{Argument} @@ -400,7 +433,13 @@ function eager_launch!(pairs::Vector{DTaskPair}) end end end - all_options = Options[pair.spec.options for pair in pairs] + # Propagate DTask return_type into options so created Thunks have chunktype for downstream inference + all_options = Options[ + let opts = pair.spec.options + isconcretetype(pair.task.metadata.return_type) ? (o = copy(opts); o.return_type = pair.task.metadata.return_type; o) : opts + end + for pair in pairs + ] # Submit the tasks #=FIXME:REALLOC=# diff --git a/src/thunk.jl b/src/thunk.jl index b7e48e420..732538b2b 100644 --- a/src/thunk.jl +++ b/src/thunk.jl @@ -9,14 +9,12 @@ Base.@kwdef mutable struct ThunkSpec fargs::Vector{Argument} = EMPTY_ARGS id::Int = 0 cache_ref::Any = nothing - affinity::Union{Pair{OSProc,Int}, Nothing} = nothing options::Union{Options, Nothing} = nothing end function unset!(spec::ThunkSpec, _) spec.fargs = EMPTY_ARGS spec.id = 0 spec.cache_ref = nothing - spec.affinity = nothing spec.options = nothing end @@ -75,7 +73,6 @@ mutable struct Thunk inputs::Vector{Argument} # TODO: Use `ImmutableArray` in 1.8 id::Int cache_ref::Any - affinity::Union{Pair{OSProc,Int}, Nothing} options::Union{Options, Nothing} # stores task options eager_accessible::Bool sch_accessible::Bool @@ -83,13 +80,13 @@ mutable struct Thunk @atomic errored::Bool # true if finished with an error result @atomic valid::Bool # true while registered with the scheduler @atomic running::Bool # true while a worker is executing this thunk - running_on::Union{OSProc,Nothing} # which OSProc is running this thunk + running_on::Union{Processor,Nothing} # which (gproc) Processor is running this thunk (an OSProc, or e.g. MPIOSProc under MPI) @atomic futures_head::Union{FutureNode,Nothing,Sealed} # Treiber list of waiting futures @atomic pending_deps::Int # count of unresolved upstream dependencies (dataflow counter) @atomic dependents_head::Union{DepNode,Nothing,Sealed} # Treiber list of downstream dependents function Thunk(spec::ThunkSpec) return new(spec.fargs, spec.id, - spec.cache_ref, spec.affinity, + spec.cache_ref, spec.options, true, true, false, false, false, false, nothing, @@ -168,7 +165,6 @@ function Thunk(f, xs...; syncdeps=nothing, id::Int=next_id(), cache_ref=nothing, - affinity=nothing, options=nothing, propagates=(), kwargs... @@ -216,7 +212,6 @@ function Thunk(f, xs...; @warn "The cache argument is deprecated, as it is now always true" maxlog=1 end spec.cache_ref = cache_ref - spec.affinity = affinity return Thunk(spec) end Serialization.serialize(io::AbstractSerializer, t::Thunk) = @@ -231,40 +226,6 @@ end Base.convert(::Type{ThunkSyncdep}, thunk::Thunk) = ThunkSyncdep(nothing, thunk) -function affinity(t::Thunk) - if t.affinity !== nothing - return t.affinity - end - - if t.cache_ref !== nothing - aff_vec = affinity(t.cache_ref) - else - aff = Dict{OSProc,Int}() - for (_, inp) in t.inputs - #if haskey(state.cache, inp) - # as = affinity(state.cache[inp]) - # for a in as - # proc, sz = a - # aff[proc] = get(aff, proc, 0) + sz - # end - #else - if isa(inp, Union{Chunk, Thunk}) - # TODO if inp is a FileRef, affinity[1] will always be OSProc(1) - proc, sz = affinity(inp) - aff[proc] = get(aff, proc, 0) + sz - end - #end - end - aff_vec = collect(aff) - end - aff_vec - #if length(aff) > 1 - # return sort!(aff_vec, by=last,rev=true) - #else - # return aff_vec - #end -end - is_task_or_chunk(x) = istask(x) function args_kwargs_to_arguments(f, args, kwargs) @@ -344,6 +305,18 @@ isweak(t) = false Base.show(io::IO, t::WeakThunk) = (print(io, "~"); Base.show(io, t.x.value)) Base.convert(::Type{WeakThunk}, t::Thunk) = WeakThunk(t) chunktype(t::WeakThunk) = chunktype(unwrap_weak_checked(t)) +# Use options.return_type when set (e.g. from mpi_propagate_chunk_types! or eager_metadata) +# so that Thunk arguments propagate type to downstream eager_metadata/execute! +function chunktype(t::Thunk) + if t.options !== nothing && t.options.return_type !== nothing && isconcretetype(t.options.return_type) + return t.options.return_type + end + # Fall back to `Any` (unknown result type) rather than `typeof(t)`: a + # `Thunk` is an internal placeholder for a not-yet-computed result, never + # the actual value type, so reporting `Thunk` as the chunk type would + # mislabel a downstream task's argument signature. + return Any +end Base.convert(::Type{ThunkSyncdep}, t::WeakThunk) = ThunkSyncdep(nothing, t) ThunkSyncdep(t::WeakThunk) = ThunkSyncdep(nothing, t) @@ -559,7 +532,7 @@ function _par(mod, ex::Expr; lazy=true, recur=true, opts=()) end args = filter(arg->!Meta.isexpr(arg, :parameters), allargs) kwargs = filter(arg->Meta.isexpr(arg, :parameters), allargs) - if !isempty(kwargs) + if !Base.isempty(kwargs) kwargs = only(kwargs).args end if body !== nothing @@ -672,6 +645,12 @@ function _spawn(args_kwargs, task_options) # N.B. Merges into task_options options_merge!(task_options, scoped_options; override=false) + # Capture the spawning context's acceleration so the scheduler restricts + # processor selection appropriately (e.g. MPI-only procs under MPI) + if task_options.acceleration === nothing + task_options.acceleration = current_acceleration() + end + # Get task queue, and don't let it propagate task_queue = get(scoped_options, :task_queue, DefaultTaskQueue())::AbstractTaskQueue filter!(prop -> prop != :task_queue, propagates) diff --git a/src/tochunk.jl b/src/tochunk.jl new file mode 100644 index 000000000..1bd209fc0 --- /dev/null +++ b/src/tochunk.jl @@ -0,0 +1,145 @@ +""" + tochunk(x, proc::Processor, space::MemorySpace, scope::AbstractScope=AnyScope(); device=nothing, type=typeof(x), rewrap=false, kwargs...) -> Chunk + +Create a chunk from data `x` which resides on `proc`, lives in memory space +`space`, and has scope `scope`. + +`proc` and/or `space` may be omitted, in which case they are derived from the +current acceleration (via `default_processor`/`default_memory_space`); `scope` +defaults to `AnyScope()`. The one-argument form `tochunk(x)` derives all three, +so that (for example) under `MPIAcceleration` the resulting `Chunk` gets an +`MPIRef` handle and `MPIMemorySpace`, rather than a rank-local `DRef` and +`CPURAMMemorySpace`. + +`space` determines how the data is stored: `tochunk_pset(x, space; ...)` builds +the appropriate handle for that space (e.g. a `DRef` for `CPURAMMemorySpace`, an +`MPIRef` under MPI). It is also recorded on the `Chunk` for aliasing and +data-movement decisions. + +`device` specifies a `MemPool.StorageDevice` (which is itself wrapped in a +`Chunk`) which will be used to manage the reference contained in the `Chunk` +generated by this function. If `device` is `nothing` (the default), the data +will be inspected to determine if it's safe to serialize; if so, the default +MemPool storage device will be used; if not, then a `MemPool.CPURAMDevice` will +be used. + +`type` can be specified manually to force the type to be `Chunk{type}`. A +placeholder chunk for remote data (`x === nothing`) must pass an explicit, +non-`Nothing` `type`. + +If `rewrap==true` and `x isa Chunk`, then the `Chunk` will be rewrapped in a +new `Chunk`. + +All other kwargs are passed directly to `MemPool.poolset`. +""" +tochunk(x::X, proc::P, space::M; kwargs...) where {X,P<:Processor,M<:MemorySpace} = + tochunk(x, proc, space, AnyScope(); kwargs...) +function tochunk(x::X, proc::P, space::M, scope::S; device=nothing, type=X, rewrap=false, kwargs...) where {X,P<:Processor,S,M<:MemorySpace} + if type === Nothing && x !== nothing + # Placeholders for remote data must carry an explicit element type + throw(ArgumentError("Chunk type cannot be Nothing. Placeholder chunks must be created with an explicit type= (e.g. tochunk(nothing, proc, space; type=Matrix{Float64})). x=$(repr(x))")) + end + if x isa Chunk + check_proc_space(x, proc, space) + return maybe_rewrap(x, proc, space, scope; type, rewrap) + end + if device === nothing + device = if Sch.walk_storage_safe(x) + MemPool.GLOBAL_DEVICE[] + else + MemPool.CPURAMDevice() + end + end + ref = tochunk_pset(x, space; device, type, kwargs...) + return Chunk{type,typeof(ref),P,S,typeof(space)}(type, domain(x), ref, proc, scope, space) +end +# Disambiguate: Chunk-specific 3-arg so kwcall(tochunk, Chunk, Processor, Scope) is not ambiguous with utils/chunks.jl +function tochunk(x::Chunk, proc::P, scope::S; rewrap=false, kwargs...) where {P<:Processor,S} + if rewrap + return remotecall_fetch(x.handle.owner) do + tochunk(MemPool.poolget(x.handle), proc, scope; kwargs...) + end + else + return x + end +end +function tochunk(x::X, proc::P, scope::S; device=nothing, type=X, rewrap=false, kwargs...) where {X,P<:Processor,S} + if type === Nothing && x !== nothing + # Placeholders for remote data must carry an explicit element type + throw(ArgumentError("Chunk type cannot be Nothing. Placeholder chunks must be created with an explicit type= (e.g. tochunk(nothing, proc, space; type=Matrix{Float64})). x=$(repr(x))")) + end + if device === nothing + device = if Sch.walk_storage_safe(x) + MemPool.GLOBAL_DEVICE[] + else + MemPool.CPURAMDevice() + end + end + if x isa Chunk + space = x.space + check_proc_space(x, proc, space) + return maybe_rewrap(x, proc, space, scope; type, rewrap) + end + space = default_memory_space(current_acceleration(), x) + ref = tochunk_pset(x, space; device, type, kwargs...) + return Chunk{type,typeof(ref),P,S,typeof(space)}(type, domain(x), ref, proc, scope, space) +end +function tochunk(x::X, space::M, scope::S; device=nothing, type=X, rewrap=false, kwargs...) where {X,M<:MemorySpace,S} + if type === Nothing && x !== nothing + # Placeholders for remote data must carry an explicit element type + throw(ArgumentError("Chunk type cannot be Nothing. Placeholder chunks must be created with an explicit type= (e.g. tochunk(nothing, proc, space; type=Matrix{Float64})). x=$(repr(x))")) + end + if device === nothing + device = if Sch.walk_storage_safe(x) + MemPool.GLOBAL_DEVICE[] + else + MemPool.CPURAMDevice() + end + end + if x isa Chunk + proc = x.processor + check_proc_space(x, proc, space) + return maybe_rewrap(x, proc, space, scope; type, rewrap) + end + proc = default_processor(current_acceleration(), x) + ref = tochunk_pset(x, space; device, type, kwargs...) + return Chunk{type,typeof(ref),typeof(proc),S,M}(type, domain(x), ref, proc, scope, space) +end +# 2-arg: avoid overwriting utils/chunks.jl's tochunk(Any, Any) and tochunk(Any); only add Processor/MemorySpace variants +# Chunk + Processor: disambiguate vs utils/chunks.jl's tochunk(x::Chunk, proc; ...) +tochunk(x::Chunk, proc::Processor; kwargs...) = tochunk(x, proc, AnyScope(); kwargs...) +tochunk(x, proc::Processor; kwargs...) = tochunk(x, proc, AnyScope(); kwargs...) +tochunk(x, space::MemorySpace; kwargs...) = tochunk(x, space, AnyScope(); kwargs...) + +# 1-arg: route through the current acceleration so the resulting Chunk gets +# the appropriate handle and memory space (e.g. MPIRef + MPIMemorySpace under +# MPIAcceleration, instead of a rank-local DRef + CPURAMMemorySpace). +# N.B. This overrides the default-argument method generated by +# utils/chunks.jl's `tochunk(x, proc=OSProc(), scope=AnyScope())`. +function tochunk(x) + accel = current_acceleration() + return tochunk(x, default_processor(accel, x), default_memory_space(accel, x), AnyScope()) +end + +check_proc_space(x, proc, space) = nothing +function check_proc_space(x::Chunk, proc, space) + if x.space !== space + throw(ArgumentError("Memory space mismatch: Chunk=$(x.space) != Requested=$space")) + end +end +function check_proc_space(x::Thunk, proc, space) + # FIXME: Validate +end +function maybe_rewrap(x, proc, space, scope; type, rewrap) + if rewrap + return remotecall_fetch(x.handle.owner) do + tochunk(MemPool.poolget(x.handle), proc, scope; kwargs...) + end + else + return x + end +end + +tochunk_pset(x, space::MemorySpace; device=nothing, type=nothing, kwargs...) = poolset(x; device, kwargs...) + +# savechunk: defined in utils/chunks.jl (fork Chunk has space field; do not duplicate here) diff --git a/src/types/acceleration.jl b/src/types/acceleration.jl new file mode 100644 index 000000000..f9aa1d86f --- /dev/null +++ b/src/types/acceleration.jl @@ -0,0 +1,3 @@ +abstract type Acceleration end + +struct DistributedAcceleration <: Acceleration end diff --git a/src/types/chunk.jl b/src/types/chunk.jl new file mode 100644 index 000000000..80cb3f74f --- /dev/null +++ b/src/types/chunk.jl @@ -0,0 +1,33 @@ +""" + Chunk + +A reference to a piece of data located on a remote worker. `Chunk`s are +typically created with `Dagger.tochunk(data)`, and the data can then be +accessed from any worker with `collect(::Chunk)`. `Chunk`s are +serialization-safe, and use distributed refcounting (provided by +`MemPool.DRef`) to ensure that the data referenced by a `Chunk` won't be GC'd, +as long as a reference exists on some worker. + +Each `Chunk` is associated with a given `Dagger.Processor`, which is (in a +sense) the processor that "owns" or contains the data. Calling +`collect(::Chunk)` will perform data movement and conversions defined by that +processor to safely serialize the data to the calling worker. + +Each `Chunk` also records the `Dagger.MemorySpace` (the `space` field) that the +data resides in — for example `CPURAMMemorySpace` for host memory, a GPU memory +space for device arrays, or `MPIMemorySpace` under `MPIAcceleration`. The memory +space is set from the `space` argument to [`tochunk`](@ref) (or derived from the +current acceleration when omitted), and is used by datadeps for aliasing and +data-movement decisions, so it must reflect where the data actually lives. + +## Constructors +See [`tochunk`](@ref). +""" +mutable struct Chunk{T, H, P<:Processor, S<:AbstractScope, M<:MemorySpace} + chunktype::Type{T} + domain + handle::H + processor::P + scope::S + space::M +end diff --git a/src/types/memory-space.jl b/src/types/memory-space.jl new file mode 100644 index 000000000..247ceccb0 --- /dev/null +++ b/src/types/memory-space.jl @@ -0,0 +1 @@ +abstract type MemorySpace end \ No newline at end of file diff --git a/src/types/processor.jl b/src/types/processor.jl new file mode 100644 index 000000000..1e333413f --- /dev/null +++ b/src/types/processor.jl @@ -0,0 +1,2 @@ +# Docstring for Processor is attached in src/processor.jl after OSProc is defined (avoids "Replacing docs" warning). +abstract type Processor end \ No newline at end of file diff --git a/src/types/scope.jl b/src/types/scope.jl new file mode 100644 index 000000000..0197fddf9 --- /dev/null +++ b/src/types/scope.jl @@ -0,0 +1 @@ +abstract type AbstractScope end \ No newline at end of file diff --git a/src/utils/chunks.jl b/src/utils/chunks.jl index 9f0c3b487..cf577182d 100644 --- a/src/utils/chunks.jl +++ b/src/utils/chunks.jl @@ -152,7 +152,9 @@ new `Chunk`. All other kwargs are passed directly to `MemPool.poolset`. """ -function tochunk(x::X, proc::P=OSProc(), scope::S=AnyScope(); device=nothing, rewrap=false, kwargs...) where {X,P,S} +# N.B. No default for `proc` here: the 1-arg `tochunk(x)` is defined in +# tochunk.jl and routes through the current acceleration. +function tochunk(x::X, proc::P, scope::S=AnyScope(); device=nothing, rewrap=false, kwargs...) where {X,P,S} if device === nothing device = if Sch.walk_storage_safe(x) MemPool.GLOBAL_DEVICE[] @@ -161,7 +163,8 @@ function tochunk(x::X, proc::P=OSProc(), scope::S=AnyScope(); device=nothing, re end end ref = poolset(x; device, kwargs...) - Chunk{X,typeof(ref),P,S}(X, domain(x), ref, proc, scope) + space = memory_space(proc) + Chunk{X,typeof(ref),P,S,typeof(space)}(X, domain(x), ref, proc, scope, space) end function tochunk(x::Chunk, proc=nothing, scope=nothing; rewrap=false, kwargs...) if rewrap @@ -185,5 +188,6 @@ function savechunk(data, dir, f) fr = FileRef(f, sz) proc = OSProc() scope = AnyScope() # FIXME: Scoped to this node - Chunk{typeof(data),typeof(fr),typeof(proc),typeof(scope)}(typeof(data), domain(data), fr, proc, scope, true) + space = memory_space(proc) + Chunk{typeof(data),typeof(fr),typeof(proc),typeof(scope),typeof(space)}(typeof(data), domain(data), fr, proc, scope, space) end diff --git a/src/utils/dagdebug.jl b/src/utils/dagdebug.jl index fe9c1d2b9..ed4749d82 100644 --- a/src/utils/dagdebug.jl +++ b/src/utils/dagdebug.jl @@ -59,7 +59,7 @@ macro opcounter(category, count=1) @gensym old opcounter_sym = Symbol(:OPCOUNTER_, cat_sym) if !isdefined(__module__, opcounter_sym) - __module__.eval(:(#=const=# $opcounter_sym = OpCounter())) + __module__.eval(:(#=const=# $opcounter_sym = $OpCounter())) end esc(quote if $(QuoteNode(cat_sym)) in $OPCOUNTER_CATEGORIES @@ -70,4 +70,7 @@ macro opcounter(category, count=1) end end) end -opcounter(mod::Module, category::Symbol) = getfield(mod, Symbol(:OPCOUNTER_, category)).value[] \ No newline at end of file +opcounter(mod::Module, category::Symbol) = getfield(mod, Symbol(:OPCOUNTER_, category)).value[] + +# No-op debug helper for tracking largest values (used alongside @opcounter) +largest_value_update!(::Any) = nothing \ No newline at end of file diff --git a/src/utils/haloarray.jl b/src/utils/haloarray.jl index 32867e973..a9bb4a0d4 100644 --- a/src/utils/haloarray.jl +++ b/src/utils/haloarray.jl @@ -195,18 +195,30 @@ function aliasing(A::HaloArray) return CombinedAliasing([aliasing(A.center), map(aliasing, A.halos)...]) end memory_space(A::HaloArray) = memory_space(A.center) - -function move_rewrap(cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, A::HaloArray) - center_chunk = move_rewrap(cache, from_proc, to_proc, from_space, to_space, A.center) - halo_chunks = ntuple(i -> move_rewrap(cache, from_proc, to_proc, from_space, to_space, A.halos[i]), length(A.halos)) - halo_width = A.halo_width - own_center = A.own_center - to_w = root_worker_id(to_proc) - return remotecall_fetch(to_w, from_proc, to_proc, from_space, to_space, center_chunk, halo_chunks, halo_width, own_center) do from_proc, to_proc, from_space, to_space, center_chunk, halo_chunks, halo_width, own_center - center_new = unwrap(center_chunk) - halos_new = ntuple(i -> unwrap(halo_chunks[i]), length(halo_chunks)) - return tochunk(HaloArray(center_new, halos_new, halo_width; own_center=own_center), to_proc) - end +# A HaloArray's chunk record must be labeled with the memory space of its backing +# storage. `value_memory_space` is what `tochunk` uses to label a task result, and +# GPU extensions only define it for device array types (CuArray, CLArray, ...); a +# GPU-backed HaloArray would otherwise fall back to the generic CPU space, so its +# chunk would be mislabeled `CPURAMMemorySpace` while wrapping GPU arrays. That +# mislabeling breaks datadeps aliasing/move tracking on every GPU backend (the +# "Aliasing mismatch!" assertion in the stencil suite). Delegate to the center. +value_memory_space(A::HaloArray) = value_memory_space(A.center) + +# Header+children: transfer center and each halo, rebuild with halo_width +move_rewrap_parts(A::HaloArray) = ((A.center, A.halos...), A.halo_width) +function move_rewrap_build(::Type{<:HaloArray}, children, halo_width) + return HaloArray(children[1], children[2:end], halo_width) +end +function move_rewrap_child_types(::Type{HA}) where {HA<:HaloArray} + A = HA.parameters[3] + H = HA.parameters[4] + return (A, H.parameters...) +end +move_rewrap_header_mode(::Type{<:HaloArray}) = :broadcast +function move_rewrap_result_type(::Type{<:HaloArray}, child_cts, halo_width) + A = child_cts[1] + H = Tuple{child_cts[2:end]...} + return HaloArray{eltype(A),ndims(A),A,H} end function Dagger.unsafe_free!(A::HaloArray) diff --git a/src/utils/span_copy.jl b/src/utils/span_copy.jl new file mode 100644 index 000000000..c1f877f56 --- /dev/null +++ b/src/utils/span_copy.jl @@ -0,0 +1,338 @@ +### Multi-span copy (CPU loop or KernelAbstractions) + +# Copy many byte spans in one shot. Absolute pointers are converted to element +# offsets from each parent base; lengths must be multiples of sizeof(T). On +# Arrays this is a tight unsafe_copyto! loop; on GPU arrays it is a single +# KernelAbstractions launch over the packed element count. +# +# Argument order is destination-first everywhere in this file, matching Base's +# `copyto!(dst, src)`: `(dst, src)`, `(dst_ptrs, src_ptrs)`, `(to, from)` and +# `(to_ainfo, from_ainfo)`. + +@kernel function _multi_span_copy_kernel!(dst, src, dst_offs, src_offs, prefix) + i = @index(Global, Linear) + # prefix[s] < i <= prefix[s+1], with prefix[1] == 0 + lo = 1 + hi = length(prefix) - 1 + s = 1 + @inbounds while lo <= hi + mid = (lo + hi) >>> 1 + if i <= prefix[mid] + hi = mid - 1 + elseif i > prefix[mid + 1] + lo = mid + 1 + else + s = mid + break + end + end + @inbounds begin + local_i = i - prefix[s] + dst[dst_offs[s] + local_i] = src[src_offs[s] + local_i] + end +end + +storage_array(x) = x +storage_array(x::AbstractArray) = x +storage_array(x::SubArray) = storage_array(parent(x)) +storage_array(x::UpperTriangular) = storage_array(parent(x)) +storage_array(x::LowerTriangular) = storage_array(parent(x)) +storage_array(x::UnitUpperTriangular) = storage_array(parent(x)) +storage_array(x::UnitLowerTriangular) = storage_array(parent(x)) +storage_array(x::Transpose) = storage_array(parent(x)) +storage_array(x::Adjoint) = storage_array(parent(x)) +storage_array(x::Diagonal) = storage_array(parent(x)) + +# Narrow, internal helper: pick which of two already-unwrapped dense storage +# arrays (at least one of which is a GPU array) to take the KA backend from. +# This is *not* used to decide whether the multi-span path is needed — that +# decision keys off the memory space (`is_device_space`), which stays correct +# for composite objects whose own type is not a GPU array type. +is_device_array(x) = x isa GPUArraysCore.AbstractGPUArray + +# Whether `space` is device (VRAM) memory. Uses the existing per-backend +# `gpu_memory_kind` signal (`:CPU` for host RAM, `:CUDA`/`:ROC`/… for devices), +# so it is authoritative even for wrappers like a `HaloArray` holding device +# buffers, or views of device arrays, whose element/container type is not itself +# an `AbstractGPUArray`. +is_device_space(space::MemorySpace) = gpu_memory_kind(space) !== :CPU + +function _ptr_to_elem_off(base::UInt64, ptr::UInt64, elsize::Int) + delta = ptr - base + @assert delta % elsize == 0 "Span pointer is not element-aligned: $ptr vs base $base (elsize $elsize)" + return UInt32(delta ÷ elsize) # 0-based +end + +function _build_span_descriptors(dst::AbstractArray{T}, src::AbstractArray{T}, + dst_ptrs, src_ptrs, lens) where T + n = length(dst_ptrs) + @assert n == length(src_ptrs) == length(lens) + elsize = sizeof(T) + dst_vec = reshape(dst, :) + src_vec = reshape(src, :) + dst_base = UInt64(pointer(dst_vec)) + src_base = UInt64(pointer(src_vec)) + dst_offs = Vector{UInt32}(undef, n) + src_offs = Vector{UInt32}(undef, n) + prefix = Vector{UInt32}(undef, n + 1) + prefix[1] = 0 + for i in 1:n + len = UInt64(lens[i]) + @assert len % elsize == 0 "Span length is not an integer multiple of the element size: $len / $elsize" + ne = UInt32(len ÷ elsize) + dst_offs[i] = _ptr_to_elem_off(dst_base, UInt64(dst_ptrs[i]), elsize) + src_offs[i] = _ptr_to_elem_off(src_base, UInt64(src_ptrs[i]), elsize) + prefix[i + 1] = prefix[i] + ne + end + return dst_vec, src_vec, dst_offs, src_offs, prefix +end + +function _device_u32(backend, data::Vector{UInt32}) + d = KernelAbstractions.allocate(backend, UInt32, length(data)) + copyto!(d, data) + return d +end + +""" + multi_span_copy!(dst, src, dst_ptrs, src_ptrs, lens) + +Copy `lens[i]` bytes from `src` at absolute pointer `src_ptrs[i]` into `dst` at +`dst_ptrs[i]`, for each span `i`. `dst` and `src` must share element type `T` +and be the dense storage backing those pointers (see `storage_array`). +""" +function multi_span_copy!(dst::AbstractArray{T}, src::AbstractArray{T}, + dst_ptrs, src_ptrs, lens) where T + isempty(dst_ptrs) && return + dst_vec, src_vec, dst_offs, src_offs, prefix = _build_span_descriptors(dst, src, dst_ptrs, src_ptrs, lens) + total = Int(prefix[end]) + total == 0 && return + + if dst isa Array && src isa Array + GC.@preserve dst_vec src_vec begin + for i in eachindex(dst_offs) + n = Int(prefix[i + 1] - prefix[i]) + n == 0 && continue + unsafe_copyto!(pointer(dst_vec, Int(dst_offs[i]) + 1), + pointer(src_vec, Int(src_offs[i]) + 1), n) + end + end + return + end + + backend = KernelAbstractions.get_backend(is_device_array(dst) ? dst : src) + dst_offs_d = _device_u32(backend, dst_offs) + src_offs_d = _device_u32(backend, src_offs) + prefix_d = _device_u32(backend, prefix) + kern = _multi_span_copy_kernel!(backend) + kern(dst_vec, src_vec, dst_offs_d, src_offs_d, prefix_d; ndrange=total) + # This synchronize is required (independent of any following DtoH): the + # kernel reads the device-resident descriptor arrays allocated just above + # (`dst_offs_d`/`src_offs_d`/`prefix_d`), which become GC-eligible on return. + # Without the sync they could be freed while the kernel is still in flight + # (use-after-free). Letting the copy run stream-ordered/async with other + # kernels would require descriptors with caller-managed, stream-scoped + # lifetime; that is deferred for now. + KernelAbstractions.synchronize(backend) + return +end + +# `spans[i] == (src_span, dst_span)` (source first, matching RemainderAliasing). +function multi_span_copy!(dst::AbstractArray{T}, src::AbstractArray{T}, + spans::Vector{Tuple{LocalMemorySpan,LocalMemorySpan}}) where T + n = length(spans) + dst_ptrs = Vector{UInt64}(undef, n) + src_ptrs = Vector{UInt64}(undef, n) + lens = Vector{UInt64}(undef, n) + for i in eachindex(spans) + src_span, dst_span = spans[i] + @assert src_span.len == dst_span.len + dst_ptrs[i] = dst_span.ptr + src_ptrs[i] = src_span.ptr + lens[i] = src_span.len + end + return multi_span_copy!(dst, src, dst_ptrs, src_ptrs, lens) +end + +# Gather `src`'s spans into a contiguous `dst` starting at element 1 / byte 0. +# `src_spans` may be any iterable of spans (e.g. a lazy `Iterators.map`). +# +# N.B. This is an internal helper (distinct name), *not* a `multi_span_gather!` +# method: when `T === UInt8` a data `Vector{UInt8}` would be indistinguishable +# from the `buf::Vector{UInt8}` byte-buffer method below, so the buffer method's +# internal call would re-dispatch to itself and recurse forever. +function _multi_span_gather_packed!(dst::AbstractArray{T}, src::AbstractArray{T}, src_spans) where T + n = length(src_spans) + dst_ptrs = Vector{UInt64}(undef, n) + src_ptrs = Vector{UInt64}(undef, n) + lens = Vector{UInt64}(undef, n) + dst_base = UInt64(pointer(reshape(dst, :))) + offset = UInt64(0) + for (i, span) in enumerate(src_spans) + len = UInt64(span_len(span)) + dst_ptrs[i] = dst_base + offset + src_ptrs[i] = UInt64(span_start(span)) + lens[i] = len + offset += len + end + return multi_span_copy!(dst, src, dst_ptrs, src_ptrs, lens) +end + +# Scatter a contiguous `src` into `dst`'s spans. Internal helper (distinct name) +# for the same reason as `_multi_span_gather_packed!` above. +function _multi_span_scatter_packed!(dst::AbstractArray{T}, src::AbstractArray{T}, dst_spans) where T + n = length(dst_spans) + dst_ptrs = Vector{UInt64}(undef, n) + src_ptrs = Vector{UInt64}(undef, n) + lens = Vector{UInt64}(undef, n) + src_base = UInt64(pointer(reshape(src, :))) + offset = UInt64(0) + for (i, span) in enumerate(dst_spans) + len = UInt64(span_len(span)) + dst_ptrs[i] = UInt64(span_start(span)) + src_ptrs[i] = src_base + offset + lens[i] = len + offset += len + end + return multi_span_copy!(dst, src, dst_ptrs, src_ptrs, lens) +end + +# Whether `x`'s storage is a single, pointer-addressable dense buffer that the +# packed (`unsafe_copyto!` / KA) span path can operate on: a host `Array`, or a +# device array (VRAM buffer, or a wrapper whose memory space is a device). A host +# `AbstractArray` that is *not* an `Array` (e.g. `SparseMatrixCSC`, whose data +# lives across separate `nzval`/`colptr`/`rowval` buffers) is not, and must use +# the per-span remainder path, which resolves each span to its owning buffer. +_span_dense_storage(x) = + x isa Array || is_device_array(x) || is_device_space(memory_space(x)) + +# Per-span remainder gather/scatter: resolve each span to the concrete buffer +# holding it (see `find_object_holding_ptr`) and copy field-by-field. Used for +# objects without a single dense backing buffer (RefValue, GPURef, +# SparseMatrixCSC, etc.). +function _gather_spans_remainder!(buf::Vector{UInt8}, src, src_spans) + offset = UInt64(1) + for span in src_spans + ptr = UInt64(span_start(span)) + len = UInt64(span_len(span)) + read_remainder!(buf, offset, src, ptr, len) + offset += len + end + return +end +function _scatter_spans_remainder!(dst, buf::Vector{UInt8}, dst_spans) + offset = UInt64(1) + for span in dst_spans + ptr = UInt64(span_start(span)) + len = UInt64(span_len(span)) + write_remainder!(buf, offset, dst, ptr, len) + offset += len + end + return +end + +# Host UInt8 buffer gather/scatter used by remainder and MPI packing paths. +function multi_span_gather!(buf::Vector{UInt8}, src::AbstractArray{T}, src_spans) where T + src_s = storage_array(src) + # Host `AbstractArray` without a single dense buffer (e.g. SparseMatrixCSC): + # fall back to the per-span remainder path rather than the packed/KA path, + # which would call `pointer` on a non-pointer-addressable array. + _span_dense_storage(src_s) || return _gather_spans_remainder!(buf, src, src_spans) + elsize = sizeof(T) + nbytes = sum(s -> UInt64(span_len(s)), src_spans; init=UInt64(0)) + @assert UInt64(length(buf)) >= nbytes + @assert nbytes % elsize == 0 + nelem = Int(nbytes ÷ elsize) + host = unsafe_wrap(Vector{T}, Ptr{T}(pointer(buf)), nelem) + if src_s isa Array + _multi_span_gather_packed!(host, src_s, src_spans) + return + end + with_context!(memory_space(src_s)) + backend = KernelAbstractions.get_backend(src_s) + packed = KernelAbstractions.allocate(backend, T, nelem) + _multi_span_gather_packed!(packed, src_s, src_spans) + copyto!(host, packed) + return +end +function multi_span_gather!(buf::Vector{UInt8}, src, src_spans) + # RefValue, GPURef, SparseMatrixCSC, etc.: keep the per-span remainder path + return _gather_spans_remainder!(buf, src, src_spans) +end + +function multi_span_scatter!(dst::AbstractArray{T}, buf::Vector{UInt8}, dst_spans) where T + dst_s = storage_array(dst) + # Host `AbstractArray` without a single dense buffer (e.g. SparseMatrixCSC): + # fall back to the per-span remainder path (see `multi_span_gather!`). + _span_dense_storage(dst_s) || return _scatter_spans_remainder!(dst, buf, dst_spans) + elsize = sizeof(T) + nbytes = sum(s -> UInt64(span_len(s)), dst_spans; init=UInt64(0)) + @assert UInt64(length(buf)) >= nbytes + @assert nbytes % elsize == 0 + nelem = Int(nbytes ÷ elsize) + host = unsafe_wrap(Vector{T}, Ptr{T}(pointer(buf)), nelem) + if dst_s isa Array + _multi_span_scatter_packed!(dst_s, host, dst_spans) + return + end + with_context!(memory_space(dst_s)) + backend = KernelAbstractions.get_backend(dst_s) + packed = KernelAbstractions.allocate(backend, T, nelem) + copyto!(packed, host) + _multi_span_scatter_packed!(dst_s, packed, dst_spans) + return +end +function multi_span_scatter!(dst, buf::Vector{UInt8}, dst_spans) + return _scatter_spans_remainder!(dst, buf, dst_spans) +end + +# Whether a `move!` must take the multi-span (KA) path instead of plain +# `copyto!`. `copyto!` is fine for host memory (any AbstractArray) and for +# device-to-device between contiguous `DenseArray`s; anything else on device +# memory (views, triangular/diagonal wrappers, composite device containers, …) +# needs the span-based KA copy, since `copyto!` would require illegal scalar +# indexing. Device-ness comes from the memory space, not the array type. +function needs_multi_span_copy(to_space::MemorySpace, from_space::MemorySpace, + to::AbstractArray, from::AbstractArray) + device = is_device_space(to_space) || is_device_space(from_space) + return device && !(to isa DenseArray && from isa DenseArray) +end + +# Copy `from`'s spans into `to`'s spans (matched pairwise by order). `to_ainfo`/ +# `from_ainfo` select which spans to move and default to the whole arrays. +function multi_span_move!(to::AbstractArray{T}, from::AbstractArray{T}, + to_ainfo=to, from_ainfo=from) where T + to_spans = memory_spans(to_ainfo) + from_spans = memory_spans(from_ainfo) + @assert length(from_spans) == length(to_spans) "Mismatched span counts for multi-span move: $(length(from_spans)) vs $(length(to_spans))" + n = length(from_spans) + dst_ptrs = Vector{UInt64}(undef, n) + src_ptrs = Vector{UInt64}(undef, n) + lens = Vector{UInt64}(undef, n) + for i in 1:n + @assert span_len(from_spans[i]) == span_len(to_spans[i]) + dst_ptrs[i] = span_start(to_spans[i]) + src_ptrs[i] = span_start(from_spans[i]) + lens[i] = span_len(from_spans[i]) + end + return multi_span_copy!(storage_array(to), storage_array(from), dst_ptrs, src_ptrs, lens) +end + +# GPU-aware AbstractArray move! (CPU spaces). Device-space methods in GPU +# extensions call needs_multi_span_copy / multi_span_move! the same way. +function move!(to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} + if needs_multi_span_copy(to_space, from_space, to, from) + multi_span_move!(to, from) + else + copyto!(to, from) + end + return +end + +function move!(::Type{<:Diagonal}, to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} + if is_device_space(to_space) || is_device_space(from_space) + multi_span_move!(to, from, aliasing(to, Diagonal), aliasing(from, Diagonal)) + else + copyto!(view(to, diagind(to)), view(from, diagind(from))) + end + return +end diff --git a/src/weakchunk.jl b/src/weakchunk.jl new file mode 100644 index 000000000..fdd688171 --- /dev/null +++ b/src/weakchunk.jl @@ -0,0 +1,32 @@ +struct WeakChunk + wid::Int + id::Int + x::WeakRef + # Cache the referent's chunktype: `chunktype(::WeakChunk)` is used by the + # scheduler's post-completion cost bookkeeping (`signature`), which can run + # after datadeps has already freed the underlying chunk (so the `WeakRef` + # would be expired). The type is fixed for the chunk's lifetime, so caching + # it keeps that path working without pinning the chunk alive. + chunktype::Type +end + +function WeakChunk(c::Chunk) + return WeakChunk(c.handle.owner, c.handle.id, WeakRef(c), chunktype(c)) +end + +unwrap_weak(c::WeakChunk) = c.x.value +function unwrap_weak_checked(c::WeakChunk) + cw = unwrap_weak(c) + # A `Nothing` chunktype marks a non-resident placeholder (e.g. an MPIRef + # chunk for data owned by another rank): its wrapper legitimately may have + # been GC'd, so only assert liveness for chunks that carry real data. + @assert chunktype(c) === Nothing || cw !== nothing "WeakChunk expired: ($(c.wid), $(c.id))" + return cw +end +wrap_weak(c::Chunk) = WeakChunk(c) +isweak(c::WeakChunk) = true +isweak(c::Chunk) = false +is_task_or_chunk(c::WeakChunk) = true +Serialization.serialize(io::AbstractSerializer, wc::WeakChunk) = + error("Cannot serialize a WeakChunk") +chunktype(c::WeakChunk) = c.chunktype From 5f05669e2ef7dfffe2307631425119ee0d623e97 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 19 Jul 2026 09:57:30 -0700 Subject: [PATCH 2/6] tests/CI: MPI CPU test suite, GPU suites, and GHA MPI job Adds the CPU MPI test suite (`test/mpi.jl`) and a shared parametrized GPU MPI suite (`test/mpi_gpu_suite.jl`) with thin per-backend wrappers (CUDA/ROCm/OpenCL/Metal/oneAPI). `test/mpienv` + `test/run_mpi.jl` launch MPI tests via MPI.jl's `mpiexec` (no system MPI needed), wired into a new `mpi-cpu` GHA job matrix (Julia 1.10/1, 2/4 ranks) alongside the existing CI jobs. Adds per-backend GPU test envs and processor-grid tests. Co-authored-by: Felipe Tome Co-authored-by: Cursor --- .github/workflows/CI.yml | 32 ++ test/array/allocation.jl | 3 + test/array/core.jl | 13 +- test/cudaenv/Project.toml | 10 + test/datadeps.jl | 34 +- test/gpu.jl | 14 +- test/metalenv/Project.toml | 10 + test/mpi.jl | 717 ++++++++++++++++++++++++++++++++++++ test/mpi_cholesky.jl | 25 ++ test/mpi_cuda.jl | 22 ++ test/mpi_gpu_suite.jl | 249 +++++++++++++ test/mpi_metal.jl | 25 ++ test/mpi_oneapi.jl | 28 ++ test/mpi_opencl.jl | 25 ++ test/mpi_rocm.jl | 45 +++ test/mpienv/Project.toml | 9 + test/oneapienv/Project.toml | 10 + test/openclenv/Project.toml | 10 + test/procgrid.jl | 103 ++++++ test/rocmenv/Project.toml | 10 + test/run_mpi.jl | 25 ++ test/runtests.jl | 1 + test/scheduler.jl | 4 +- 23 files changed, 1408 insertions(+), 16 deletions(-) create mode 100644 test/cudaenv/Project.toml create mode 100644 test/metalenv/Project.toml create mode 100644 test/mpi.jl create mode 100644 test/mpi_cholesky.jl create mode 100644 test/mpi_cuda.jl create mode 100644 test/mpi_gpu_suite.jl create mode 100644 test/mpi_metal.jl create mode 100644 test/mpi_oneapi.jl create mode 100644 test/mpi_opencl.jl create mode 100644 test/mpi_rocm.jl create mode 100644 test/mpienv/Project.toml create mode 100644 test/oneapienv/Project.toml create mode 100644 test/openclenv/Project.toml create mode 100644 test/procgrid.jl create mode 100644 test/rocmenv/Project.toml create mode 100644 test/run_mpi.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index eed72ba58..0b25e5576 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -231,3 +231,35 @@ jobs: } else { await github.rest.issues.createComment({ owner, repo, issue_number, body }); } + + mpi-cpu: + name: MPI CPU - Julia ${{ matrix.version }} / ${{ matrix.ranks }} ranks + runs-on: ubuntu-latest + timeout-minutes: 60 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') }} + strategy: + fail-fast: false + matrix: + version: + - '1.10' # LTS + - '1.11' + - '1' # latest stable + ranks: + - 2 + - 4 + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + arch: x64 + - uses: julia-actions/cache@v3 + # test/mpienv pulls MPI (MPICH_jll by default) plus the checked-out Dagger; + # no system MPI installation is required. We `develop` the checkout rather + # than relying on the `[sources]` entry because `[sources]` is a Julia 1.11+ + # feature and is silently ignored on 1.10 (LTS), which would otherwise + # resolve a stale registered Dagger lacking the MPI API (`accelerate!`). + - name: Instantiate CPU MPI environment + run: julia --project=test/mpienv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + - name: Run CPU MPI tests + run: julia --project=test/mpienv test/run_mpi.jl ${{ matrix.ranks }} 2 test/mpi.jl diff --git a/test/array/allocation.jl b/test/array/allocation.jl index 0173eddfa..83dba36d2 100644 --- a/test/array/allocation.jl +++ b/test/array/allocation.jl @@ -235,6 +235,9 @@ end function chunk_processors(Ad::DArray) [Dagger.processor(fetch(Ad.chunks[idx]; raw=true)) for idx in CartesianIndices(size(domainchunks(Ad)))] end +function tile_processors(proc_grid::Dagger.AbstractProcGrid{N}, block_grid::Tuple{Vararg{Int,N}}) where N + return [Dagger.procgrid_processor(proc_grid, I) for I in CartesianIndices(block_grid)] +end function tile_processors(proc_grid::AbstractArray{<:Dagger.Processor,N}, block_grid::Tuple{Vararg{Int,N}}) where N reps = Int.(ceil.(block_grid ./ size(proc_grid))) tiled = repeat(proc_grid, reps...) diff --git a/test/array/core.jl b/test/array/core.jl index 0cdcb0456..e0e4d698f 100644 --- a/test/array/core.jl +++ b/test/array/core.jl @@ -116,15 +116,10 @@ end using MemPool -@testset "affinity" begin +@testset "datasize" begin x = Dagger.tochunk([1:10;]) - aff = Dagger.affinity(x) - @test aff[1] == Dagger.OSProc(myid()) - @test aff[2] == sizeof(Int)*10 + @test Dagger.datasize(x) == sizeof(Int)*10 @test Dagger.tochunk(x) === x - f = MemPool.FileRef("/tmp/d", aff[2]) - aff = Dagger.affinity(f) - #@test length(aff) == 3 - @test (aff[1]).pid in procs() - @test aff[2] == sizeof(Int)*10 + f = MemPool.FileRef("/tmp/d", Dagger.datasize(x)) + @test Dagger.datasize(f) == sizeof(Int)*10 end diff --git a/test/cudaenv/Project.toml b/test/cudaenv/Project.toml new file mode 100644 index 000000000..6e245a12a --- /dev/null +++ b/test/cudaenv/Project.toml @@ -0,0 +1,10 @@ +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Dagger = {path = "../.."} diff --git a/test/datadeps.jl b/test/datadeps.jl index 291e10e70..3d4bdbe56 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -61,6 +61,33 @@ end end end + @testset "View of View" begin + outer = view(chunk1, 2:5, 3:6) + nested = view(outer, 1:2, 2:3) + direct = view(chunk1, 2:3, 4:5) + @test nested isa ChunkView + @test nested.chunk === chunk1 + @test nested.slices == direct.slices + @test fetch(nested) == fetch(direct) + + # Colon parent dims pass sub-slices through + outer_colon = view(chunk1, :, 2:5) + nested_colon = view(outer_colon, 3:4, 1:2) + @test nested_colon.slices == view(chunk1, 3:4, 2:3).slices + + # Int parent dim is dropped; nested view indexes remaining dims only + outer_drop = view(chunk1, 3, 1:8) + nested_drop = view(outer_drop, 2:5) + @test nested_drop.slices == view(chunk1, 3, 2:5).slices + @test fetch(nested_drop) == fetch(view(chunk1, 3, 2:5)) + + @test_throws DimensionMismatch view(outer, :) + @test_throws DimensionMismatch view(outer, :, :, :) + @test_throws ArgumentError view(outer, 1:5, 1:2) # out of outer range → composed OOB + @test_throws ArgumentError view(outer_drop, 0) + @test_throws ArgumentError view(outer_drop, 9) + end + @test fetch(v1) == fetch(chunk1) @test Dagger.memory_space(v1) == Dagger.memory_space(chunk1) @@ -90,15 +117,16 @@ end end function test_move_rewrap_aliasing(obj, dest_space) + accel = Dagger.current_acceleration() src_space = Dagger.memory_space(obj) from_proc = first(Dagger.processors(src_space)) to_proc = first(Dagger.processors(dest_space)) # move_rewrap like generate_slot! - dummy_backing = Dagger.tochunk(Dagger.AliasedObjectCacheStore()) - cache = Dagger.AliasedObjectCache(dest_space, dummy_backing) + dummy_backing = Dagger.tochunk(Dagger.AliasedObjectCacheStore(accel)) + cache = Dagger.AliasedObjectCache(accel, dest_space, dummy_backing) - dest_obj_chunk = Dagger.move_rewrap(cache, from_proc, to_proc, src_space, dest_space, obj) + dest_obj_chunk = Dagger.remotecall_endpoint_toplevel(Dagger.move_rewrap, accel, cache, from_proc, to_proc, src_space, dest_space, obj) # VERIFICATION: Check that source and destination have compatible memory spans # Use the chunk directly for aliasing so it handles remote workers correctly diff --git a/test/gpu.jl b/test/gpu.jl index 61db4745c..474a0bd08 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -808,5 +808,17 @@ end end end end + + @testset "Datadeps" begin + local_scope = Dagger.scope(worker=1, cl_device=1) + A = rand(Float32, 8, 8) + ref = copy(A) + Dagger.with_options(;scope=local_scope) do + Dagger.spawn_datadeps() do + Dagger.@spawn addarray!(Dagger.InOut(A)) + end + end + @test A ≈ ref .+ 1 + end end -end \ No newline at end of file +end diff --git a/test/metalenv/Project.toml b/test/metalenv/Project.toml new file mode 100644 index 000000000..b3d5e9632 --- /dev/null +++ b/test/metalenv/Project.toml @@ -0,0 +1,10 @@ +[deps] +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +Metal = "dde4c033-4e86-420c-a63e-0dd931031962" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Dagger = {path = "../.."} diff --git a/test/mpi.jl b/test/mpi.jl new file mode 100644 index 000000000..f5e7ffd6f --- /dev/null +++ b/test/mpi.jl @@ -0,0 +1,717 @@ +# Unified CPU MPI test suite (datadeps + linalg smokes). +# +# Acceleration-independent logic — aliasing algebra (will_alias, span math), +# dependency-ordering semantics (In/Out/InOut/Deps), remainder computation — +# is covered by the Distributed suite (test/datadeps.jl) and is intentionally +# NOT re-tested here. This suite covers only MPI-exclusive machinery: +# +# - SPMD-uniform chunk creation (deterministic placement, MPIRef handles) +# - check_uniform / compare_all agreement +# - acceleration-aware wrapping of non-Chunk arguments (rank-0 ownership) +# - the SPMD-symmetric endpoint (MPIWireValue move_rewrap, aliased-object +# cache sharing, rank-stamped aliasing spans) [mirrors "Aliased Object +# Copying"] +# - cross-rank copies at execution time (identity, view/remainder, and +# non-identity dep_mod transfers) +# - collect/fetch under uniform execution +# - result-type broadcast for untyped task results +# - Cholesky / LU DArray smokes +# +# GPU coverage lives in test/mpi_cuda.jl (CUDA) and test/mpi_rocm.jl (ROCm). +# +# Known gaps (not tested): NaiveScheduler/UltraScheduler use rank-local cost +# measurements and are not deterministic across ranks. +# +# Run: mpiexec -n 4 julia --project --threads=2 test/mpi.jl + +using Dagger, MPI, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +# MPI-specific types live in the MPIExt extension (loaded via `using MPI`), +# not in Dagger core. Reference them through the extension module. +const MPIExt = Base.get_extension(Dagger, :MPIExt) + +include(joinpath(@__DIR__, "util.jl")) + +Dagger.accelerate!(:mpi) +Dagger.check_uniformity!(true) +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) +const nranks = MPI.Comm_size(comm) +const accel = Dagger.current_acceleration() + +mpi_procs() = sort(collect(Dagger.get_processors(MPIExt.MPIClusterProc(comm))); + by=p->(p.rank, Dagger.short_name(p))) +proc_for_rank(r) = first(filter(p->p.rank == r, mpi_procs())) +space_for_rank(r) = MPIExt.MPIMemorySpace(Dagger.CPURAMMemorySpace(1), comm, r) +rank_scope(r) = Dagger.ExactScope(proc_for_rank(r)) + +# Rank pairs for slot tests (mirrors Distributed's worker pairs), including a +# same-rank pair to check that origin slots alias the original data +const rank_pairs = if nranks >= 3 + [(0, 1), (1, 0), (1, 2), (1, 1)] +else + [(0, nranks-1), (nranks-1, 0), (0, 0)] +end + +inc!(X) = (X .+= 1; nothing) +add1!(X) = (X .+= 1; nothing) +scale2!(X) = (X .*= 2; nothing) +sum_into!(r, X) = (r[] = Int(sum(X)); nothing) +function upper_double!(X) + for j in axes(X, 2), i in 1:j + X[i, j] *= 2 + end + return nothing +end +function diag_inc!(X) + for i in 1:minimum(size(X)) + X[i, i] += 1 + end + return nothing +end +function strict_upper_double!(X) + for j in axes(X, 2), i in 1:j-1 + X[i, j] *= 2 + end + return nothing +end +function lower_double!(X) + for j in axes(X, 2), i in j:size(X, 1) + X[i, j] *= 2 + end + return nothing +end +function strict_lower_double!(X) + for j in axes(X, 2), i in j+1:size(X, 1) + X[i, j] *= 2 + end + return nothing +end +untyped_result(x) = x > 0 ? Int(x) : Float64(x) +concrete_fail(x) = x > 0 ? error("Test concrete") : x + 1 # inferred Int, throws at runtime +nothrow_add(x) = x + 1 # concrete Int, proven nothrow → execute! skips status bcast +mut_ref!(R) = (R[] .= 0; nothing) +exec_rank() = MPI.Comm_rank(MPI.COMM_WORLD) # rank actually executing a task + +function fresh_cache(dst_space) + backing = Dagger._with_default_acceleration() do + Dagger.tochunk(Dagger.AliasedObjectCacheStore(accel)) + end + return Dagger.AliasedObjectCache(accel, dst_space, backing) +end + +# Generate a slot for `obj` (a Chunk in src space) in dst space, like +# generate_slot! does, and return it +function make_slot(cache, src_rank, dst_rank, obj) + src_space = space_for_rank(src_rank) + dst_space = space_for_rank(dst_rank) + return Dagger.remotecall_endpoint_toplevel(Dagger.move_rewrap, accel, cache, + proc_for_rank(src_rank), proc_for_rank(dst_rank), + src_space, dst_space, obj) +end + +@testset "MPI" begin + +@testset "accelerate! idempotency and switch" begin + @test Dagger.current_acceleration() isa MPIExt.MPIAcceleration + Dagger.accelerate!(:mpi) + @test Dagger.current_acceleration() isa MPIExt.MPIAcceleration + Dagger.accelerate!(:distributed) + @test Dagger.current_acceleration() isa Dagger.DistributedAcceleration + # Diverge default RNGs across ranks; re-init must re-sync them. + Random.seed!(UInt64(rank + 1)) + Dagger.accelerate!(:mpi) + @test Dagger.current_acceleration() isa MPIExt.MPIAcceleration + probe = rand(UInt64) + @test Dagger.check_uniform(probe) + MPI.Barrier(comm) +end + +@testset "Datadeps" begin + +@testset "Uniform chunk creation" begin + Random.seed!(1) + A = rand(16, 16) + # distribute path + DA = DArray(A, Blocks(4, 4)) + owners = Int[] + for t in vec(DA.chunks) + c = fetch(t; raw=true) + h = c.handle + @test h isa MPIExt.MPIRef + @test Dagger.check_uniform(h) + @test Dagger.check_uniform(c.space) + @test Dagger.check_uniform(c.processor) + # The owner (and only the owner) holds the payload + @test (h.rank == rank) == (h.innerRef !== nothing) + push!(owners, h.rank) + end + # Deterministic round-robin placement covers all ranks + @test sort(unique(owners)) == collect(0:nranks-1) + @test Dagger.check_uniform(hash(owners)) + + # alloc path + DZ = zeros(Blocks(4, 4), Float64, 16, 16) + for t in vec(DZ.chunks) + c = fetch(t; raw=true) + @test c.handle isa MPIExt.MPIRef + @test Dagger.check_uniform(c.handle) + end +end + +@testset "check_uniform" begin + @test Dagger.check_uniform(42) + @test Dagger.check_uniform(hash((1, :a, "x"))) + # Rank-dependent values must be detected on every rank + @test_throws ArgumentError Dagger.check_uniform(rank) + # The compare stream stays aligned after a detected failure + @test Dagger.check_uniform(7) +end + +@testset "Non-Chunk argument wrapping" begin + # Non-Chunk mutable arguments wrap as rank-0-owned MPIRef chunks + x = Ref(11) + c = Dagger.tochunk(x) + @test c.handle isa MPIExt.MPIRef + @test c.space isa MPIExt.MPIMemorySpace + @test c.space.rank == 0 + @test Dagger.check_uniform(c.handle) + @test (rank == 0) == (c.handle.innerRef !== nothing) + + # As datadeps arguments: rank 0's copy is the source of truth and is + # written back at region end, even when the task runs on another rank + r = Ref(0) + A = ones(4, 4) + exec_rank = min(1, nranks-1) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(exec_rank) sum_into!(InOut(r), In(A)) + end + if rank == 0 + @test r[] == 16 + end +end + +# Mirrors test/datadeps.jl "Aliased Object Copying", via the SPMD endpoint +@testset "Aliased Object Copying" begin + for (src, dst) in rank_pairs + @testset "Rank $src -> $dst" begin + @testset "Array" begin + Random.seed!(100 + 10*src + dst) + A = rand(4, 4) # replicated construction; owner rank holds the chunk payload + obj = Dagger.tochunk(A, proc_for_rank(src), space_for_rank(src)) + cache = fresh_cache(space_for_rank(dst)) + slot = make_slot(cache, src, dst, obj) + @test Dagger.memory_space(slot) == space_for_rank(dst) + @test Dagger.check_uniform(slot.handle) + @test Dagger.chunktype(slot) == Matrix{Float64} + + src_ainfo = Dagger.aliasing(accel, obj, identity) + dst_ainfo = Dagger.aliasing(accel, slot, identity) + ss, ds = Dagger.memory_spans(src_ainfo), Dagger.memory_spans(dst_ainfo) + @test length(ss) == length(ds) + @test all(Dagger.span_len(a) == Dagger.span_len(b) for (a, b) in zip(ss, ds)) + if src == dst + # Origin slots wrap (alias) the original data + @test Dagger.will_alias(src_ainfo, dst_ainfo) + else + # Cross-rank spans are stamped with their owner rank and + # must never alias, even with coinciding SPMD addresses + @test !Dagger.will_alias(src_ainfo, dst_ainfo) + end + # The slot holds a copy of the source data on the destination + if rank == dst + @test Dagger.unwrap(slot) == A + end + end + + @testset "SubArray + parent sharing" begin + Random.seed!(200 + 10*src + dst) + A = rand(8, 8) + vA = view(A, 1:4, 1:8) + vB = view(A, 5:8, 1:8) + objA = Dagger.tochunk(vA, proc_for_rank(src), space_for_rank(src)) + objB = Dagger.tochunk(vB, proc_for_rank(src), space_for_rank(src)) + cache = fresh_cache(space_for_rank(dst)) + slotA = make_slot(cache, src, dst, objA) + slotB = make_slot(cache, src, dst, objB) + @test Dagger.check_uniform(slotA.handle) + @test Dagger.check_uniform(slotB.handle) + @test Dagger.chunktype(slotA) <: SubArray + + aA = Dagger.aliasing(accel, slotA, identity) + aB = Dagger.aliasing(accel, slotB, identity) + # Both views must share one parent allocation in the + # destination space (aliased-object cache reuse) + @test aA isa Dagger.StridedAliasing + @test aB isa Dagger.StridedAliasing + @test aA.base_ptr == aB.base_ptr + # ... while the views themselves remain disjoint + @test !Dagger.will_alias(aA, aB) + if rank == dst + @test Dagger.unwrap(slotA) == vA + @test Dagger.unwrap(slotB) == vB + @test parent(Dagger.unwrap(slotA)) === parent(Dagger.unwrap(slotB)) + end + end + + @testset "UpperTriangular wrapper" begin + Random.seed!(300 + 10*src + dst) + A = rand(4, 4) + U = UpperTriangular(A) + obj = Dagger.tochunk(U, proc_for_rank(src), space_for_rank(src)) + cache = fresh_cache(space_for_rank(dst)) + slot = make_slot(cache, src, dst, obj) + @test Dagger.check_uniform(slot.handle) + @test Dagger.chunktype(slot) <: UpperTriangular + if rank == dst + @test Dagger.unwrap(slot) == U + end + end + + @testset "ChunkView" begin + Random.seed!(400 + 10*src + dst) + A = rand(8, 8) + obj = Dagger.tochunk(A, proc_for_rank(src), space_for_rank(src)) + cvA = view(obj, 1:4, 1:8) + cvB = view(obj, 5:8, 1:8) + @test cvA isa Dagger.ChunkView + cache = fresh_cache(space_for_rank(dst)) + slotA = make_slot(cache, src, dst, cvA) + slotB = make_slot(cache, src, dst, cvB) + @test Dagger.memory_space(slotA) == space_for_rank(dst) + @test Dagger.check_uniform(slotA.handle) + @test Dagger.check_uniform(slotB.handle) + @test Dagger.chunktype(slotA) <: SubArray + + src_ainfo = Dagger.aliasing(accel, cvA, identity) + dst_ainfo = Dagger.aliasing(accel, slotA, identity) + ss, ds = Dagger.memory_spans(src_ainfo), Dagger.memory_spans(dst_ainfo) + @test length(ss) == length(ds) + @test all(Dagger.span_len(a) == Dagger.span_len(b) for (a, b) in zip(ss, ds)) + if src == dst + @test Dagger.will_alias(src_ainfo, dst_ainfo) + else + @test !Dagger.will_alias(src_ainfo, dst_ainfo) + end + + # Both views share one backing-chunk allocation on the + # destination (aliased-object cache reuse), but don't overlap + aA = Dagger.aliasing(accel, slotA, identity) + aB = Dagger.aliasing(accel, slotB, identity) + @test aA isa Dagger.StridedAliasing + @test aA.base_ptr == aB.base_ptr + @test !Dagger.will_alias(aA, aB) + if rank == dst + @test Dagger.unwrap(slotA) == view(A, 1:4, 1:8) + @test Dagger.unwrap(slotB) == view(A, 5:8, 1:8) + @test parent(Dagger.unwrap(slotA)) === parent(Dagger.unwrap(slotB)) + end + + # Nested ChunkView flattens to the same slices as a direct view + cv_nested = view(cvA, 2:3, 1:4) + cv_direct = view(obj, 2:3, 1:4) + @test cv_nested isa Dagger.ChunkView + @test cv_nested.chunk === obj + @test cv_nested.slices == cv_direct.slices + slot_nested = make_slot(fresh_cache(space_for_rank(dst)), src, dst, cv_nested) + slot_direct = make_slot(fresh_cache(space_for_rank(dst)), src, dst, cv_direct) + @test Dagger.chunktype(slot_nested) <: SubArray + if rank == dst + @test Dagger.unwrap(slot_nested) == Dagger.unwrap(slot_direct) + @test Dagger.unwrap(slot_nested) == view(A, 2:3, 1:4) + end + end + + @testset "HaloArray" begin + H = Dagger.HaloArray{Int,2}((4, 4), (1, 1)) + fill!(H.center, 7) + foreach(h->fill!(h, -1), H.halos) + obj = Dagger.tochunk(H, proc_for_rank(src), space_for_rank(src)) + cache = fresh_cache(space_for_rank(dst)) + slot = make_slot(cache, src, dst, obj) + @test Dagger.check_uniform(slot.handle) + @test Dagger.chunktype(slot) <: Dagger.HaloArray + + src_ainfo = Dagger.aliasing(accel, obj, identity) + dst_ainfo = Dagger.aliasing(accel, slot, identity) + ss, ds = Dagger.memory_spans(src_ainfo), Dagger.memory_spans(dst_ainfo) + @test length(ss) == length(ds) + @test all(Dagger.span_len(a) == Dagger.span_len(b) for (a, b) in zip(ss, ds)) + if src != dst + @test !Dagger.will_alias(src_ainfo, dst_ainfo) + end + if rank == dst + H2 = Dagger.unwrap(slot) + @test H2.center == H.center + @test all(h2 == h for (h2, h) in zip(H2.halos, H.halos)) + end + end + end + end +end + +@testset "Cross-rank aliasing isolation" begin + # Same-shaped arrays owned by different ranks must never alias + Random.seed!(4) + A = rand(4, 4) + B = rand(4, 4) + r2 = min(1, nranks-1) + cA = Dagger.tochunk(A, proc_for_rank(0), space_for_rank(0)) + cB = Dagger.tochunk(B, proc_for_rank(r2), space_for_rank(r2)) + aA = Dagger.aliasing(accel, cA, identity) + aB = Dagger.aliasing(accel, cB, identity) + @test Dagger.check_uniform(aA) + @test Dagger.check_uniform(aB) + @test !Dagger.will_alias(aA, aB) +end + +# Numeric dataflow across ranks: full-array, view (remainder copies), and +# triangular-dep_mod (non-identity cross-rank move!) writes, with sequential +# (program-order) semantics +@testset "Cross-rank dataflow" begin + Random.seed!(7) + A = rand(8, 8) + ref = copy(A) + vA = view(A, 3:6, 3:6) + + r1 = min(1, nranks-1) + r2 = min(2, nranks-1) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) add1!(InOut(A)) + Dagger.@spawn scope=rank_scope(r2) scale2!(InOut(vA)) + Dagger.@spawn scope=rank_scope(0) upper_double!(Deps(A, InOut(UpperTriangular))) + end + + # Serial reference, in program order + ref .+= 1 + view(ref, 3:6, 3:6) .*= 2 + upper_double!(ref) + + # Rank 0 owns the original (non-Chunk) argument and receives the write-back + if rank == 0 + @test A ≈ ref + end +end + +@testset "ChunkView dataflow" begin + # Views over a DArray block as datadeps arguments, mixed with whole-chunk + # access (exercises cross-rank remainder copies through the shared backing + # allocation), mirroring the Distributed "ChunkView -> Aliasing" testset + Random.seed!(8) + A = rand(8, 8) + DA = DArray(A, Blocks(4, 4)) + c = fetch(DA.chunks[1, 1]; raw=true) + @test c.handle.rank == 0 # deterministic placement puts block (1,1) on rank 0 + cv_top = view(c, 1:2, 1:4) + cv_bot = view(c, 3:4, 1:4) + @test cv_top isa Dagger.ChunkView + + r1 = min(1, nranks-1) + r2 = min(2, nranks-1) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) add1!(InOut(cv_top)) + Dagger.@spawn scope=rank_scope(r2) scale2!(InOut(cv_bot)) + Dagger.@spawn scope=rank_scope(0) add1!(InOut(c)) + end + + ref_blk = A[1:4, 1:4] + ref_blk[1:2, :] .+= 1 + ref_blk[3:4, :] .*= 2 + ref_blk .+= 1 + # Collective uniform fetches: identical on every rank + @test fetch(c) ≈ ref_blk + @test fetch(cv_top) ≈ ref_blk[1:2, :] +end + +@testset "DTask arguments" begin + # DTasks as datadeps arguments resolve to their (rank-local) result chunks + Random.seed!(9) + t = Dagger.@spawn rand(4, 4) + # Collective fetch; copy since the owner's fetch aliases the stored result + B = copy(fetch(t)) + exec_rank = min(1, nranks-1) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(exec_rank) add1!(InOut(t)) + end + @test fetch(t) ≈ B .+ 1 +end + +@testset "Dep-mod coverage" begin + # Diagonal dep_mod with cross-rank placement + Random.seed!(10) + A = rand(6, 6) + ref = copy(A) + r1 = min(1, nranks-1) + r2 = min(2, nranks-1) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) diag_inc!(Deps(A, InOut(Diagonal))) + Dagger.@spawn scope=rank_scope(r2) add1!(InOut(A)) + Dagger.@spawn scope=rank_scope(0) diag_inc!(Deps(A, InOut(Diagonal))) + end + diag_inc!(ref); ref .+= 1; diag_inc!(ref) + if rank == 0 + @test A ≈ ref + end + + # Full triangular wrapper matrix: each dep_mod narrows the cross-rank + # move! to its region (mirrors test/datadeps.jl's (Unit)Upper/Lower set) + Random.seed!(11) + B = rand(6, 6) + refB = copy(B) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) upper_double!(Deps(B, InOut(UpperTriangular))) + Dagger.@spawn scope=rank_scope(r2) strict_lower_double!(Deps(B, InOut(UnitLowerTriangular))) + Dagger.@spawn scope=rank_scope(0) strict_upper_double!(Deps(B, InOut(UnitUpperTriangular))) + Dagger.@spawn scope=rank_scope(r1) lower_double!(Deps(B, InOut(LowerTriangular))) + Dagger.@spawn scope=rank_scope(r2) diag_inc!(Deps(B, InOut(Diagonal))) + end + upper_double!(refB); strict_lower_double!(refB); strict_upper_double!(refB) + lower_double!(refB); diag_inc!(refB) + if rank == 0 + @test B ≈ refB + end +end + +@testset "Field aliasing" begin + # Symbol dep_mods narrow the dependency (and cross-rank move!) to one + # field of a wrapped object (mirrors test/datadeps.jl "Field aliasing") + Random.seed!(12) + X = Ref(rand(100)) + r1 = min(1, nranks-1) + t = Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) mut_ref!(Deps(X, InOut(:x))) + Dagger.@spawn scope=rank_scope(0) getfield(Deps(X, In(:x)), :x) + end + # Collective fetch of the read task's result: program-order RAW means it + # observed the mutation + @test all(==(0), fetch(t)) + # Rank 0 owns the wrapped Ref and receives the write-back + if rank == 0 + @test all(==(0), X[]) + end +end + +@testset "Scope semantics" begin + # Outer scope: get_compute_scope() restricts where datadeps tasks run + target = min(1, nranks-1) + ts = Dagger.with_options(;scope=rank_scope(target)) do + Dagger.spawn_datadeps() do + [Dagger.@spawn exec_rank() for _ in 1:4] + end + end + @test all(==(target), fetch.(ts)) + + # Inner scope: an unsatisfiable task scope throws uniformly on all ranks + @test_throws Dagger.Sch.SchedulingException Dagger.spawn_datadeps() do + Dagger.@spawn scope=Dagger.ExactScope(Dagger.ThreadProc(1, 5000)) 1+1 + end + # The system remains usable afterwards + A = ones(4, 4) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(target) add1!(InOut(A)) + end + if rank == 0 + @test A == fill(2.0, 4, 4) + end +end + +# Reduced dataflow verification (mirrors test/datadeps.jl's logging-based +# test_dataflow suite): under SPMD every rank plans the same copy tasks, so +# each rank checks its own logs for the exact set of `move!` tasks inserted +function with_logs(f) + Dagger.enable_logging!(;tasknames=true, all_task_deps=false) + try + f() + return Dagger.fetch_logs!() + finally + Dagger.disable_logging!() + end +end +function count_move_tasks(logs) + n = 0 + for w in keys(logs), name in logs[w][:tasknames] + if name isa String && occursin("move!", name) + n += 1 + end + end + return n +end +read_only(X) = (sum(X); nothing) + +@testset "Copy planning" begin + r1 = min(1, nranks-1) + cross = r1 != 0 # single-rank runs degenerate to zero copies + + # Read-only cross-rank arg: one copy-in, no writeback + A = ones(4, 4) + logs = with_logs() do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) read_only(In(A)) + end + end + @test count_move_tasks(logs) == (cross ? 1 : 0) + @test Dagger.check_uniform(count_move_tasks(logs)) + + # Same-rank InOut: the origin slot aliases the original, zero copies + B = ones(4, 4) + logs = with_logs() do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(0) add1!(InOut(B)) + end + end + @test count_move_tasks(logs) == 0 + + # Cross-rank InOut: exactly one copy-in plus one writeback + C = ones(4, 4) + logs = with_logs() do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) add1!(InOut(C)) + end + end + @test count_move_tasks(logs) == (cross ? 2 : 0) + + # RAW chain on one rank: the intermediate copy is elided (slot up-to-date) + D = ones(4, 4) + logs = with_logs() do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(r1) add1!(InOut(D)) + Dagger.@spawn scope=rank_scope(r1) scale2!(InOut(D)) + end + end + @test count_move_tasks(logs) == (cross ? 2 : 0) + if rank == 0 + @test D == fill(4.0, 4, 4) + end +end + +@testset "Error propagation" begin + # A throwing task poisons its result-status broadcast, so every rank + # rethrows and SPMD control flow stays aligned + @test_throws_unwrap (Dagger.DTaskFailedException, ErrorException) Dagger.spawn_datadeps() do + Dagger.@spawn error("Test") + end + # Concrete inferred return type still needs a status broadcast when the + # compiler does not prove :nothrow (type skip alone cannot silence errors) + @test_throws Exception Dagger.spawn_datadeps() do + Dagger.@spawn concrete_fail(1) + end + t = Dagger.@spawn concrete_fail(1) + @test_throws Exception fetch(t) + # Concrete + nothrow: no status/type traffic; all ranks still agree + r = Dagger.spawn_datadeps() do + Dagger.@spawn nothrow_add(41) + end + @test fetch(r) == 42 + # The system remains usable afterwards + A = ones(4, 4) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=rank_scope(min(1, nranks-1)) add1!(InOut(A)) + end + if rank == 0 + @test A == fill(2.0, 4, 4) + end +end + +@testset "collect/fetch SPMD" begin + Random.seed!(3) + A = rand(12, 12) + DA = DArray(A, Blocks(3, 3)) + # Datadeps tree collect: identical on every rank + @test collect(DA) == A + + # fetch resolves chunk records locally (owners keep their payloads) + fa = fetch(DA) + for c in vec(fa.chunks) + @test c isa Dagger.Chunk + h = c.handle + @test (h.rank == rank) == (h.innerRef !== nothing) + end + + # Mutation round-trip through a datadeps region + Dagger.spawn_datadeps() do + for c in vec(DA.chunks) + Dagger.@spawn inc!(InOut(c)) + end + end + @test collect(DA) == A .+ 1 +end + +@testset "Untyped task results" begin + # Non-concrete inferred return types rely on the owner's result-status + # broadcast carrying the actual result type to all ranks + t1 = Dagger.@spawn untyped_result(3.0) + @test fetch(t1) === 3 + t2 = Dagger.@spawn untyped_result(-3.0) + @test fetch(t2) === -3.0 +end + +@testset "Nothing task results" begin + # Void results must stay typed as Nothing (not rewritten to Any) + t = Dagger.@spawn (() -> nothing)() + @test fetch(t) === nothing + @test Dagger.chunktype(fetch(t; raw=true)) === Nothing +end + +@testset "Production mode (uniformity checks off)" begin + # With checks disabled, ranks run planning and execution fully + # asynchronously: no compare_all broadcasts keep them in lockstep, so this + # exercises tag matching and task-overlap under maximal divergence + # (a serialized-execution regression here deadlocks in copy cycles) + Dagger.check_uniformity!(false) + try + Random.seed!(40) + N, B = 256, 64 + A = rand(N, N); A = A * A'; A[diagind(A)] .+= N + local U + for _ in 1:2 + DA = DArray(A, Blocks(B, B)) + U = collect(cholesky(DA).U) + end + @test U ≈ cholesky(A).U + finally + Dagger.check_uniformity!(true) + end +end + +end # @testset "Datadeps" + +@testset "Cholesky" begin + Random.seed!(42) + A = rand(Float64, 64, 64) + A = A * A' + A[diagind(A)] .+= size(A, 1) + B = copy(A) + + DA = zeros(Blocks(16, 16), Float64, 64, 64) + copyto!(DA, A) + LinearAlgebra._chol!(DA, UpperTriangular) + U_dist = UpperTriangular(collect(DA)) + + C = cholesky(B) + err = norm(U_dist - C.U) / norm(C.U) + recon = norm(U_dist' * U_dist - B) / norm(B) + @test err < 1e-12 + @test recon < 1e-12 +end + +@testset "LU" begin + Random.seed!(1234) + A = randn(100, 100) + orig_A = copy(A) + DA = DArray(A, Blocks(25, 25)) + + F = lu!(DA, RowMaximum(); check=false) + lu!(A, RowMaximum()) + + @test norm(collect(A) - collect(DA)) / norm(collect(A)) < 1e-12 + DAc = collect(DA) + p = LinearAlgebra.ipiv2perm(collect(F.ipiv), size(DAc, 1)) + LtU = UnitLowerTriangular(DAc) * UpperTriangular(DAc) + @test norm(LtU - orig_A[p, :]) / norm(orig_A) < 1e-12 +end + +end # @testset "MPI" + +MPI.Barrier(comm) +Core.println("[$rank] MPI suite OK") diff --git a/test/mpi_cholesky.jl b/test/mpi_cholesky.jl new file mode 100644 index 000000000..5559ec96a --- /dev/null +++ b/test/mpi_cholesky.jl @@ -0,0 +1,25 @@ +using Dagger, MPI, LinearAlgebra, Random + +Dagger.accelerate!(:mpi) +Dagger.check_uniformity!(true) +comm = MPI.COMM_WORLD +rank = MPI.Comm_rank(comm) + +Random.seed!(42) +A = rand(Float64, 64, 64) +A = A * A' +A[diagind(A)] .+= size(A, 1) +B = copy(A) + +DA = zeros(Blocks(16, 16), Float64, 64, 64) +copyto!(DA, A) +LinearAlgebra._chol!(DA, UpperTriangular) +U_dist = UpperTriangular(collect(DA)) + +C = cholesky(B) +err = norm(U_dist - C.U) / norm(C.U) +recon = norm(U_dist' * U_dist - B) / norm(B) +Core.println("[$rank] chol err=$err recon=$recon") +@assert err < 1e-12 && recon < 1e-12 +MPI.Barrier(comm) +Core.println("[$rank] chol OK") diff --git a/test/mpi_cuda.jl b/test/mpi_cuda.jl new file mode 100644 index 000000000..eb943c1b8 --- /dev/null +++ b/test/mpi_cuda.jl @@ -0,0 +1,22 @@ +# MPI × GPU datadeps suite for CUDA. See test/mpi_gpu_suite.jl for the shared +# logic; this file only supplies the CUDA-specific config. +# +# Run (env must provide Dagger, MPI, CUDA; see test/cudaenv): +# mpiexec -n 2 julia --project=test/cudaenv --threads=2 test/mpi_cuda.jl + +using Dagger, MPI, CUDA, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +include(joinpath(@__DIR__, "mpi_gpu_suite.jl")) + +const CUDAExt = Base.get_extension(Dagger, :CUDAExt) +@assert CUDAExt !== nothing "CUDAExt failed to load" + +run_mpi_gpu_suite((; + name = "CUDA", + DeviceProc = CUDAExt.CuArrayDeviceProc, + elt = Float64, + subarray_depmod = true, + matmul = true, + cholesky = true, +)) diff --git a/test/mpi_gpu_suite.jl b/test/mpi_gpu_suite.jl new file mode 100644 index 000000000..4102081a6 --- /dev/null +++ b/test/mpi_gpu_suite.jl @@ -0,0 +1,249 @@ +# Shared MPI × GPU datadeps test suite. +# +# Every backend (CUDA, ROCm, OpenCL, Metal, oneAPI) drives the exact same SPMD +# datadeps logic; only the array/device/memory-space names and the set of +# supported features differ. Rather than copy this suite per backend, each +# `test/mpi_.jl` loads its backend, then calls `run_mpi_gpu_suite` +# with a small config describing those differences. +# +# Data plane is host-staged through P2P (no GPU-aware MPI required). +# +# Config fields (NamedTuple; feature fields default to off / absent): +# name :: String backend label used in testset names / logs +# DeviceProc :: Type the backend's device Processor type +# elt :: Type element type for host arrays (Float32/Float64) +# subarray_depmod :: Bool run SubArray + dep_mod / dep-mod coverage +# matmul :: Bool run the DArray matmul smoke +# cholesky :: Bool run the DArray cholesky smoke +# remap :: NamedTuple or absent — MPI remap / memory-kind hooks: +# make_space :: Function () -> a VRAM memory space (device 1) +# device_field :: Symbol field naming the device index (:device/:device_id) +# kind :: Symbol expected gpu_memory_kind +# test_ipc :: Bool also assert !ipc_eligible(space, space) +# extra :: Function or absent — extra backend-specific checks + +using Dagger, MPI, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +const MPIExt = Base.get_extension(Dagger, :MPIExt) + +# Broadcast-only mutation helpers (scalar indexing is illegal on GPU arrays) +add1!(X) = (X .+= 1; nothing) +scale2!(X) = (X .*= 2; nothing) +function upper_double!(X) + for d in 0:size(X, 1)-1 + view(X, diagind(X, d)) .*= 2 + end + return nothing +end +function diag_inc!(X) + view(X, diagind(X)) .+= 1 + return nothing +end +function strict_upper_double!(X) + for d in 1:size(X, 1)-1 + view(X, diagind(X, d)) .*= 2 + end + return nothing +end +function lower_double!(X) + for d in 0:size(X, 1)-1 + view(X, diagind(X, -d)) .*= 2 + end + return nothing +end +function strict_lower_double!(X) + for d in 1:size(X, 1)-1 + view(X, diagind(X, -d)) .*= 2 + end + return nothing +end + +function run_mpi_gpu_suite(cfg) + DeviceProc = cfg.DeviceProc + elt = cfg.elt + + # Init MPI (via the MPI acceleration) before any MPI routine is called. + Dagger.accelerate!(:mpi) + Dagger.check_uniformity!(true) + + comm = MPI.COMM_WORLD + rank = MPI.Comm_rank(comm) + nranks = MPI.Comm_size(comm) + + mpi_procs() = sort(collect(Dagger.get_processors(MPIExt.MPIClusterProc(comm))); + by=p->(p.rank, Dagger.short_name(p))) + gpu_proc_for_rank(r) = first(filter(p->p.rank == r && p.innerProc isa DeviceProc, mpi_procs())) + cpu_proc_for_rank(r) = first(filter(p->p.rank == r && p.innerProc isa Dagger.ThreadProc, mpi_procs())) + gpu_scope(r) = Dagger.ExactScope(gpu_proc_for_rank(r)) + cpu_scope(r) = Dagger.ExactScope(cpu_proc_for_rank(r)) + all_gpu_scope() = Dagger.UnionScope([gpu_scope(r) for r in 0:nranks-1]...) + # GPU processors are not default-enabled; regions mixing CPU and GPU tasks + # must opt in via an outer scope that includes both kinds + mixed_scope() = Dagger.UnionScope(vcat([cpu_scope(r) for r in 0:nranks-1], + [gpu_scope(r) for r in 0:nranks-1])...) + +@testset "MPI $(cfg.name) Datadeps" begin + +@testset "GPU processors visible per rank" begin + for r in 0:nranks-1 + p = gpu_proc_for_rank(r) + @test p.innerProc isa DeviceProc + @test Dagger.check_uniform(p) + end +end + +@testset "Cross-rank GPU dataflow" begin + # CPU-origin data mutated by GPU tasks on different ranks, in program + # order, with write-back to the rank-0 CPU original + Random.seed!(20) + A = rand(elt, 8, 8) + ref = copy(A) + r1 = min(1, nranks-1) + Dagger.with_options(;scope=mixed_scope()) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=gpu_scope(r1) add1!(InOut(A)) + Dagger.@spawn scope=gpu_scope(0) scale2!(InOut(A)) + end + end + ref .+= 1 + ref .*= 2 + if rank == 0 + @test A ≈ ref + end +end + +@testset "Mixed CPU/GPU dataflow" begin + # The same slot chain crosses CPU and GPU spaces (HtoD and DtoH inside + # one region), across ranks + Random.seed!(21) + A = rand(elt, 8, 8) + ref = copy(A) + r1 = min(1, nranks-1) + Dagger.with_options(;scope=mixed_scope()) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=gpu_scope(r1) add1!(InOut(A)) + Dagger.@spawn scope=cpu_scope(0) scale2!(InOut(A)) + Dagger.@spawn scope=gpu_scope(0) add1!(InOut(A)) + end + end + ref .+= 1; ref .*= 2; ref .+= 1 + if rank == 0 + @test A ≈ ref + end +end + +if get(cfg, :subarray_depmod, false) + @testset "GPU SubArray + dep_mod dataflow" begin + # View remainders and non-identity dep_mods exercise multi-span KA pack + # / copy paths across ranks (mirrors test/mpi.jl Cross-rank dataflow) + Random.seed!(24) + A = rand(elt, 8, 8) + ref = copy(A) + vA = view(A, 3:6, 3:6) + r1 = min(1, nranks-1) + r2 = min(2, nranks-1) % nranks + Dagger.with_options(;scope=mixed_scope()) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=gpu_scope(r1) add1!(InOut(A)) + Dagger.@spawn scope=gpu_scope(r2) scale2!(InOut(vA)) + Dagger.@spawn scope=gpu_scope(0) upper_double!(Deps(A, InOut(UpperTriangular))) + end + end + ref .+= 1 + view(ref, 3:6, 3:6) .*= 2 + upper_double!(ref) + if rank == 0 + @test A ≈ ref + end + end + + @testset "GPU dep-mod coverage" begin + Random.seed!(25) + A = rand(elt, 6, 6) + ref = copy(A) + r1 = min(1, nranks-1) + r2 = min(2, nranks-1) % nranks + Dagger.with_options(;scope=mixed_scope()) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=gpu_scope(r1) diag_inc!(Deps(A, InOut(Diagonal))) + Dagger.@spawn scope=gpu_scope(r2) add1!(InOut(A)) + Dagger.@spawn scope=gpu_scope(0) diag_inc!(Deps(A, InOut(Diagonal))) + end + end + diag_inc!(ref); ref .+= 1; diag_inc!(ref) + if rank == 0 + @test A ≈ ref + end + + Random.seed!(26) + B = rand(elt, 6, 6) + refB = copy(B) + Dagger.with_options(;scope=mixed_scope()) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=gpu_scope(r1) upper_double!(Deps(B, InOut(UpperTriangular))) + Dagger.@spawn scope=gpu_scope(r2) strict_lower_double!(Deps(B, InOut(UnitLowerTriangular))) + Dagger.@spawn scope=gpu_scope(0) strict_upper_double!(Deps(B, InOut(UnitUpperTriangular))) + Dagger.@spawn scope=gpu_scope(r1) lower_double!(Deps(B, InOut(LowerTriangular))) + Dagger.@spawn scope=gpu_scope(r2) diag_inc!(Deps(B, InOut(Diagonal))) + end + end + upper_double!(refB); strict_lower_double!(refB); strict_upper_double!(refB) + lower_double!(refB); diag_inc!(refB) + if rank == 0 + @test B ≈ refB + end + end +end + +if get(cfg, :matmul, false) + @testset "GPU matmul (DArray)" begin + Random.seed!(22) + A = rand(elt, 8, 8) + B = rand(elt, 8, 8) + DA = DArray(A, Blocks(4, 4)) + DB = DArray(B, Blocks(4, 4)) + DC = zeros(Blocks(4, 4), elt, 8, 8) + Dagger.with_options(;scope=all_gpu_scope()) do + mul!(DC, DA, DB) + end + C = collect(DC) + @test C ≈ A * B + @test Dagger.check_uniform(hash(C)) + end +end + +if get(cfg, :cholesky, false) + @testset "GPU cholesky (DArray)" begin + Random.seed!(23) + A = rand(elt, 8, 8) + A = A * A' + A[diagind(A)] .+= size(A, 1) + DA = DArray(A, Blocks(4, 4)) + Dagger.with_options(;scope=all_gpu_scope()) do + @test collect(cholesky(DA).U) ≈ cholesky(A).U + end + end +end + +remap = get(cfg, :remap, nothing) +if remap !== nothing + @testset "MPI remap / memory kind hooks" begin + space = remap.make_space() + remapped = Dagger.mpi_remap_space(space, 3) + @test remapped.owner == 3 + @test getfield(remapped, remap.device_field) == 1 + @test Dagger.gpu_memory_kind(space) === remap.kind + if get(remap, :test_ipc, false) + @test !Dagger.ipc_eligible(space, space) + end + extra = get(remap, :extra, nothing) + extra !== nothing && extra() + end +end + +end # @testset "MPI $(cfg.name) Datadeps" + + MPI.Barrier(comm) + Core.println("[$rank] MPI $(cfg.name) suite OK") +end diff --git a/test/mpi_metal.jl b/test/mpi_metal.jl new file mode 100644 index 000000000..23a068399 --- /dev/null +++ b/test/mpi_metal.jl @@ -0,0 +1,25 @@ +# MPI × GPU datadeps suite for Metal. See test/mpi_gpu_suite.jl for the shared +# logic; this file only supplies the Metal-specific config. +# +# Run (macOS; env must provide Dagger, MPI, Metal; see test/metalenv): +# mpiexec -n 2 julia --project=test/metalenv --threads=2 test/mpi_metal.jl + +using Dagger, MPI, Metal, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +include(joinpath(@__DIR__, "mpi_gpu_suite.jl")) + +const MetalExt = Base.get_extension(Dagger, :MetalExt) +@assert MetalExt !== nothing "MetalExt failed to load" + +run_mpi_gpu_suite((; + name = "Metal", + DeviceProc = MetalExt.MtlArrayDeviceProc, + elt = Float32, + remap = (; + make_space = () -> MetalExt.MetalVRAMMemorySpace(1, 1), + device_field = :device_id, + kind = :Metal, + test_ipc = true, + ), +)) diff --git a/test/mpi_oneapi.jl b/test/mpi_oneapi.jl new file mode 100644 index 000000000..2202097a7 --- /dev/null +++ b/test/mpi_oneapi.jl @@ -0,0 +1,28 @@ +# MPI × GPU datadeps suite for oneAPI. See test/mpi_gpu_suite.jl for the shared +# logic; this file only supplies the oneAPI-specific config. +# +# GPU-direct requires Intel MPI + I_MPI_OFFLOAD or DAGGER_MPI_GPU_DIRECT=1; +# otherwise the data plane is host-staged through P2P. +# +# Run (env must provide Dagger, MPI, oneAPI; see test/oneapienv): +# mpiexec -n 2 julia --project=test/oneapienv --threads=2 test/mpi_oneapi.jl + +using Dagger, MPI, oneAPI, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +include(joinpath(@__DIR__, "mpi_gpu_suite.jl")) + +const IntelExt = Base.get_extension(Dagger, :IntelExt) +@assert IntelExt !== nothing "IntelExt failed to load" + +run_mpi_gpu_suite((; + name = "oneAPI", + DeviceProc = IntelExt.oneArrayDeviceProc, + elt = Float32, + matmul = true, + remap = (; + make_space = () -> IntelExt.IntelVRAMMemorySpace(1, 1), + device_field = :device_id, + kind = :oneAPI, + ), +)) diff --git a/test/mpi_opencl.jl b/test/mpi_opencl.jl new file mode 100644 index 000000000..ab2816dcb --- /dev/null +++ b/test/mpi_opencl.jl @@ -0,0 +1,25 @@ +# MPI × GPU datadeps suite for OpenCL. See test/mpi_gpu_suite.jl for the shared +# logic; this file only supplies the OpenCL-specific config. +# +# Run (env must provide Dagger, MPI, OpenCL; see test/openclenv): +# mpiexec -n 2 julia --project=test/openclenv --threads=2 test/mpi_opencl.jl + +using Dagger, MPI, OpenCL, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +include(joinpath(@__DIR__, "mpi_gpu_suite.jl")) + +const OpenCLExt = Base.get_extension(Dagger, :OpenCLExt) +@assert OpenCLExt !== nothing "OpenCLExt failed to load" + +run_mpi_gpu_suite((; + name = "OpenCL", + DeviceProc = OpenCLExt.CLArrayDeviceProc, + elt = Float32, + remap = (; + make_space = () -> OpenCLExt.CLMemorySpace(1, 1), + device_field = :device, + kind = :OpenCL, + test_ipc = true, + ), +)) diff --git a/test/mpi_rocm.jl b/test/mpi_rocm.jl new file mode 100644 index 000000000..76d8ce9ae --- /dev/null +++ b/test/mpi_rocm.jl @@ -0,0 +1,45 @@ +# MPI × GPU datadeps suite for ROCm. See test/mpi_gpu_suite.jl for the shared +# logic; this file only supplies the ROCm-specific config. +# +# Run (env must provide Dagger, MPI, AMDGPU; see test/rocmenv): +# mpiexec -n 2 julia --project=test/rocmenv --threads=2 test/mpi_rocm.jl + +using Dagger, MPI, AMDGPU, LinearAlgebra, Random, Test +using Dagger: In, Out, InOut, Deps + +include(joinpath(@__DIR__, "mpi_gpu_suite.jl")) + +const ROCExt = Base.get_extension(Dagger, :ROCExt) +@assert ROCExt !== nothing "ROCExt failed to load" + +function rocm_gpu_direct_check() + # Default is library-detected; force-off must stick for this process + old = get(ENV, "DAGGER_MPI_GPU_DIRECT", nothing) + try + ENV["DAGGER_MPI_GPU_DIRECT"] = "0" + ROCExt.MPI_GPU_DIRECT[] = nothing + A = AMDGPU.ones(Float32, 4) + @test Dagger.mpi_device_direct(A) == false + finally + if old === nothing + delete!(ENV, "DAGGER_MPI_GPU_DIRECT") + else + ENV["DAGGER_MPI_GPU_DIRECT"] = old + end + ROCExt.MPI_GPU_DIRECT[] = nothing + end +end + +run_mpi_gpu_suite((; + name = "ROCm", + DeviceProc = ROCExt.ROCArrayDeviceProc, + elt = Float64, + matmul = true, + cholesky = true, + remap = (; + make_space = () -> ROCExt.ROCVRAMMemorySpace(1, 1), + device_field = :device_id, + kind = :ROC, + extra = rocm_gpu_direct_check, + ), +)) diff --git a/test/mpienv/Project.toml b/test/mpienv/Project.toml new file mode 100644 index 000000000..71e617954 --- /dev/null +++ b/test/mpienv/Project.toml @@ -0,0 +1,9 @@ +[deps] +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Dagger = {path = "../.."} diff --git a/test/oneapienv/Project.toml b/test/oneapienv/Project.toml new file mode 100644 index 000000000..30e7631f3 --- /dev/null +++ b/test/oneapienv/Project.toml @@ -0,0 +1,10 @@ +[deps] +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" + +[sources] +Dagger = {path = "../.."} diff --git a/test/openclenv/Project.toml b/test/openclenv/Project.toml new file mode 100644 index 000000000..b7718eca7 --- /dev/null +++ b/test/openclenv/Project.toml @@ -0,0 +1,10 @@ +[deps] +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Dagger = {path = "../.."} diff --git a/test/procgrid.jl b/test/procgrid.jl new file mode 100644 index 000000000..1321ca44e --- /dev/null +++ b/test/procgrid.jl @@ -0,0 +1,103 @@ +using Test +using Dagger +using Dagger: CyclicProcGrid, BlockedProcGrid, DenseProcGrid, + build_procgrid, procgrid_processor, materialize_procgrid, + DistributedAcceleration + +function legacy_dense_procgrid(assignment, sizeA, blocksize, availprocs) + N = length(sizeA) + np = length(availprocs) + nblocks = ntuple(i -> cld(sizeA[i], blocksize[i]), N) + if assignment == :arbitrary + return nothing + elseif assignment == :blockrow + shape = ntuple(i -> i == 1 ? Int(ceil(sizeA[1] / blocksize[1])) : 1, N) + rows_per_proc, extra = divrem(Int(ceil(sizeA[1] / blocksize[1])), np) + counts = [rows_per_proc + (i <= extra ? 1 : 0) for i in 1:np] + return reshape(vcat(fill.(availprocs, counts)...), shape) + elseif assignment == :blockcol + shape = ntuple(i -> i == N ? Int(ceil(sizeA[N] / blocksize[N])) : 1, N) + cols_per_proc, extra = divrem(Int(ceil(sizeA[N] / blocksize[N])), np) + counts = [cols_per_proc + (i <= extra ? 1 : 0) for i in 1:np] + return reshape(vcat(fill.(availprocs, counts)...), shape) + elseif assignment == :cyclicrow + shape = ntuple(i -> i == 1 ? np : 1, N) + return reshape(availprocs, shape) + elseif assignment == :cycliccol + shape = ntuple(i -> i == N ? np : 1, N) + return reshape(availprocs, shape) + end + error("unknown assignment") +end + +function legacy_lookup(grid, I::CartesianIndex) + grid[CartesianIndex(mod1.(Tuple(I), size(grid))...)] +end + +availprocs = collect(Dagger.compatible_processors()) +filter!(p -> p isa Dagger.ThreadProc, availprocs) +sort!(availprocs, by=x -> (x.owner, x.tid)) +np = length(availprocs) +accel = DistributedAcceleration() + +@testset "build_procgrid :arbitrary Distributed" begin + @test build_procgrid(:arbitrary, (16, 16), (4, 4), accel) === nothing +end + +for assignment in (:blockrow, :blockcol, :cyclicrow, :cycliccol) + @testset "parity $assignment" begin + sizeA = (41, 35) + blocksize = (4, 3) + nblocks = ntuple(i -> cld(sizeA[i], blocksize[i]), 2) + pg = build_procgrid(assignment, sizeA, blocksize, accel) + legacy = legacy_dense_procgrid(assignment, sizeA, blocksize, availprocs) + dense = materialize_procgrid(pg, nblocks) + for I in CartesianIndices(nblocks) + @test procgrid_processor(pg, I) == legacy_lookup(legacy, I) + @test dense[I] == legacy_lookup(legacy, I) + end + # tiling: larger block grid than procgrid shape + big = (nblocks[1] + 2, nblocks[2] + 1) + for I in CartesianIndices(big) + @test procgrid_processor(pg, I) == legacy_lookup(legacy, I) + end + end +end + +@testset "CyclicProcGrid MPI round-robin parity" begin + nblocks = (4, 4) + procs = availprocs[1:min(4, np)] + pg = CyclicProcGrid(procs, nblocks) + legacy = Array{Dagger.Processor, 2}(undef, nblocks) + for (i, I) in enumerate(CartesianIndices(legacy)) + legacy[I] = procs[mod1(i, length(procs))] + end + for I in CartesianIndices(nblocks) + @test procgrid_processor(pg, I) == legacy[I] + end +end + +@testset "BlockedProcGrid" begin + np_test = min(np, 3) + if np_test >= 2 + shape = (6, 1) + counts = fill(2, np_test) + procs = availprocs[1:np_test] + pg = BlockedProcGrid(procs, counts, shape, 1) + legacy = reshape(vcat(fill.(procs, counts)...), shape) + for I in CartesianIndices((8, 1)) + @test procgrid_processor(pg, I) == legacy_lookup(legacy, I) + end + end +end + +@testset "DenseProcGrid" begin + np_test = min(np, 4) + if np_test >= 4 + grid = reshape(availprocs[1:4], (2, 2)) + pg = DenseProcGrid(grid) + for I in CartesianIndices((5, 5)) + @test procgrid_processor(pg, I) == legacy_lookup(grid, I) + end + end +end diff --git a/test/rocmenv/Project.toml b/test/rocmenv/Project.toml new file mode 100644 index 000000000..d350b4c4a --- /dev/null +++ b/test/rocmenv/Project.toml @@ -0,0 +1,10 @@ +[deps] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Dagger = {path = "../.."} diff --git a/test/run_mpi.jl b/test/run_mpi.jl new file mode 100644 index 000000000..4a5471102 --- /dev/null +++ b/test/run_mpi.jl @@ -0,0 +1,25 @@ +# Launch an MPI test script through the MPI.jl-provided `mpiexec`, so no +# system MPI installation is required (MPI.jl uses its configured binary, +# MPICH_jll by default). The child processes inherit the active project. +# +# Usage: +# julia --project= test/run_mpi.jl