diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 26e3a7a2b..d83131171 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -41,6 +41,7 @@ steps: queue: "rocm" env: CI_USE_ROCM: "1" + JULIA_NUM_THREADS: "2" - label: Julia 1.11 (oneAPI) key: test-oneapi @@ -55,6 +56,8 @@ steps: queue: "oneapi" env: CI_USE_ONEAPI: "1" + # oneMKL native-queue interop aborts with the v2 adapter on Xe2. + SYCL_UR_USE_LEVEL_ZERO_V2: "0" - label: Julia 1.11 (Metal) key: test-metal @@ -95,21 +98,19 @@ steps: julia --project=test/rocmenv -e 'using Pkg; Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/TimespanLogging")]); Pkg.instantiate()' julia --project=test/rocmenv test/mpi/run.jl 2 2 test/mpi/rocm.jl - # Disabled: has failed on every build since at least 2026-09-16, including on master (e.g. - # ZE_RESULT_ERROR_INVALID_ARGUMENT and "Non-uniform value" errors in the dense matvec tests). - # Re-enable once test/mpi/oneapi.jl passes. - # - # - label: Julia 1.11 (oneAPI, MPI) - # key: test-oneapi-mpi - # <<: *mpigputest - # plugins: - # - JuliaCI/julia#v1: - # version: "1.11" - # agents: - # queue: "oneapi" - # command: | - # julia --project=test/oneapienv -e 'using Pkg; Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/TimespanLogging")]); Pkg.instantiate()' - # julia --project=test/oneapienv test/mpi/run.jl 2 2 test/mpi/oneapi.jl + - label: Julia 1.11 (oneAPI, MPI) + key: test-oneapi-mpi + <<: *mpigputest + plugins: + - JuliaCI/julia#v1: + version: "1.11" + agents: + queue: "oneapi" + env: + SYCL_UR_USE_LEVEL_ZERO_V2: "0" + command: | + julia --project=test/oneapienv -e 'using Pkg; Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/TimespanLogging")]); Pkg.instantiate()' + julia --project=test/oneapienv test/mpi/run.jl 2 2 test/mpi/oneapi.jl # Disabled: has failed on every build since at least 2026-09-16, including on master (the # process aborts with Metal's "commit command buffer with uncommitted encoder" assertion). @@ -337,6 +338,7 @@ steps: command: "bash .buildkite/run_gpu_benchmarks.sh {{matrix}}" env: BENCHMARK_GPU: "oneapi" + SYCL_UR_USE_LEVEL_ZERO_V2: "0" BENCHMARK_VARIANT: "gpu" BENCHMARK_TEST_STEPS: "test-oneapi" BENCHMARK_CI_THREADS: "2" @@ -367,6 +369,7 @@ steps: command: "bash .buildkite/run_gpu_benchmarks.sh {{matrix}}" env: BENCHMARK_GPU: "oneapi" + SYCL_UR_USE_LEVEL_ZERO_V2: "0" BENCHMARK_VARIANT: "gpu-distributed" BENCHMARK_TEST_STEPS: "test-oneapi" BENCHMARK_CI_THREADS: "2" @@ -381,7 +384,9 @@ steps: artifacts: - "benchmark-results-*/**/*" - # Disabled along with the "Julia 1.11 (oneAPI, MPI)" step it depends on. + # Disabled: the two-rank linalg benchmark segfaults in oneMKL/SYCL at + # N=1024 with 128-element tiles, including with the legacy Level Zero adapter. + # Functional MPI tests pass; re-enable after the benchmark matrix also passes. # # - label: "Benchmarks (oneAPI, gpu-mpi, {{matrix}})" # key: "bench-oneapi-gpu-mpi" diff --git a/AGENTS.md b/AGENTS.md index 0385d5384..f501f5063 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -686,3 +686,176 @@ lesson. running; MPICH had placed each rank in its own process group. After a timeout, inspect rank PIDs/PGIDs and stop those groups too. Killing only the wrapper's group is not sufficient for this launcher. + +67. **A Distributed GPU transfer must preserve the source's ownership.** + Sending a raw device array in a slot-creation RPC deserializes it on the + receiver while `from_proc` still names the sender. The subsequent GPU + `move` then activates a remote worker's context locally and asserts; some + backends cannot serialize the array in the first place. Send an owner-side + `Chunk` and use the backend's Chunk transport instead. The same rule + applies when fetching a remote source for in-place copies. Pin regression + inputs and consumers to different workers selecting device 1: a GPU scope + alone does not guarantee that a transfer actually happens. + +68. **CUDA IPC must export staging memory and return an owned copy.** + CUDA's pooled allocations cannot be exported with `cuIpcGetMemHandle`, and + returning an imported mapping makes the receiver alias the sender's data + and depend on its lifetime. Use the shared `ipc_export`/`ipc_materialize` + hooks, keep the staging token on the sender until the receiver finishes, + and release it in `finally`. Use `CUDADRV` for closing handles too: CUDA 6 + moved the driver API into CUDACore. Test repeated transfers in both + directions and mutate the destination to verify source independence. + +69. **Two GPUs of one process are not one device, and may not even be + peers.** Several paths assumed "same worker" meant "one kernel may touch + both buffers": the same-worker device remainder copy, `collect`'s + in-process `cat`, and CUDA `pointer()` (which takes ownership for the + *active* device and throws without P2P). `CUDA.pin` is also only deduped + per context while host registration is process-wide. Key direct device + paths on equal memory spaces, not equal workers, read addresses without + ownership side effects, and test with tiles on two devices of one + process: the one-GPU-per-worker suites cannot see any of this. + +70. **Reading a GPU task's result needs Dagger's stream, not the caller's.** + GPU `execute!` returns without synchronizing: the kernels are still queued + on Dagger's per-device stream. A plain `fetch` + `Array(x)` from the + caller copies on the caller's own task-local stream, and is correct only + if the array library synchronizes the previous owner when another stream + touches a buffer. CUDA, AMDGPU and OpenCL do; oneAPI does not, and it also + copies on the *calling task's* device. A green CUDA run therefore does not + prove the wait exists. Before reading a tile on the host, call + `gpu_synchronize(chunk.processor)` and copy under `with_context`, on the + chunk's owner (as `collect` does via `_collect_host_tile`). Note that + IntelExt's hooks sync only the calling task's stream, and each oneAPI + `execute!` runs on a fresh task, so oneAPI's own hook is not sufficient yet. + +71. **Never unregister host memory from a finalizer that can block.** + `CUDA.pin` unregisters in a GC finalizer that takes a `ReentrantLock` + (as does `AMDGPU.Mem.unpin`, which ROCExt's finalizer used to call). When + that lock is contended the finalizer throws "task switch not allowed from + inside gc finalizer", the unregistration is lost, and GC frees the memory + anyway: the range stays registered with the driver, and a later + allocation at that address fails to register or is DMA'd through stale + pages. Multi-GPU GEMM segfaulted a few iterations in, with a backtrace + that blew the stack while printing (`jl_static_show` recursion) and showed + nothing of the cause; the "error in running finalizer" lines before it + were the only clue. (On ROCm, 2-GPU Datadeps tests segfaulted the same + way, inside `hipMemcpyWithStream`.) CUDAExt's and ROCExt's `pin_buffer!` + now register buffers themselves, and their finalizers only `trylock`, + re-arming themselves on contention (which keeps the buffer alive until a + later GC can unregister it). Anything that must run before + memory is freed has to be finalizer-safe: no blocking locks, no yields. + +72. **Two GPUs without peer access: never let CUDA.jl stage the copy.** + Without P2P, `copyto!` between devices goes through a fresh *pageable* + `Vector` with a synchronous DtoH: 1.4 GB/s on PCIe-attached L40S, where + `cuMemcpyPeerAsync` (which works without peer access; the driver pipelines + it through its own pinned buffers) reached 22 GB/s. A 4-GPU GEMM spent + 27 s per call in copies. Check `can_access_peer` before assuming a copy + path is cheap, and measure the copy primitive in isolation first -- it + took one 30-line script to find a 15x. (HIP does not stage this way: + between two RX 6900 XTs without peer access, AMDGPU's `copyto!` and + `hipMemcpyPeerAsync` both reached ~5.5 GB/s, vs 1.4 GB/s through a + pageable host buffer, and `copyto!` blocks until the copy is done.) + +73. **One stream per device serializes every transfer with that device's + compute.** Copies enqueued on the device stream wait behind its kernels, + and an event recorded on the *source's* stream waits for kernels that + merely read the data being sent (a GEMM's panel owner is busy reading the + panel it is asked to send). Peer copies now run on a per-device copy stream + ordered only against the allocations they touch (`BUFFER_EVENTS`, fed by + Datadeps' `writes` option), and are host-synced so every existing + synchronization point stays valid; 4-GPU GEMM went from ~3.9 s to ~3.1 s. + (ROCExt does the same with HIP events and `hipMemcpyPeerAsync`.) + Bookkeeping keyed by address outlives the allocation: a peer copy's + "already complete, skip the next write" flag, left on short-lived packing + buffers, swallowed the first write to the next allocation at that address, + and a later copy out of it read stale data. Name the task the skip is for + (`filled_by`), or keep entries to plain events. + Two traps found on the way: Datadeps spawns copy tasks `meta`, so + `execute!` sees `Chunk`s, not arrays (unwrap before classifying, or every + copy looks like an opaque argument and falls back to whole-stream waits); + and `unsafe_free!` tasks hand you a freed array (check `data.freed`). + +74. **Level Zero IPC handles cannot be serialized between processes.** Intel's + opaque `ze_ipc_mem_handle_t` includes a process-local file descriptor. + MPI's serialized metadata copies the descriptor number, not its underlying + resource, so `zeMemOpenIpcHandle` fails with `ZE_RESULT_ERROR_INVALID_ARGUMENT` + on the receiving rank. A usable import needs descriptor passing, such as + Unix-domain `SCM_RIGHTS`, before reconstructing the handle. Keep Intel MPI + on host-staged transport until that exists; a same-node/device check or + `DAGGER_IPC=1` cannot supply the missing descriptor. See + https://github.com/intel/compute-runtime/issues/448. + +75. **ROCm Distributed tests need a spare Julia thread for host calls.** + AMDGPU's persistent host-call tasks poll HSA signals with + `Libc.systemsleep` and `yield`, which can monopolize a sticky task's only + thread (lesson 12). With one default thread the GPU suite stopped making + progress in `hostcall_host_wait`, without an exception; with two threads + the host-call task remained alive while the other thread continued into + the stencil suite. Give both the driver and Distributed workers two + default threads in ROCm CI. Increasing only the driver's thread count + leaves remote workers exposed to the same starvation. + +76. **GPU event tracking must check a chunk's acceleration ownership.** + `root_worker_id(::MPIRef)` and MPI processors intentionally return `myid()` + on every rank; Distributed's worker-id test therefore calls every MPI chunk + local. A GPU copy task executes on both endpoint ranks, and recording its + arguments then tries to unwrap the remote endpoint's placeholder and throws + "MPIRef is not owned by this rank". Check the handle with + `is_local(current_acceleration(), chunk.handle)` before inspecting storage. + The existing cross-rank GPU dataflow tests exercise both endpoints; a + Distributed-only suite cannot catch this mistake. + +77. **Reduce apparent GPU ordering failures below Julia before adding fences.** + This host's Ryzen 9600X `gfx1036` GPU with HIP 7.2 returns zeros after an + asynchronous kernel or device copy followed by a short pause, even when + the same stream is synchronized before reading. A native HIP C++ program + reproduces it without Julia, KernelAbstractions, or Dagger; synchronous + allocations and disabling SDMA do not fix it. Setting + `AMD_SERIALIZE_KERNEL=3 AMD_SERIALIZE_COPY=3` makes that reproducer pass. + Per-halo synchronization happened to hide the failure, but would impose + a production cost for a host runtime problem. Use the serialization flags + only for local correctness validation and report them with the results; + they are not a valid configuration for performance measurements. + + +78. **oneMKL queue interop can fail below Dagger on Intel Xe2 GPUs.** + With oneAPI.jl 2.9.2 and its 2025.3 support libraries, this Arc Pro B60's + default Level Zero v2 adapter aborts with + `UR_RESULT_ERROR_UNSUPPORTED_FEATURE` when SYCL wraps a native Level Zero + queue for GEMM. A standalone `oneArray` `mul!` reproduces the abort; worker + EOF errors are its aftermath. `SYCL_UR_USE_LEVEL_ZERO_V2=0` selects the + legacy adapter and passes that reproducer. Diagnose native BLAS failures + separately from Dagger scheduling and record the adapter used in validation. + + +79. **A passing GPU suite does not certify dependency finalizers.** + Julia logs finalizer exceptions without failing the suite or its exit code. + AMDGPU 2.8.0's `src/cache.jl` `library_state` accepts a `destroy_handle` + callback but its task finalizer calls undefined `destroy_handle!`, producing + an `UndefVarError` while the two-rank numerical suite passes all assertions. + Check stderr as well as test counts, and attribute the callback to its + owning package before changing Dagger's task or GPU resource lifecycle. + + +80. **Installing an external GPU compiler is not the same as selecting it.** + GPUCompiler 1.23 defaults `GCNCompilerTarget.backend` to `:external` only + when `AMDGPU_LLVM_Backend_jll` is already loaded, not merely installed. + AMDGPU 2.1 on Julia 1.11 otherwise uses the embedded LLVM backend, which + computes wrong mixed-boundary indices on this `gfx1036` host despite + correct scalar arguments and valid-looking LLVM IR. Loading the external + LLVM 22 backend makes the unmodified kernel compute the correct indices. + Load it on every ROCm test/benchmark worker before kernel compilation; + adding it only to the driver's environment leaves remote codegen unchanged. + Keep both comparison revisions on the same backend. + +81. **Small Intel MPI smokes do not validate the benchmark workload.** On this + Arc Pro B60 with oneAPI 2.9.2, two-rank functional tests pass all 171/174 + assertions with host staging and the legacy Level Zero adapter, and the + MPI array benchmark passes all 28 leaves at sizes 256/1024. The linalg + worker nevertheless segfaults on `A*A` at size 1024 with 128-element tiles, + inside SYCL's kernel-info lookup called by oneMKL GEMM. That stack identifies + the failing native call, not its root cause. Keep Intel MPI functional tests + enabled separately from the benchmark jobs; a green small matmul test is + insufficient evidence to enable the larger tiled benchmark matrix. diff --git a/benchmark/backend.jl b/benchmark/backend.jl index 9c5a041b1..dd28af14b 100644 --- a/benchmark/backend.jl +++ b/benchmark/backend.jl @@ -19,6 +19,12 @@ if benchmark_accel !== nothing if benchmark_accel == "opencl" && get(ENV, "BENCHMARK_OPENCL_SOFTWARE", "false") == "true" @everywhere using pocl_jll end + if benchmark_accel == "amdgpu" + # CI installs this explicitly; recent AMDGPU also loads it itself. + @everywhere if Base.find_package("AMDGPU_LLVM_Backend_jll") !== nothing + using AMDGPU_LLVM_Backend_jll + end + end @eval @everywhere using $(cfg.package) end const benchmark_compute_scope = benchmark_accel === nothing ? nothing : diff --git a/benchmark/ci.jl b/benchmark/ci.jl index 79f7777ee..a5f9a0c77 100644 --- a/benchmark/ci.jl +++ b/benchmark/ci.jl @@ -69,6 +69,7 @@ const gpu_packages = Dict("cuda" => "CUDA", "amdgpu" => "AMDGPU", for spec in split(get(ENV, "BENCHMARK", ""), ';'), method in split(last(split(spec, ':')), ',') for accel in split(method, '+')[2:end] push!(EXTRA_PKGS, gpu_packages[accel]) + accel == "amdgpu" && push!(EXTRA_PKGS, "AMDGPU_LLVM_Backend_jll") end end if get(ENV, "BENCHMARK_OPENCL_SOFTWARE", "false") == "true" diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 7b773686f..0fbfd76b7 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -16,7 +16,7 @@ if isdefined(Base, :get_extension) else import ..CUDA end -import CUDA: CuDevice, CuContext, CuStream, CuArray, CUDABackend +import CUDA: CuDevice, CuContext, CuStream, CuEvent, CuArray, CUDABackend import CUDA: devices, attribute, context, context!, stream, stream! import CUDA: CUBLAS, CUSOLVER @@ -37,8 +37,12 @@ end function Dagger.aliasing(x::CuArray{T}) where T space = Dagger.memory_space(x) S = typeof(space) - cuptr = pointer(x) - rptr = Dagger.RemotePtr{Cvoid}(UInt64(cuptr), space) + # N.B. Not `pointer(x)`: that takes stream ownership of the buffer for the + # calling task's active device, which throws when that device cannot reach + # `x` (e.g. aliasing RPCs on a multi-GPU worker) and re-stamps ownership as + # a side effect. Aliasing only needs the address. + # (`x.offset` counts elements; pointer arithmetic counts bytes) + rptr = Dagger.RemotePtr{Cvoid}(UInt64(_raw_cuaddr(x)), space) return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end @@ -142,6 +146,17 @@ struct AllocateUndef{S} end (::AllocateUndef{S})(T, dims::Dims{N}) where {S,N} = CuArray{S,N}(undef, dims) Dagger.allocate_array_func(::CuArrayDeviceProc, ::Dagger.AllocateUndef{S}) where S = AllocateUndef{S}() +# Uninitialized Datadeps slots: without these, a slot in another GPU's memory +# was created by copying the whole array there (staged through the host when +# the devices have no peer access), only for Datadeps' copy-to phase to +# overwrite it with a second copy. See `Dagger.slot_may_be_uninit`. +Dagger.can_alloc_uninit(::CUDAVRAMMemorySpace, ::Type{<:CuArray{T}}) where {T} = isbitstype(T) +function Dagger.alloc_uninit(space::CUDAVRAMMemorySpace, ::Type{A}, dims::Dims{N}) where {T,N,A<:CuArray{T,N}} + return with_context(space) do + _note_alloc!(space.device, A(undef, dims)) + end +end + # In-place # N.B. These methods assume that later operations will implicitly or # explicitly synchronize with their associated stream @@ -150,16 +165,35 @@ function Dagger.move!(to_space::Dagger.CPURAMMemorySpace, from_space::CUDAVRAMMe sync_with_context(from_space) with_context!(from_space) end - copyto!(to, from) + if from isa DenseArray + copyto!(to, from) + else + # A strided device view: gather it with a kernel first, since + # `copyto!` would index it element by element from the host + dense = similar(from, size(from)) + dense .= from + copyto!(to, Array(dense)) + end # N.B. DtoH will synchronize return end function Dagger.move!(to_space::CUDAVRAMMemorySpace, from_space::Dagger.CPURAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} with_context!(to_space) - copyto!(to, from) + if to isa DenseArray + copyto!(to, from) + else + # A strided device view: upload densely, then scatter with a kernel + dense = CuArray{T,N}(undef, size(from)) + copyto!(dense, from isa DenseArray ? from : collect(from)) + to .= dense + end return end function Dagger.move!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} + if to_space != from_space && to_space.owner == from_space.owner == myid() + peer_move!(to_space, from_space, to, from) + return + end sync_with_context(from_space) with_context!(to_space) if Dagger.needs_multi_span_copy(to_space, from_space, to, from) @@ -170,11 +204,392 @@ function Dagger.move!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemoryS return end +### Same-process copies between two GPUs +# Two GPUs of one process need not be peers (PCIe-attached GPUs behind a host +# bridge often are not). CUDA.jl's `copyto!` then stages through a fresh +# *pageable* host `Vector` with a synchronous DtoH (~1.4 GB/s on L40S over +# PHB), and a KA kernel touching both buffers is illegal outright. +# `cuMemcpyPeerAsync` works without peer access -- the driver pipelines it +# through its own pinned buffers (~22 GB/s on the same machine). +# +# The copy runs on the destination's `COPY_STREAMS` entry rather than on the +# device stream every Dagger task uses, so it overlaps that device's kernels: +# a GEMM rotating panels between GPUs otherwise spends every transfer with +# the destination's compute stalled behind it. Ordering is per allocation +# (`BUFFER_EVENTS`): the copy stream waits only for the source's last write +# and the destination's last write and read, not for everything queued on +# either device -- the source GPU is usually busy *reading* the very panel +# being sent. The copying task then waits on the host for the copy alone, so +# when Datadeps sees the copy task finish the copy is complete: later +# consumers, writers of the source, frees and host reads need no further +# ordering, and every existing synchronization point stays valid. +# Only addresses are read (as `Dagger.aliasing` does): `pointer` would take +# CUDA.jl stream ownership for the active device, which throws without P2P. +_buffer_key(x::CuArray) = UInt(convert(CUDA.CuPtr{Cvoid}, x.data[].mem)) +_raw_cuaddr(x::CuArray) = _buffer_key(x) + UInt(x.offset * Base.elsize(x)) +_cudevice(x::CuArray) = x.data[].mem.dev.handle +Dagger.data_address(x::CuArray) = UInt64(_raw_cuaddr(x)) + +const COPY_STREAMS = Dict{Int, CuStream}() + +# Outstanding device-stream work per allocation, keyed by base address. +# `write`/`read` are events recorded on the owning device's stream after the +# last task that wrote/read the allocation there; `nothing` means nothing +# outstanding. An allocation with no entry has an unknown history, and a copy +# waits on its whole device stream. Entries for freed allocations go stale; +# a stale event only makes a copy into a reused address wait longer. +# +# A peer copy leaves its destination complete, and the copy task's own record +# in `execute!` need not then add a device-stream event to it, which would +# make later copies out of it wait for unrelated kernels. So the copy stamps +# the entry with its task (`filled_by`), and only that task's record skips +# the write. The stamp must name the task: state like "skip the next write" +# outlives a freed allocation, and swallows the first write to the next +# allocation at its address -- a copy out of that one then reads stale data. +mutable struct BufferEvents + write::Union{CuEvent,Nothing} + read::Union{CuEvent,Nothing} + # The `COPY_TASK` whose peer copy last completed the allocation (0: none) + filled_by::UInt64 +end +# The `move!` task running on this thread of control, numbered per `execute!` +# (0 outside one). Scoped, so `move!`'s own helper tasks (a remotecall to this +# worker runs on a fresh task) still see it. +const COPY_TASK = Dagger.ScopedValue{UInt64}(UInt64(0)) +const COPY_TASK_COUNT = Threads.Atomic{UInt64}(0) +const BUFFER_EVENTS = Dict{UInt, BufferEvents}() +# Per device, the event after the last task with an argument that may hold +# device memory but is not an array of this backend (so any allocation of +# the device may have been touched); copies wait on it as well. +const UNTRACKED_EVENTS = Dict{Int, CuEvent}() +const BUFFER_EVENTS_LOCK = ReentrantLock() +const BUFFER_EVENTS_PRUNE_AT = Ref(4096) + +function _record_on(dev::Int, stream::CuStream) + return context!(CONTEXTS[dev]) do + ev = CUDA.CuEvent(CUDA.EVENT_DISABLE_TIMING) + CUDA.record(ev, stream) + ev + end +end + +_tracked_buffer(x::CuArray) = x +function _tracked_buffer(x) + x isa AbstractArray || return nothing + s = Dagger.storage_array(x) + return s isa CuArray ? s : nothing +end +_maybe_device_memory(x) = + !(isbits(x) || x isa Union{Symbol,AbstractString,Function,Type,Module} || + (x isa Array && isbitstype(eltype(x)))) + +_local_value(x) = x +# A remote chunk is in no allocation of this process +_local_value(x::Chunk) = Dagger.is_local(Dagger.current_acceleration(), x.handle) ? unwrap(x) : nothing + +# Caller holds `BUFFER_EVENTS_LOCK` +function _note_use!(dev::Int, ev::CuEvent, x, written::Bool, copy_task::UInt64=UInt64(0)) + buf = _tracked_buffer(x) + if buf === nothing + _maybe_device_memory(x) && (UNTRACKED_EVENTS[dev] = ev) + return + end + # e.g. Datadeps' `unsafe_free!` tasks; the stale entry is pruned later + buf.data.freed && return + _cudevice(buf) == dev || return + key = _buffer_key(buf) + entry = get(BUFFER_EVENTS, key, nothing) + if written + if entry === nothing + BUFFER_EVENTS[key] = BufferEvents(ev, nothing, 0) + elseif copy_task != 0 && entry.filled_by == copy_task + # Filled by this very task's peer copy, which is complete + else + entry.write = ev + entry.read = nothing + entry.filled_by = 0 + end + elseif entry !== nothing + # Reads on one stream are ordered, so the latest stands for all + entry.read = ev + end + # (A read of an allocation with no entry stays unknown: a later copy + # into it then waits on the whole device stream, which covers the read.) + return +end + +_event_pending(ev::Union{CuEvent,Nothing}) = ev !== nothing && !CUDA.isdone(ev) +function _prune_buffer_events!() + length(BUFFER_EVENTS) < BUFFER_EVENTS_PRUNE_AT[] && return + filter!(BUFFER_EVENTS) do (_, entry) + _event_pending(entry.write) || _event_pending(entry.read) + end + BUFFER_EVENTS_PRUNE_AT[] = max(4096, 2 * length(BUFFER_EVENTS)) + return +end + +# Record what a task launched on `dev` used, once its work is enqueued. +# `writes` lists the written positional arguments (the `writes` option); +# `nothing` (no Datadeps) treats every argument as possibly written. +function _record_task!(dev::Int, @nospecialize(f), @nospecialize(args::Tuple), @nospecialize(result), + writes::Union{Vector{Int},Nothing}, copy_task::UInt64) + ev = _record_on(dev, STREAMS[dev]) + @lock BUFFER_EVENTS_LOCK begin + if Dagger.is_move_task(f) && length(args) == 5 + # move!(dep_mod, to_space, from_space, to, from), whose array + # arguments arrive as `Chunk`s (Datadeps spawns copies `meta`) + _note_use!(dev, ev, _local_value(args[4]), true, copy_task) + _note_use!(dev, ev, _local_value(args[5]), false) + else + for i in eachindex(args) + _note_use!(dev, ev, args[i], writes === nothing || i in writes) + end + end + _note_use!(dev, ev, result, true) + _prune_buffer_events!() + end + return +end + +# A fresh allocation on `dev`'s stream: copies into it wait for the +# allocation only (it may reuse memory freed by work still queued there) +function _note_alloc!(dev::Int, x::CuArray) + ev = _record_on(dev, STREAMS[dev]) + @lock BUFFER_EVENTS_LOCK begin + BUFFER_EVENTS[_buffer_key(x)] = BufferEvents(ev, nothing, 0) + end + return x +end +# A fresh allocation with no outstanding device-stream work +function _note_complete!(x::CuArray) + @lock BUFFER_EVENTS_LOCK begin + BUFFER_EVENTS[_buffer_key(x)] = BufferEvents(nothing, nothing, 0) + end + return x +end + +# `x`'s outstanding events now, before ordering work after them +function _events_snapshot(x::CuArray) + @lock BUFFER_EVENTS_LOCK begin + entry = get(BUFFER_EVENTS, _buffer_key(x), nothing) + return entry === nothing ? (nothing, nothing) : (entry.write, entry.read) + end +end +# `x` was filled by work that completed after the events in `snapshot`: clear +# those, but keep any recorded since (another task may write another part of +# the same allocation meanwhile), and stamp the entry for this task's record. +function _note_filled!(x::CuArray, snapshot) + copy_task = COPY_TASK[] + @lock BUFFER_EVENTS_LOCK begin + key = _buffer_key(x) + entry = get(BUFFER_EVENTS, key, nothing) + if entry === nothing + BUFFER_EVENTS[key] = BufferEvents(nothing, nothing, copy_task) + else + entry.write === snapshot[1] && (entry.write = nothing) + entry.read === snapshot[2] && (entry.read = nothing) + entry.filled_by = copy_task + end + end + return x +end + +# Events a copy reading `from` (on device `s`) into `to` (on device `d`) must +# wait for, and the snapshot of `to`'s events they include (`_note_filled!`) +function _copy_dependencies(s::Int, from::CuArray, d::Int, to::CuArray) + evs = CuEvent[] + local to_snapshot + @lock BUFFER_EVENTS_LOCK begin + entry = get(BUFFER_EVENTS, _buffer_key(from), nothing) + if entry === nothing + push!(evs, _record_on(s, STREAMS[s])) + elseif entry.write !== nothing + push!(evs, entry.write) + end + haskey(UNTRACKED_EVENTS, s) && push!(evs, UNTRACKED_EVENTS[s]) + entry = get(BUFFER_EVENTS, _buffer_key(to), nothing) + if entry === nothing + push!(evs, _record_on(d, STREAMS[d])) + to_snapshot = (nothing, nothing) + else + entry.write !== nothing && push!(evs, entry.write) + entry.read !== nothing && push!(evs, entry.read) + to_snapshot = (entry.write, entry.read) + end + haskey(UNTRACKED_EVENTS, d) && push!(evs, UNTRACKED_EVENTS[d]) + end + # Work CUDA.jl tracks on some other stream (e.g. code outside Dagger) + for (dev, buf) in ((s, from), (d, to)) + managed = buf.data[] + if managed.dirty && managed.stream != STREAMS[dev] && + managed.stream != COPY_STREAMS[dev] + push!(evs, _record_on(dev, managed.stream)) + end + end + return evs, to_snapshot +end + +# Copy `lens[i]` bytes from `src_ptrs[i]` (inside `from`, on `from_space`) +# to `dst_ptrs[i]` (inside `to`, on `to_space`); returns once complete. +function peer_copy_spans!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, + to::CuArray, from::CuArray, dst_ptrs, src_ptrs, lens) + isempty(lens) && return + d, s = to_space.device, from_space.device + @assert _cudevice(to) == d && _cudevice(from) == s + dst_ctx = CONTEXTS[d] + src_ctx = CONTEXTS[s] + copy_stream = COPY_STREAMS[d] + deps, to_snapshot = _copy_dependencies(s, from, d, to) + # Only raw addresses reach the copy: keep both arrays (e.g. the packed + # staging buffers, referenced nowhere else) alive until it completes + GC.@preserve to from begin + context!(dst_ctx) do + for ev in deps + CUDADRV.cuStreamWaitEvent(copy_stream, ev, 0) + end + for i in eachindex(lens) + len = lens[i] + len == 0 && continue + CUDADRV.cuMemcpyPeerAsync(CUDA.CuPtr{Cvoid}(dst_ptrs[i]), dst_ctx, + CUDA.CuPtr{Cvoid}(src_ptrs[i]), src_ctx, + len, copy_stream) + end + done = CUDA.CuEvent(CUDA.EVENT_DISABLE_TIMING) + CUDA.record(done, copy_stream) + CUDA.synchronize(done) + end + end + # Everything the copy waited for is complete too + _note_filled!(to, to_snapshot) + return +end + +# Spans below this average size are packed into one contiguous buffer on each +# side first: every peer copy without P2P is a host-staged transfer with its +# own fixed cost, so a halo's thousands of short column segments would +# otherwise cost thousands of round trips. +const PEER_PACK_MAX_AVG_BYTES = 256 * 1024 +const PEER_PACK_MIN_SPANS = 16 + +function peer_copy_span_pairs!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, + to_s::CuArray{T}, from_s::CuArray{T}, + dst_ptrs::Vector{UInt64}, src_ptrs::Vector{UInt64}, + lens::Vector{UInt64}) where T + n = length(lens) + total = sum(lens; init=UInt64(0)) + total == 0 && return + if n < PEER_PACK_MIN_SPANS || total ÷ n >= PEER_PACK_MAX_AVG_BYTES + peer_copy_spans!(to_space, from_space, to_s, from_s, dst_ptrs, src_ptrs, lens) + return + end + elsize = sizeof(T) + @assert total % elsize == 0 + nelem = Int(total ÷ elsize) + # Gather on the source device, copy one buffer, scatter on the destination. + # (`multi_span_copy!` synchronizes its device after each kernel.) + src_packed = with_context(from_space) do + packed = CuArray{T}(undef, nelem) + base = UInt64(_raw_cuaddr(packed)) + offs = UInt64(0) + packed_ptrs = Vector{UInt64}(undef, n) + for i in 1:n + packed_ptrs[i] = base + offs + offs += lens[i] + end + Dagger.multi_span_copy!(packed, from_s, packed_ptrs, src_ptrs, lens) + _note_complete!(packed) + end + dst_packed = with_context(() -> _note_alloc!(to_space.device, CuArray{T}(undef, nelem)), to_space) + peer_copy_spans!(to_space, from_space, dst_packed, src_packed, + UInt64[_raw_cuaddr(dst_packed)], UInt64[_raw_cuaddr(src_packed)], UInt64[total]) + to_snapshot = _events_snapshot(to_s) + with_context(to_space) do + base = UInt64(_raw_cuaddr(dst_packed)) + offs = UInt64(0) + packed_ptrs = Vector{UInt64}(undef, n) + for i in 1:n + packed_ptrs[i] = base + offs + offs += lens[i] + end + Dagger.multi_span_copy!(to_s, dst_packed, dst_ptrs, packed_ptrs, lens) + end + # The scatter ran on (and synchronized) the stream `to_s`'s events are on. + # The packing buffers die here: forget them, so their addresses come back + # with no history. + _note_filled!(to_s, to_snapshot) + @lock BUFFER_EVENTS_LOCK begin + delete!(BUFFER_EVENTS, _buffer_key(src_packed)) + delete!(BUFFER_EVENTS, _buffer_key(dst_packed)) + end + return +end + +function peer_move!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, + to::AbstractArray{T}, from::AbstractArray{T}) where T + if to isa CuArray && from isa CuArray + @assert length(to) == length(from) + peer_copy_spans!(to_space, from_space, to, from, UInt64[_raw_cuaddr(to)], + UInt64[_raw_cuaddr(from)], UInt64[sizeof(from)]) + return + end + to_spans = Dagger.memory_spans(Dagger.aliasing(to)) + from_spans = Dagger.memory_spans(Dagger.aliasing(from)) + @assert length(to_spans) == length(from_spans) + n = length(to_spans) + dst_ptrs = Vector{UInt64}(undef, n) + src_ptrs = Vector{UInt64}(undef, n) + lens = Vector{UInt64}(undef, n) + for i in 1:n + @assert Dagger.span_len(to_spans[i]) == Dagger.span_len(from_spans[i]) + dst_ptrs[i] = UInt64(Dagger.span_start(to_spans[i])) + src_ptrs[i] = UInt64(Dagger.span_start(from_spans[i])) + lens[i] = UInt64(Dagger.span_len(from_spans[i])) + end + peer_copy_span_pairs!(to_space, from_space, Dagger.storage_array(to), + Dagger.storage_array(from), dst_ptrs, src_ptrs, lens) + return +end + +# Out-of-place copy of `x` to another GPU of this process (complete on return). +function _peer_copy_out(to_proc::CuArrayDeviceProc, from_proc::CuArrayDeviceProc, x::CuArray) + to_space = only(Dagger.memory_spaces(to_proc)) + from_space = only(Dagger.memory_spaces(from_proc)) + to_arr = with_context(() -> _note_alloc!(to_proc.device, similar(x)), to_proc) + peer_move!(to_space, from_space, to_arr, x) + return to_arr +end + +# Remainder copies between two GPUs of this process (see `remainders.jl`). +function Dagger.device_remainder_copy!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, + to_s::CuArray{T}, from_s::CuArray{T}, + spans::Vector{Tuple{Dagger.LocalMemorySpan,Dagger.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 1:n + 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 + peer_copy_span_pairs!(to_space, from_space, to_s, from_s, dst_ptrs, src_ptrs, lens) + return true +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) + buf, token = pin_host!(x) + try + arr = adapt(CuArray, buf) + CUDA.synchronize() + return arr + finally + unpin_host!(token) + end end arr = adapt(CuArray, x) CUDA.synchronize() @@ -186,9 +601,20 @@ function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Chunk) to_w = Dagger.root_worker_id(to_proc) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) + # A chunk labelled with a host processor can already hold device memory: + # Datadeps rebuilding a view of a host array on a GPU passes the parent + # it already moved there. Page-locking that as host memory throws. + cpu_data isa CuArray && return Dagger.move(from_proc, to_proc, cpu_data) with_context(to_proc) do if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) - Dagger.pin_buffer!(:CUDA, cpu_data) + buf, token = pin_host!(cpu_data) + try + arr = adapt(CuArray, buf) + CUDA.synchronize() + return arr + finally + unpin_host!(token) + end end arr = adapt(CuArray, cpu_data) CUDA.synchronize() @@ -244,52 +670,43 @@ function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x return arr 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(CUDA.synchronize, from_proc) - return with_context(to_proc) do - to_arr = similar(from_arr) - copyto!(to_arr, from_arr) - CUDA.synchronize() - return to_arr - end - elseif Dagger.same_node(Dagger.current_acceleration(), from_proc.owner, to_proc.owner) && from_proc.device == to_proc.device - # Same node, we can use IPC - ipc_handle, eT, shape = remotecall_fetch(from_proc.owner, x) do x - arr = unwrap(x) - ipc_handle_ref = Ref{CUDA.CUipcMemHandle}() - GC.@preserve arr begin - CUDA.cuIpcGetMemHandle(ipc_handle_ref, pointer(arr)) - end - (ipc_handle_ref[], eltype(arr), size(arr)) - end - r_ptr = Ref{CUDA.CUdeviceptr}() - CUDA.device!(from_proc.device) do - CUDA.cuIpcOpenMemHandle(r_ptr, ipc_handle, CUDA.CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS) - end - ptr = Base.unsafe_convert(CUDA.CuPtr{eT}, r_ptr[]) - arr = unsafe_wrap(CuArray, ptr, shape; own=false) - finalizer(arr) do arr - CUDA.cuIpcCloseMemHandle(pointer(arr)) - end - if from_proc.device != to_proc.device - return CUDA.device!(to_proc.device) do - to_arr = similar(arr) - copyto!(to_arr, arr) - to_arr + return _peer_copy_out(to_proc, from_proc, unwrap(x)) + elseif ipc_enabled() && Dagger.datasize(x) >= Dagger.IPC_MIN_BYTES[] && + from_proc.device == to_proc.device && + Dagger.same_node(Dagger.current_acceleration(), from_proc.owner, to_proc.owner) + # Pool allocations cannot be exported directly. Keep the exportable + # staging buffer alive on the sender until the receiver has copied it + # into its own allocation; no shared mapping may escape this call. + ref = remotecall_fetch(from_proc.owner, to_proc, x) do to_proc, x + info, token = Dagger.ipc_export(unwrap(x)) + try + return remotecall_fetch(to_proc.owner, to_proc, info) do to_proc, info + with_context(to_proc) do + received = Dagger.ipc_materialize(info) + MemPool.poolset(received; device=MemPool.CPURAMDevice()) + end + end + finally + Dagger.ipc_release!(token) end - else - return arr end + return poolget(ref) else - # Different node, use DtoH, serialization, HtoD (pinned host staging) + # Small transfers, disabled IPC, or different nodes: stage through host. 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 - Dagger.pin_buffer!(:CUDA, host_copy) - return CuArray(host_copy) + buf, token = pin_host!(host_copy) + try + arr = CuArray(buf) + CUDA.synchronize() + return arr + finally + unpin_host!(token) + end end end end @@ -299,20 +716,20 @@ function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x with_context(CUDA.synchronize, from_proc) return x elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) - with_context(CUDA.synchronize, from_proc) - return with_context(to_proc) do - to_arr = similar(x) - copyto!(to_arr, x) - CUDA.synchronize() - return to_arr - end + return _peer_copy_out(to_proc, from_proc, x) else host_copy = with_context(from_proc) do return Dagger.pinned_host_array(x) end return with_context(to_proc) do - Dagger.pin_buffer!(:CUDA, host_copy) - return CuArray(host_copy) + buf, token = pin_host!(host_copy) + try + arr = CuArray(buf) + CUDA.synchronize() + return arr + finally + unpin_host!(token) + end end end end @@ -326,10 +743,22 @@ Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Chunk{T}) where { function Dagger.execute!(proc::CuArrayDeviceProc, f, args...; kwargs...) @nospecialize f args kwargs tls = Dagger.get_tls() + spec = tls.task_spec + writes = spec isa Dagger.Sch.TaskSpec && spec.options !== nothing ? + spec.options.writes : nothing task = Threads.@spawn begin Dagger.set_tls!(tls) with_context!(proc) - result = Base.@invokelatest f(args...; kwargs...) + copy_task = Dagger.is_move_task(f) ? + Threads.atomic_add!(COPY_TASK_COUNT, UInt64(1)) + UInt64(1) : UInt64(0) + result = if copy_task == 0 + Base.@invokelatest f(args...; kwargs...) + else + Dagger.with(COPY_TASK => copy_task) do + Base.@invokelatest f(args...; kwargs...) + end + end + _record_task!(proc.device, f, args, result, writes, copy_task) # N.B. Synchronization must be done when accessing result or args return result end @@ -464,11 +893,129 @@ Dagger.mpi_device_direct(x::CUDA.StridedCuArray) = mpi_gpu_direct_enabled() 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 +# Page-lock host staging buffers (~2x DtoH/HtoD bandwidth), unregistering +# them from a GC finalizer. Dagger.gpu_memory_kind(::CuArray) = :CUDA Dagger.gpu_memory_kind(::CUDAVRAMMemorySpace) = :CUDA -Dagger.pin_buffer!(::Val{:CUDA}, buf::DenseArray) = (CUDA.pin(buf); nothing) +# N.B. Not `CUDA.pin`: its finalizer takes a `ReentrantLock`, and a finalizer +# that finds the lock contended throws "task switch not allowed from inside gc +# finalizer". The unregistration is then lost, while GC frees the memory +# anyway: a freed range stays registered with the driver, a later allocation +# at that address fails to register or is DMA'd through stale pages, and +# multi-GPU runs crashed with a segfault a few iterations in. Our finalizer +# only `trylock`s, and on contention re-arms itself, which keeps the buffer +# alive until a later GC can unregister it. +# Host registration is process-wide, so a buffer already registered through +# another device's context (one host array moved to two GPUs) is page-locked +# either way, and is skipped. +const PIN_LOCK = ReentrantLock() +# Base pointer => (registration, context) for buffers pinned by `pin_buffer!` +const FINALIZER_PINS = Dict{Ptr{Cvoid}, Tuple{Any,CuContext}}() +function Dagger.pin_buffer!(::Val{:CUDA}, buf::DenseArray) + isempty(buf) && return + ptr = convert(Ptr{Cvoid}, pointer(buf)) + @lock PIN_LOCK begin + CUDADRV.is_pinned(ptr) && return + mem = try_register(ptr, sizeof(buf)) + mem === nothing && return + FINALIZER_PINS[ptr] = (mem, context()) + end + finalizer(_unpin_finalizer, buf) + return +end +function _unpin_finalizer(buf) + if !trylock(PIN_LOCK) + finalizer(_unpin_finalizer, buf) + return + end + try + entry = pop!(FINALIZER_PINS, convert(Ptr{Cvoid}, pointer(buf)), nothing) + entry === nothing && return + mem, ctx = entry + context!(ctx) do + CUDADRV.unregister(mem) + end + catch err + Core.println("Dagger CUDAExt: failed to unregister a pinned host buffer: ", err) + finally + unlock(PIN_LOCK) + end + return +end + +# Scoped host registration for one-shot HtoD sources. +# `CUDA.pin` ties unregistration to a GC finalizer, so a registration can +# outlive the transfer it was made for and collide with whatever the +# allocator places at that address next (ALREADY_REGISTERED, or INVALID_VALUE +# if the stale finalizer unregisters a newer owner's range mid-copy). Here the +# registration is refcounted by us, keyed by base pointer (registration is +# process-wide, not per context), and released by `unpin_host!` as soon as the +# copy has completed, while the caller still holds the buffer. +# A range already registered by anyone else (exact match through +# `Dagger.pin_buffer!`, or an overlapping one) is never adopted or extended: +# the caller gets a private fresh copy to register instead, since a partially +# registered range is invalid for CUDA copies. +mutable struct HostPinEntry + mem::Any + ctx::CuContext + size::Int + count::Int +end +const HOST_PINS = Dict{Ptr{Cvoid}, HostPinEntry}() + +function try_register(ptr, sz) + try + return CUDADRV.register(CUDADRV.HostMemory, ptr, sz) + catch err + err isa CUDA.CuError && err.code == CUDA.ERROR_HOST_MEMORY_ALREADY_REGISTERED || + rethrow() + return nothing + end +end + +# Returns (buf, token): copy from `buf` (which may be a private copy of the +# argument), then call `unpin_host!(token)` once the copy has completed. +# `token === nothing` means nothing to release. +function pin_host!(buf::DenseArray) + isempty(buf) && return buf, nothing + ptr = convert(Ptr{Cvoid}, pointer(buf)) + sz = sizeof(buf) + @lock PIN_LOCK begin + entry = get(HOST_PINS, ptr, nothing) + if entry !== nothing && entry.size == sz + entry.count += 1 + return buf, ptr + end + # registration fails on *any* overlap, including a registration + # strictly inside the range that endpoint probes cannot see + mem = entry === nothing ? try_register(ptr, sz) : nothing + if mem === nothing + # Registered by someone else, or overlapping a registration we + # cannot reason about (ours of a different size, a + # finalizer-managed one): use a fresh, unregistered copy. + buf = copy(buf) + ptr = convert(Ptr{Cvoid}, pointer(buf)) + sz = sizeof(buf) + mem = CUDADRV.register(CUDADRV.HostMemory, ptr, sz) + end + HOST_PINS[ptr] = HostPinEntry(mem, context(), sz, 1) + end + return buf, ptr +end +function unpin_host!(token) + token === nothing && return + @lock PIN_LOCK begin + entry = HOST_PINS[token] + entry.count -= 1 + if entry.count == 0 + delete!(HOST_PINS, token) + context!(entry.ctx) do + CUDADRV.unregister(entry.mem) + end + end + end + return +end # 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 @@ -547,6 +1094,7 @@ function __init__() CONTEXTS[dev.handle] = ctx context!(ctx) do STREAMS[dev.handle] = stream() + COPY_STREAMS[dev.handle] = CuStream(; flags=CUDA.STREAM_NON_BLOCKING) end return proc end diff --git a/ext/IntelExt.jl b/ext/IntelExt.jl index 12632104d..1c207eaee 100644 --- a/ext/IntelExt.jl +++ b/ext/IntelExt.jl @@ -180,13 +180,28 @@ function Dagger.move!(to_space::Dagger.CPURAMMemorySpace, from_space::IntelVRAMM sync_with_context(from_space) with_context!(from_space) end - copyto!(to, from) + if from isa DenseArray + copyto!(to, from) + else + # A strided device view: gather it with a kernel first, since + # `copyto!` would index it element by element from the host + dense = similar(from, size(from)) + dense .= from + copyto!(to, Array(dense)) + end # N.B. DtoH will synchronize return end function Dagger.move!(to_space::IntelVRAMMemorySpace, from_space::Dagger.CPURAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} with_context!(to_space) - copyto!(to, from) + if to isa DenseArray + copyto!(to, from) + else + # A strided device view: upload densely, then scatter with a kernel + dense = oneArray{T,N}(undef, size(from)) + copyto!(dense, from isa DenseArray ? from : collect(from)) + to .= dense + end return end function Dagger.move!(to_space::IntelVRAMMemorySpace, from_space::IntelVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} @@ -216,6 +231,10 @@ function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::Chunk) to_w = Dagger.root_worker_id(to_proc) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) + # A chunk labelled with a host processor can already hold device memory: + # Datadeps rebuilding a view of a host array on a GPU passes the parent + # it already moved there. Treating that as host memory throws. + cpu_data isa oneArray && return Dagger.move(from_proc, to_proc, cpu_data) with_context(to_proc) do if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) Dagger.pin_buffer!(:oneAPI, cpu_data) @@ -434,18 +453,13 @@ 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() +# Level Zero export/import primitives. Stage into a dedicated device allocation +# so an imported mapping does not depend on the original array's lifetime. +# Level Zero's opaque handle contains a process-local file descriptor. MPI's +# serialized metadata transports the integer, not the descriptor (SCM_RIGHTS), +# so zeMemOpenIpcHandle on the receiver cannot import it. Use host staging until +# the transport supports descriptor passing; DAGGER_IPC cannot make this safe. +Dagger.ipc_eligible(::IntelVRAMMemorySpace, ::IntelVRAMMemorySpace) = false const oneL0 = oneAPI.oneL0 diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 11ecf9313..bcd3cc623 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -43,9 +43,7 @@ function Dagger.memory_space(x::MtlArray) return MetalVRAMMemorySpace(myid(), device_id) end _device_id(dev::MtlDevice) = findfirst(other_dev->other_dev === dev, Metal.devices()) -function Dagger.aliasing(x::MtlArray{T}) where T - space = Dagger.memory_space(x) - S = typeof(space) +function _mtl_addr(x::MtlArray) mtl_ptr = pointer(x) # Metal ≥1.10 defines `UInt(::MtlPtr)` as the GPU virtual address, which # `span_copy` uses via `UInt64(pointer(x))`. Metal <1.10 has no such method; @@ -55,9 +53,17 @@ function Dagger.aliasing(x::MtlArray{T}) where T else Metal.contents(mtl_ptr.buffer) + mtl_ptr.offset end - rptr = Dagger.RemotePtr{Cvoid}(UInt64(addr), space) + return UInt64(addr) +end +function Dagger.aliasing(x::MtlArray{T}) where T + space = Dagger.memory_space(x) + S = typeof(space) + rptr = Dagger.RemotePtr{Cvoid}(_mtl_addr(x), space) return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end +# Views locate themselves from their parent's address, which must match the +# address the parent's own aliasing reports +Dagger.data_address(x::MtlArray) = _mtl_addr(x) # MPI (SPMD) integration: stamp owning rank on VRAM spaces Dagger.mpi_remap_space(space::MetalVRAMMemorySpace, owner::Int) = @@ -186,13 +192,28 @@ function Dagger.move!(to_space::Dagger.CPURAMMemorySpace, from_space::MetalVRAMM sync_with_context(from_space) with_context!(from_space) end - copyto!(to, from) + if from isa DenseArray + copyto!(to, from) + else + # A strided device view: gather it with a kernel first, since + # `copyto!` would index it element by element from the host + dense = similar(from, size(from)) + dense .= from + copyto!(to, Array(dense)) + end # N.B. DtoH will synchronize return end function Dagger.move!(to_space::MetalVRAMMemorySpace, from_space::Dagger.CPURAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} with_context!(to_space) - copyto!(to, from) + if to isa DenseArray + copyto!(to, from) + else + # A strided device view: upload densely, then scatter with a kernel + dense = MtlArray{T,N}(undef, size(from)) + copyto!(dense, from isa DenseArray ? from : collect(from)) + to .= dense + end return end function Dagger.move!(to_space::MetalVRAMMemorySpace, from_space::MetalVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} @@ -222,6 +243,10 @@ function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::Chunk) to_w = Dagger.root_worker_id(to_proc) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) + # A chunk labelled with a host processor can already hold device memory: + # Datadeps rebuilding a view of a host array on a GPU passes the parent + # it already moved there. Treating that as host memory throws. + cpu_data isa MtlArray && return Dagger.move(from_proc, to_proc, cpu_data) with_context(to_proc) do if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) Dagger.pin_buffer!(:Metal, cpu_data) diff --git a/ext/OpenCLExt.jl b/ext/OpenCLExt.jl index 3d4711730..74cbb1479 100644 --- a/ext/OpenCLExt.jl +++ b/ext/OpenCLExt.jl @@ -180,13 +180,28 @@ function Dagger.move!(to_space::Dagger.CPURAMMemorySpace, from_space::CLMemorySp _sync_with_context(from_space) with_context!(from_space) end - copyto!(to, from) + if from isa DenseArray + copyto!(to, from) + else + # A strided device view: gather it with a kernel first, since + # `copyto!` would index it element by element from the host + dense = similar(from, size(from)) + dense .= from + copyto!(to, Array(dense)) + end # N.B. DtoH will synchronize return end function Dagger.move!(to_space::CLMemorySpace, from_space::Dagger.CPURAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} with_context!(to_space) - copyto!(to, from) + if to isa DenseArray + copyto!(to, from) + else + # A strided device view: upload densely, then scatter with a kernel + dense = CLArray{T,N}(undef, size(from)) + copyto!(dense, from isa DenseArray ? from : collect(from)) + to .= dense + end return end function Dagger.move!(to_space::CLMemorySpace, from_space::CLMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} @@ -216,6 +231,10 @@ function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Chunk) to_w = Dagger.root_worker_id(to_proc) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) + # A chunk labelled with a host processor can already hold device memory: + # Datadeps rebuilding a view of a host array on a GPU passes the parent + # it already moved there. Treating that as host memory throws. + cpu_data isa CLArray && return Dagger.move(from_proc, to_proc, cpu_data) with_context(to_proc) do if cpu_data isa DenseArray && isbitstype(eltype(cpu_data)) Dagger.pin_buffer!(:OpenCL, cpu_data) diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index a2b77a68e..805613702 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -17,6 +17,7 @@ else import ..AMDGPU end import AMDGPU: HIPDevice, HIPContext, HIPStream, ROCArray, ROCBackend +import AMDGPU.HIP: HIPEvent import AMDGPU: devices, context, context!, stream, stream! import AMDGPU: rocBLAS, rocSOLVER @@ -35,11 +36,34 @@ Dagger.memory_space(x::ROCArray) = function Dagger.aliasing(x::ROCArray{T}) where T space = Dagger.memory_space(x) S = typeof(space) - gpu_ptr = pointer(x) - rptr = Dagger.RemotePtr{Cvoid}(UInt64(gpu_ptr), space) + # N.B. Not `pointer(x)`: that takes AMDGPU stream ownership of `x` for the + # calling task's stream, synchronizing the stream `x` was last used on and + # re-stamping ownership as a side effect. Aliasing only needs the address. + rptr = Dagger.RemotePtr{Cvoid}(UInt64(_raw_rocaddr(x)), space) return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x))) end +# Device addresses for aliasing, read without `pointer` (see above). +function _buffer_key(x::ROCArray) + mem = x.buf[].mem + return UInt(mem isa AMDGPU.Mem.HIPBuffer ? mem.ptr : mem.dev_ptr) +end +# `x.offset` counts elements in older AMDGPU.jl (e.g. 2.1) and bytes in newer +# ones (e.g. 2.8). Ask this version's own `derive`, once: deriving a view one +# `Float64` in reports an offset of 1 or of 8. +const OFFSET_IN_BYTES = Ref{Union{Bool,Nothing}}(nothing) +function _offset_in_bytes() + inbytes = OFFSET_IN_BYTES[] + inbytes === nothing || return inbytes + probe = ROCArray{Float64}(undef, 0) + inbytes = AMDGPU.GPUArrays.derive(Float64, probe, (0,), 1).offset == sizeof(Float64) + OFFSET_IN_BYTES[] = inbytes + return inbytes +end +_raw_rocaddr(x::ROCArray) = + _buffer_key(x) + UInt(_offset_in_bytes() ? x.offset : x.offset * Base.elsize(x)) +Dagger.data_address(x::ROCArray) = UInt64(_raw_rocaddr(x)) + # 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) @@ -47,23 +71,59 @@ 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 +# Page-lock host staging buffers (~2x DtoH/HtoD bandwidth), unregistering +# them from a GC finalizer. Dagger.gpu_memory_kind(::ROCArray) = :ROC Dagger.gpu_memory_kind(::ROCVRAMMemorySpace) = :ROC +# N.B. Not `AMDGPU.Mem.pin`/`unpin`: they take AMDGPU's registration lock (a +# `ReentrantLock`), and a finalizer that finds it contended throws "task switch +# not allowed from inside gc finalizer". The unregistration is then lost, while +# GC frees the memory anyway: a freed range stays registered with the driver, +# and a later HtoD copy from memory the allocator placed there segfaulted +# inside `hipMemcpyWithStream`. We register with HIP directly, under our own +# `PIN_LOCK`, and our finalizer only `trylock`s, and on contention re-arms +# itself, which keeps the buffer alive until a later GC can unregister it. +const PIN_LOCK = ReentrantLock() +# Base pointers of buffers registered by `pin_buffer!` +const FINALIZER_PINS = Set{Ptr{Cvoid}}() function Dagger.pin_buffer!(::Val{:ROC}, buf::DenseArray) + isempty(buf) && return 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 + @lock PIN_LOCK begin + # Registered by someone else, and page-locked either way + AMDGPU.Mem.is_pinned(ptr) && return + AMDGPU.HIP.hipHostRegister(ptr, sizeof(buf), AMDGPU.HIP.hipHostRegisterMapped) + push!(FINALIZER_PINS, ptr) + end + finalizer(_unpin_finalizer, buf) + return +end +function _unpin_finalizer(buf) + if !trylock(PIN_LOCK) + finalizer(_unpin_finalizer, buf) + return end - return nothing + try + ptr = Ptr{Cvoid}(pointer(buf)) + ptr in FINALIZER_PINS || return + delete!(FINALIZER_PINS, ptr) + AMDGPU.HIP.hipHostUnregister(ptr) + catch err + Core.println("Dagger ROCExt: failed to unregister a pinned host buffer: ", err) + finally + unlock(PIN_LOCK) + end + return end +# N.B. HtoD sources are not registered: HIP stages pageable uploads through +# its own pinned buffers at the same bandwidth (~11 GB/s here, 1-64 MiB), and +# registering the caller's array for one copy corrupted data. Under 2-GPU +# Datadeps GEMM, tiles read back NaN in 6 of 6 stress runs while +# `hipHostRegister`/`hipHostUnregister` bracketed each upload of a host tile +# that other tasks were concurrently reading (a private copy, which nobody +# else touches, was clean but 10x slower to upload). + function Dagger.unsafe_free!(x::ROCArray) AMDGPU.unsafe_free!(x) return @@ -210,16 +270,35 @@ function Dagger.move!(to_space::Dagger.CPURAMMemorySpace, from_space::ROCVRAMMem sync_with_context(from_space) with_context!(from_space) end - copyto!(to, from) + if from isa DenseArray + copyto!(to, from) + else + # A strided device view: gather it with a kernel first, since + # `copyto!` would index it element by element from the host + dense = similar(from, size(from)) + dense .= from + copyto!(to, Array(dense)) + end # N.B. DtoH will synchronize return end function Dagger.move!(to_space::ROCVRAMMemorySpace, from_space::Dagger.CPURAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} with_context!(to_space) - copyto!(to, from) + if to isa DenseArray + copyto!(to, from) + else + # A strided device view: upload densely, then scatter with a kernel + dense = ROCArray{T,N}(undef, size(from)) + copyto!(dense, from isa DenseArray ? from : collect(from)) + to .= dense + end return end function Dagger.move!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N} + if to_space != from_space && to_space.owner == from_space.owner == myid() + peer_move!(to_space, from_space, to, from) + return + end sync_with_context(from_space) with_context!(to_space) if Dagger.needs_multi_span_copy(to_space, from_space, to, from) @@ -230,12 +309,373 @@ function Dagger.move!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpa return end +### Same-process copies between two GPUs +# Unlike CUDA.jl, HIP does not stage these copies through pageable host +# memory: without peer access (RX 6900 XTs over PCIe have none), AMDGPU's +# `copyto!` and `hipMemcpyPeerAsync` both reach ~5.5 GB/s on PCIe 3.0. But the +# old path synchronized the source's whole device stream first and queued the +# copy on the destination's device stream, so every transfer waited for both +# GPUs' unrelated kernels, and Datadeps' remainder copies took a host round +# trip instead. +# +# The copy runs on the destination's `COPY_STREAMS` entry rather than on the +# device stream every Dagger task uses, so it overlaps that device's kernels. +# Ordering is per allocation (`BUFFER_EVENTS`): the copy stream waits only for +# the source's last write and the destination's last write and read, not for +# everything queued on either device -- the source GPU is usually busy +# *reading* the very panel being sent. The copying task then waits on the host +# for the copy alone, so when Datadeps sees the copy task finish the copy is +# complete: later consumers, writers of the source, frees and host reads need +# no further ordering, and every existing synchronization point stays valid. +# Only raw addresses are read (as `Dagger.aliasing` does): `pointer` would +# take AMDGPU stream ownership and synchronize the array's last stream. +_rocdevice(x::ROCArray) = AMDGPU.device(x).device_id + +const COPY_STREAMS = Dict{Int, HIPStream}() + +# Outstanding device-stream work per allocation, keyed by base address. +# `write`/`read` are events recorded on the owning device's stream after the +# last task that wrote/read the allocation there; `nothing` means nothing +# outstanding. An allocation with no entry has an unknown history, and a copy +# waits on its whole device stream. Entries for freed allocations go stale; +# a stale event only makes a copy into a reused address wait longer. +# +# A peer copy leaves its destination complete, and the copy task's own record +# in `execute!` need not then add a device-stream event to it, which would +# make later copies out of it wait for unrelated kernels. So the copy stamps +# the entry with its task (`filled_by`), and only that task's record skips +# the write. The stamp must name the task: state like "skip the next write" +# outlives a freed allocation, and swallows the first write to the next +# allocation at its address -- a copy out of that one then reads stale data. +mutable struct BufferEvents + write::Union{HIPEvent,Nothing} + read::Union{HIPEvent,Nothing} + # The `COPY_TASK` whose peer copy last completed the allocation (0: none) + filled_by::UInt64 +end +# The `move!` task running on this thread of control, numbered per `execute!` +# (0 outside one). Scoped, so `move!`'s own helper tasks (a remotecall to this +# worker runs on a fresh task) still see it. +const COPY_TASK = Dagger.ScopedValue{UInt64}(UInt64(0)) +const COPY_TASK_COUNT = Threads.Atomic{UInt64}(0) +const BUFFER_EVENTS = Dict{UInt, BufferEvents}() +# Per device, the event after the last task with an argument that may hold +# device memory but is not an array of this backend (so any allocation of +# the device may have been touched); copies wait on it as well. +const UNTRACKED_EVENTS = Dict{Int, HIPEvent}() +const BUFFER_EVENTS_LOCK = ReentrantLock() +const BUFFER_EVENTS_PRUNE_AT = Ref(4096) + +# Events are created on the stream's device +_record_on(dev::Int, stream::HIPStream) = + with_context(() -> HIPEvent(stream), dev) + +_tracked_buffer(x::ROCArray) = x +function _tracked_buffer(x) + x isa AbstractArray || return nothing + s = Dagger.storage_array(x) + return s isa ROCArray ? s : nothing +end +_maybe_device_memory(x) = + !(isbits(x) || x isa Union{Symbol,AbstractString,Function,Type,Module} || + (x isa Array && isbitstype(eltype(x)))) + +_local_value(x) = x +# A remote chunk is in no allocation of this process +_local_value(x::Chunk) = Dagger.is_local(Dagger.current_acceleration(), x.handle) ? unwrap(x) : nothing + +# Caller holds `BUFFER_EVENTS_LOCK` +function _note_use!(dev::Int, ev::HIPEvent, x, written::Bool, copy_task::UInt64=UInt64(0)) + buf = _tracked_buffer(x) + if buf === nothing + _maybe_device_memory(x) && (UNTRACKED_EVENTS[dev] = ev) + return + end + # e.g. Datadeps' `unsafe_free!` tasks; the stale entry is pruned later + buf.buf.freed && return + _rocdevice(buf) == dev || return + key = _buffer_key(buf) + entry = get(BUFFER_EVENTS, key, nothing) + if written + if entry === nothing + BUFFER_EVENTS[key] = BufferEvents(ev, nothing, 0) + elseif copy_task != 0 && entry.filled_by == copy_task + # Filled by this very task's peer copy, which is complete + else + entry.write = ev + entry.read = nothing + entry.filled_by = 0 + end + elseif entry !== nothing + # Reads on one stream are ordered, so the latest stands for all + entry.read = ev + end + # (A read of an allocation with no entry stays unknown: a later copy + # into it then waits on the whole device stream, which covers the read.) + return +end + +_event_pending(ev::Union{HIPEvent,Nothing}) = ev !== nothing && !AMDGPU.HIP.isdone(ev) +function _prune_buffer_events!() + length(BUFFER_EVENTS) < BUFFER_EVENTS_PRUNE_AT[] && return + filter!(BUFFER_EVENTS) do (_, entry) + _event_pending(entry.write) || _event_pending(entry.read) + end + BUFFER_EVENTS_PRUNE_AT[] = max(4096, 2 * length(BUFFER_EVENTS)) + return +end + +# Record what a task launched on `dev` used, once its work is enqueued. +# `writes` lists the written positional arguments (the `writes` option); +# `nothing` (no Datadeps) treats every argument as possibly written. +function _record_task!(dev::Int, @nospecialize(f), @nospecialize(args::Tuple), @nospecialize(result), + writes::Union{Vector{Int},Nothing}, copy_task::UInt64) + ev = _record_on(dev, STREAMS[dev]) + @lock BUFFER_EVENTS_LOCK begin + if Dagger.is_move_task(f) && length(args) == 5 + # move!(dep_mod, to_space, from_space, to, from), whose array + # arguments arrive as `Chunk`s (Datadeps spawns copies `meta`) + _note_use!(dev, ev, _local_value(args[4]), true, copy_task) + _note_use!(dev, ev, _local_value(args[5]), false) + else + for i in eachindex(args) + _note_use!(dev, ev, args[i], writes === nothing || i in writes) + end + end + _note_use!(dev, ev, result, true) + _prune_buffer_events!() + end + return +end + +# A fresh allocation on `dev`'s stream: copies into it wait for the +# allocation only (it may reuse memory freed by work still queued there) +function _note_alloc!(dev::Int, x::ROCArray) + ev = _record_on(dev, STREAMS[dev]) + @lock BUFFER_EVENTS_LOCK begin + BUFFER_EVENTS[_buffer_key(x)] = BufferEvents(ev, nothing, 0) + end + return x +end +# A fresh allocation with no outstanding device-stream work +function _note_complete!(x::ROCArray) + @lock BUFFER_EVENTS_LOCK begin + BUFFER_EVENTS[_buffer_key(x)] = BufferEvents(nothing, nothing, 0) + end + return x +end + +# `x`'s outstanding events now, before ordering work after them +function _events_snapshot(x::ROCArray) + @lock BUFFER_EVENTS_LOCK begin + entry = get(BUFFER_EVENTS, _buffer_key(x), nothing) + return entry === nothing ? (nothing, nothing) : (entry.write, entry.read) + end +end +# `x` was filled by work that completed after the events in `snapshot`: clear +# those, but keep any recorded since (another task may write another part of +# the same allocation meanwhile), and stamp the entry for this task's record. +function _note_filled!(x::ROCArray, snapshot) + copy_task = COPY_TASK[] + @lock BUFFER_EVENTS_LOCK begin + key = _buffer_key(x) + entry = get(BUFFER_EVENTS, key, nothing) + if entry === nothing + BUFFER_EVENTS[key] = BufferEvents(nothing, nothing, copy_task) + else + entry.write === snapshot[1] && (entry.write = nothing) + entry.read === snapshot[2] && (entry.read = nothing) + entry.filled_by = copy_task + end + end + return x +end + +# Events a copy reading `from` (on device `s`) into `to` (on device `d`) must +# wait for, and the snapshot of `to`'s events they include (`_note_filled!`) +function _copy_dependencies(s::Int, from::ROCArray, d::Int, to::ROCArray) + evs = HIPEvent[] + local to_snapshot + @lock BUFFER_EVENTS_LOCK begin + entry = get(BUFFER_EVENTS, _buffer_key(from), nothing) + if entry === nothing + push!(evs, _record_on(s, STREAMS[s])) + elseif entry.write !== nothing + push!(evs, entry.write) + end + haskey(UNTRACKED_EVENTS, s) && push!(evs, UNTRACKED_EVENTS[s]) + entry = get(BUFFER_EVENTS, _buffer_key(to), nothing) + if entry === nothing + push!(evs, _record_on(d, STREAMS[d])) + to_snapshot = (nothing, nothing) + else + entry.write !== nothing && push!(evs, entry.write) + entry.read !== nothing && push!(evs, entry.read) + to_snapshot = (entry.write, entry.read) + end + haskey(UNTRACKED_EVENTS, d) && push!(evs, UNTRACKED_EVENTS[d]) + end + # Work AMDGPU tracks on some other stream (e.g. code outside Dagger) + for (dev, buf) in ((s, from), (d, to)) + managed = buf.buf[] + if managed.dirty && managed.stream.stream != STREAMS[dev].stream && + managed.stream.stream != COPY_STREAMS[dev].stream + push!(evs, _record_on(dev, managed.stream)) + end + end + return evs, to_snapshot +end + +# Copy `lens[i]` bytes from `src_ptrs[i]` (inside `from`, on `from_space`) +# to `dst_ptrs[i]` (inside `to`, on `to_space`); returns once complete. +function peer_copy_spans!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, + to::ROCArray, from::ROCArray, dst_ptrs, src_ptrs, lens) + isempty(lens) && return + d, s = to_space.device_id, from_space.device_id + @assert _rocdevice(to) == d && _rocdevice(from) == s + copy_stream = COPY_STREAMS[d] + deps, to_snapshot = _copy_dependencies(s, from, d, to) + # Only raw addresses reach the copy: keep both arrays (e.g. the packed + # staging buffers, referenced nowhere else) alive until it completes + GC.@preserve to from begin + with_context(d) do + for ev in deps + AMDGPU.HIP.hipStreamWaitEvent(copy_stream.stream, ev.handle, 0) + end + # (HIP numbers devices from 0, AMDGPU from 1) + for i in eachindex(lens) + len = lens[i] + len == 0 && continue + AMDGPU.HIP.hipMemcpyPeerAsync(Ptr{Cvoid}(dst_ptrs[i]), d - 1, + Ptr{Cvoid}(src_ptrs[i]), s - 1, + len, copy_stream.stream) + end + AMDGPU.HIP.synchronize(HIPEvent(copy_stream)) + end + end + # Everything the copy waited for is complete too + _note_filled!(to, to_snapshot) + return +end + +# Spans below this average size are packed into one contiguous buffer on each +# side first: every peer copy without P2P is a host-staged transfer with its +# own fixed cost, so a halo's thousands of short column segments would +# otherwise cost thousands of round trips. +const PEER_PACK_MAX_AVG_BYTES = 256 * 1024 +const PEER_PACK_MIN_SPANS = 16 + +function peer_copy_span_pairs!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, + to_s::ROCArray{T}, from_s::ROCArray{T}, + dst_ptrs::Vector{UInt64}, src_ptrs::Vector{UInt64}, + lens::Vector{UInt64}) where T + n = length(lens) + total = sum(lens; init=UInt64(0)) + total == 0 && return + if n < PEER_PACK_MIN_SPANS || total ÷ n >= PEER_PACK_MAX_AVG_BYTES + peer_copy_spans!(to_space, from_space, to_s, from_s, dst_ptrs, src_ptrs, lens) + return + end + elsize = sizeof(T) + @assert total % elsize == 0 + nelem = Int(total ÷ elsize) + # Gather on the source device, copy one buffer, scatter on the destination. + # (`multi_span_copy!` synchronizes its device after each kernel.) + src_packed = with_context(from_space) do + packed = ROCArray{T}(undef, nelem) + base = UInt64(_raw_rocaddr(packed)) + offs = UInt64(0) + packed_ptrs = Vector{UInt64}(undef, n) + for i in 1:n + packed_ptrs[i] = base + offs + offs += lens[i] + end + Dagger.multi_span_copy!(packed, from_s, packed_ptrs, src_ptrs, lens) + _note_complete!(packed) + end + dst_packed = with_context(() -> _note_alloc!(to_space.device_id, ROCArray{T}(undef, nelem)), to_space) + peer_copy_spans!(to_space, from_space, dst_packed, src_packed, + UInt64[_raw_rocaddr(dst_packed)], UInt64[_raw_rocaddr(src_packed)], UInt64[total]) + to_snapshot = _events_snapshot(to_s) + with_context(to_space) do + base = UInt64(_raw_rocaddr(dst_packed)) + offs = UInt64(0) + packed_ptrs = Vector{UInt64}(undef, n) + for i in 1:n + packed_ptrs[i] = base + offs + offs += lens[i] + end + Dagger.multi_span_copy!(to_s, dst_packed, dst_ptrs, packed_ptrs, lens) + end + # The scatter ran on (and synchronized) the stream `to_s`'s events are on. + # The packing buffers die here: forget them, so their addresses come back + # with no history. + _note_filled!(to_s, to_snapshot) + @lock BUFFER_EVENTS_LOCK begin + delete!(BUFFER_EVENTS, _buffer_key(src_packed)) + delete!(BUFFER_EVENTS, _buffer_key(dst_packed)) + end + return +end + +function peer_move!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, + to::AbstractArray{T}, from::AbstractArray{T}) where T + if to isa ROCArray && from isa ROCArray + @assert length(to) == length(from) + peer_copy_spans!(to_space, from_space, to, from, UInt64[_raw_rocaddr(to)], + UInt64[_raw_rocaddr(from)], UInt64[sizeof(from)]) + return + end + to_spans = Dagger.memory_spans(Dagger.aliasing(to)) + from_spans = Dagger.memory_spans(Dagger.aliasing(from)) + @assert length(to_spans) == length(from_spans) + n = length(to_spans) + dst_ptrs = Vector{UInt64}(undef, n) + src_ptrs = Vector{UInt64}(undef, n) + lens = Vector{UInt64}(undef, n) + for i in 1:n + @assert Dagger.span_len(to_spans[i]) == Dagger.span_len(from_spans[i]) + dst_ptrs[i] = UInt64(Dagger.span_start(to_spans[i])) + src_ptrs[i] = UInt64(Dagger.span_start(from_spans[i])) + lens[i] = UInt64(Dagger.span_len(from_spans[i])) + end + peer_copy_span_pairs!(to_space, from_space, Dagger.storage_array(to), + Dagger.storage_array(from), dst_ptrs, src_ptrs, lens) + return +end + +# Out-of-place copy of `x` to another GPU of this process (complete on return). +# The source device is read from `x` itself, which is what the copy reads. +function _peer_copy_out(to_proc::ROCArrayDeviceProc, x::ROCArray) + to_space = only(Dagger.memory_spaces(to_proc)) + from_space = ROCVRAMMemorySpace(myid(), _rocdevice(x)) + to_arr = with_context(() -> _note_alloc!(to_proc.device_id, similar(x)), to_proc) + peer_move!(to_space, from_space, to_arr, x) + return to_arr +end + +# Remainder copies between two GPUs of this process (see `remainders.jl`). +function Dagger.device_remainder_copy!(to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, + to_s::ROCArray{T}, from_s::ROCArray{T}, + spans::Vector{Tuple{Dagger.LocalMemorySpan,Dagger.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 1:n + 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 + peer_copy_span_pairs!(to_space, from_space, to_s, from_s, dst_ptrs, src_ptrs, lens) + return true +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 @@ -246,10 +686,11 @@ function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Chunk) to_w = Dagger.root_worker_id(to_proc) @assert myid() == to_w cpu_data = remotecall_fetch(unwrap, from_w, x) + # A chunk labelled with a host processor can already hold device memory: + # Datadeps rebuilding a view of a host array on a GPU passes the parent + # it already moved there. Treating that as host memory throws. + cpu_data isa ROCArray && return Dagger.move(from_proc, to_proc, cpu_data) 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 @@ -304,15 +745,7 @@ function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, return arr 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) - dev = AMDGPU.device(from_arr) - with_context(AMDGPU.synchronize, dev.device_id) - return with_context(to_proc) do - to_arr = similar(from_arr) - copyto!(to_arr, from_arr) - AMDGPU.synchronize() - to_arr - end + return _peer_copy_out(to_proc, unwrap(x)) else # Different node, use DtoH, serialization, HtoD (pinned host staging) host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x @@ -321,8 +754,9 @@ function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, end end return with_context(to_proc) do - Dagger.pin_buffer!(:ROC, host_copy) - return ROCArray(host_copy) + arr = ROCArray(host_copy) + AMDGPU.synchronize() + return arr end end end @@ -332,14 +766,7 @@ function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, with_context(AMDGPU.synchronize, from_proc) return x elseif Dagger.root_worker_id(from_proc) == Dagger.root_worker_id(to_proc) - dev = AMDGPU.device(x) - with_context(AMDGPU.synchronize, dev.device_id) - return with_context(to_proc) do - to_arr = similar(x) - copyto!(to_arr, x) - AMDGPU.synchronize() - to_arr - end + return _peer_copy_out(to_proc, x) else host_copy = remotecall_fetch(from_proc.owner, from_proc, x) do from_proc, x return with_context(from_proc) do @@ -347,8 +774,9 @@ function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, end end return with_context(to_proc) do - Dagger.pin_buffer!(:ROC, host_copy) - return ROCArray(host_copy) + arr = ROCArray(host_copy) + AMDGPU.synchronize() + return arr end end end @@ -362,10 +790,22 @@ Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Chunk{T}) where function Dagger.execute!(proc::ROCArrayDeviceProc, f, args...; kwargs...) @nospecialize f args kwargs tls = Dagger.get_tls() + spec = tls.task_spec + writes = spec isa Dagger.Sch.TaskSpec && spec.options !== nothing ? + spec.options.writes : nothing task = Threads.@spawn begin Dagger.set_tls!(tls) with_context!(proc) - result = Base.@invokelatest f(args...; kwargs...) + copy_task = Dagger.is_move_task(f) ? + Threads.atomic_add!(COPY_TASK_COUNT, UInt64(1)) + UInt64(1) : UInt64(0) + result = if copy_task == 0 + Base.@invokelatest f(args...; kwargs...) + else + Dagger.with(COPY_TASK => copy_task) do + Base.@invokelatest f(args...; kwargs...) + end + end + _record_task!(proc.device_id, f, args, result, writes, copy_task) # N.B. Synchronization must be done when accessing result or args return result end @@ -553,6 +993,13 @@ function Dagger.ipc_materialize(info::HIPIpcInfo{T,N}) where {T,N} return Dagger.ipc_copyto!(dest, info) end +# A stream that does not synchronize with the null stream, on the current device +function _nonblocking_stream() + ref = Ref{AMDGPU.HIP.hipStream_t}() + AMDGPU.HIP.hipStreamCreateWithFlags(ref, AMDGPU.HIP.hipStreamNonBlocking) + return HIPStream(ref[]) +end + const DEVICES = Dict{Int, HIPDevice}() const CONTEXTS = Dict{Int, HIPContext}() const STREAMS = Dict{Int, HIPStream}() @@ -569,6 +1016,7 @@ function __init__() CONTEXTS[dev.device_id] = ctx context!(ctx) do STREAMS[dev.device_id] = HIPStream() + COPY_STREAMS[dev.device_id] = _nonblocking_stream() end return proc end diff --git a/src/array/darray.jl b/src/array/darray.jl index 8a80359ae..cf6d36593 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -300,14 +300,41 @@ function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} # # Sparse tiles (`DSparseArray`) are unwrapped to host storage before `cat`: # GPU sparse types (CuSparse/ROCSparse/DeviceSparseMatrixCSC) disallow the - # scalar indexing that generic `cat` would use. + # scalar indexing that generic `cat` would use. Dense GPU tiles are brought + # to host too: tiles may live on different devices (with no peer access), + # which a device-side `cat` cannot combine, and the result is a host + # `Array` regardless. dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:N] - tiles = asyncmap(a.chunks) do c - x = fetch(c) - x isa DSparseArray ? _sparse_collect(x.mat) : x - end + tiles = asyncmap(_collect_host_tile, a.chunks) return _collect_dense(T, Val(N), treereduce_nd(dimcatfuncs, tiles)) end + +# Host copy of a finished tile for `collect`'s in-process `cat`. GPU `execute!` +# does not synchronize, so a tile's kernels may still be running on Dagger's +# stream for that device. The caller's stream is a different one, and only +# some array libraries synchronize the previous owner when another stream +# touches a buffer (oneAPI does not). Wait on the tile's own processor and copy +# under its context instead, on the owner: the synchronization has to happen +# where the stream lives, and a host copy is also cheaper to send. +_collect_host_tile(t::DTask) = _collect_host_tile(fetch(t; raw=true)) +function _collect_host_tile(chunk::Chunk) + owner = root_worker_id(chunk) + owner == myid() && return _collect_host_tile_local(chunk) + return remotecall_fetch(_collect_host_tile_local, owner, chunk) +end +# Chunks may also hold plain values (e.g. `sort`'s result tiles) +_collect_host_tile(x) = _collect_host_tile_value(x) +function _collect_host_tile_local(chunk::Chunk) + proc = chunk.processor + x = fetch(chunk) + x isa DSparseArray || x isa GPUArraysCore.AbstractGPUArray || return x + gpu_synchronize(proc) + return with_context(() -> _collect_host_tile_value(x), proc) +end +_collect_host_tile_value(x) = + x isa DSparseArray ? _sparse_collect(x.mat) : + x isa GPUArraysCore.AbstractGPUArray ? Array(x) : x + Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A)) Base.wait(A::DArray) = foreach(wait, A.chunks) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index e020d864b..f8de25146 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1507,6 +1507,17 @@ function remotecall_endpoint_transfer(f, accel::DistributedAcceleration, from_pr end end +function remotecall_endpoint_transfer(f, accel::DistributedAcceleration, from_proc, to_proc, from_space, to_space, data::GPUArraysCore.AbstractGPUArray) + if root_worker_id(to_proc) == myid() + return f(accel, from_proc, to_proc, from_space, to_space, data) + end + # Keep the device allocation on its owner. Serializing a raw GPU array + # loses that ownership (and some backends cannot serialize it at all). + # A Chunk lets move dispatch to the backend's cross-worker transport. + chunk = tochunk(data, from_proc, from_space) + return remotecall_endpoint_transfer(f, accel, from_proc, to_proc, from_space, to_space, chunk) +end + #============================================================================== move_rewrap header + children protocol diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 91e400adf..e22866428 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -247,6 +247,23 @@ map_or_ntuple(f, xs::Vector) = map(f, 1:length(xs)) # N.B. Accept any `Tuple` (typed specs produce heterogeneous tuples of # `TypedArgument{T}`, not a homogeneous `NTuple{N,T}`). @inline map_or_ntuple(@specialize(f), xs::Tuple) = ntuple(f, Val(length(xs))) +# Indices, among the positional arguments `execute!` will receive (the function +# excluded), of the arguments a task writes; see the `writes` option. +function written_positions(task_arg_ws, deps_vec) + writes = nothing + npos = 0 + for idx in 2:length(task_arg_ws) + arg_ws = task_arg_ws[idx] + ispositional(arg_ws.pos) || continue + npos += 1 + if any(di->deps_vec[di].writedep, arg_deps_range(arg_ws)) + writes = something(writes, Int[]) + push!(writes, npos) + end + end + return something(writes, Int[]) +end + function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int; ownership=nothing) where typed @specialize spec fargs @@ -422,6 +439,7 @@ 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.writes = written_positions(task_arg_ws, deps_vec) if uniform new_spec.options.occupancy = Dict(Any=>0) end diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 28a9ed615..509da2f8e 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -542,11 +542,24 @@ function enqueue_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copied=true) end +""" + device_remainder_copy!(to_space, from_space, to, from, spans) -> Bool + +Copy remainder `spans` (`(source, dest)` pairs) between two device memory +spaces of this process without staging through the host, returning `true`, or +return `false` to let the caller take the host-staged path. GPU extensions +define methods for their own space and array types. +""" +device_remainder_copy!(to_space::MemorySpace, from_space::MemorySpace, to, from, spans) = false + # Main copy function for RemainderAliasing function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space::MemorySpace, to::Chunk, from::Chunk) where S - # 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) + # Same-device: copy spans directly with one KA launch. Not merely + # same-worker: one kernel reading `from` and writing `to` needs both on + # one device (two GPUs of a worker need not be peer-accessible), so other + # device pairs stage through the host below. + if from_space == to_space && root_worker_id(to_space) == myid() && + is_device_space(to_space) from_s = storage_array(unwrap(from)) to_s = storage_array(unwrap(to)) with_context!(from_space) @@ -554,6 +567,15 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: multi_span_copy!(to_s, from_s, dep_mod.spans) return end + # Two devices of this process: let the backend copy device-to-device + # (`device_remainder_copy!`) instead of round-tripping through a host buffer. + if from_space != to_space && root_worker_id(to_space) == myid() && + root_worker_id(from_space) == myid() && + is_device_space(to_space) && is_device_space(from_space) + from_s = storage_array(unwrap(from)) + to_s = storage_array(unwrap(to)) + device_remainder_copy!(to_space, from_space, to_s, from_s, dep_mod.spans) && return + end # N.B. Whole-object containers (e.g. `DSparseArray`) never reach here: their # storage reallocates on write, so they can't be span-copied. They are routed # to a whole-object `FullCopy` in `compute_remainder_for_arg!` instead. diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 63d18cf05..bfc5afb8f 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -180,11 +180,22 @@ function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::Chun to_raw = unwrap(to) from_w = root_worker_id(from_space) # TODO: Use dep_mod to fetch with less memory usage - from_raw = to_w == from_w ? unwrap(from) : remotecall_fetch(unwrap, from_w, from) + from_raw = to_w == from_w ? unwrap(from) : fetch_copy_source(from, to) move!(dep_mod, to_space, from_space, to_raw, from_raw) end return end + +# Device backends track copy tasks separately from user kernels so a completed +# transfer does not acquire dependencies on unrelated compute-stream work. +is_move_task(f) = f === move! + +fetch_copy_source(from::Chunk, to::Chunk) = + remotecall_fetch(unwrap, root_worker_id(from), from) +# GPU allocations must travel through the backend's transport, which stages +# or imports them on the receiving worker without serializing device handles. +fetch_copy_source(from::Chunk{<:GPUArraysCore.AbstractGPUArray}, to::Chunk) = + move(to.processor, from) function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::Base.RefValue{T}, from::Base.RefValue{T}) where {T} to[] = from[] return @@ -744,16 +755,30 @@ function _memory_spans(a::StridedAliasing{T,N,S}, spans, ptr, dim) where {T,N,S} return spans end +""" + data_address(x::DenseArray) -> UInt64 + +The address of `x`'s first element, for aliasing analysis. Defaults to +`pointer(x)`; GPU extensions override it where `pointer` has side effects +(CUDA.jl's and AMDGPU.jl's take stream ownership for the *active* stream, and +CUDA.jl's throws when that device cannot reach `x` without peer access), or +where it differs from the address their `aliasing` reports (Metal). +""" +data_address(x) = UInt64(pointer(x)) + function aliasing(x::SubArray{T,N}) where {T,N} if isbitstype(T) p = parent(x) space = memory_space(p) S = typeof(space) - parent_ptr = RemotePtr{Cvoid}(UInt64(pointer(p)), space) - ptr = RemotePtr{Cvoid}(UInt64(pointer(x)), space) + base = data_address(p) + parent_ptr = RemotePtr{Cvoid}(base, space) NA = ndims(p) raw_inds = parentindices(x) inds = ntuple(i->raw_inds[i] isa Integer ? (raw_inds[i]:raw_inds[i]) : UnitRange(raw_inds[i]), NA) + # The view's first element, located from the parent's address + first_offset = isempty(x) ? 0 : sum(ntuple(i->(first(inds[i]) - 1) * stride(p, i), NA)) + ptr = RemotePtr{Cvoid}(base + UInt64(first_offset * sizeof(T)), space) sz = ntuple(i->length(inds[i]), NA) return StridedAliasing{T,NA,S}(parent_ptr, ptr, diff --git a/src/options.jl b/src/options.jl index bf01364c3..672b97714 100644 --- a/src/options.jl +++ b/src/options.jl @@ -32,6 +32,7 @@ Stores per-task options to be passed to the scheduler. - `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. +- `writes::Union{Vector{Int},Nothing}=nothing`: (Datadeps, set automatically) Indices, among the positional arguments passed to `execute!`, of the arguments this task writes. Lets asynchronous device backends order their own transfers against only the producers of the data they move. `nothing` means unknown. - `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 @@ -73,6 +74,8 @@ Base.@kwdef mutable struct Options acceleration::Union{Acceleration,Nothing} = nothing + writes::Union{Vector{Int},Nothing} = nothing + return_type::Union{Type,Nothing} = nothing end Options(::Nothing) = Options() diff --git a/test/gpu.jl b/test/gpu.jl index 7ace19cca..3f2c42444 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -13,6 +13,7 @@ end # header for the acceleration/backend matrix they cover. include(joinpath(@__DIR__, "array", "sparse_defs.jl")) include(joinpath(@__DIR__, "array", "linalg", "matvec_defs.jl")) +include(joinpath(@__DIR__, "gpu_distributed.jl")) @testset "Dense matvec ($name)" for (name, scope) in GPU_SCOPES Dagger.with_options(; scope) do @@ -158,6 +159,71 @@ end end end + # Two GPUs of one process: copies between them take the peer-copy + # path on a separate copy stream, ordered only against the producers + # and prior users of the allocations involved. Each case below would + # read stale or overwritten data if one of those orderings were lost. + ndevices > 1 && @testset "Cross-GPU copies, one process ($(nameof(typeof(sched))))" for sched in + (Dagger.RoundRobinScheduler(), Dagger.GreedyScheduler()) + g1 = Dagger.scope(worker=1, cuda_gpu=1) + g2 = Dagger.scope(worker=1, cuda_gpu=2) + local_scope = Dagger.scope(worker=1, cuda_gpus=[1, 2]) + # A region plans over its ambient scope's processors + region(f) = Dagger.with_options(; scope=local_scope) do + Dagger.spawn_datadeps(f; scheduler=sched) + end + n = 1024 + + # Writes alternating between the GPUs: each pulls the other's write + A = zeros(Float32, n, n) + region() do + for i in 1:8 + Dagger.@spawn scope=(isodd(i) ? g1 : g2) addarray!(InOut(A)) + end + end + @test all(==(8f0), A) + + # A writer on one GPU racing readers of its previous value on the + # other: each reader must see exactly the preceding write, and the + # next write (and the next copy into the reader's replica) must + # wait for the read + A = zeros(Float32, n, n) + Bs = [zeros(Float32, n, n) for _ in 1:6] + region() do + for i in 1:6 + Dagger.@spawn scope=g1 addarray!(InOut(A)) + Dagger.@spawn scope=g2 copyto!(Out(Bs[i]), In(A)) + end + end + @test all(i->all(==(Float32(i)), Bs[i]), 1:6) + @test all(==(6f0), A) + + # Strided views: many short spans, packed on each side + A = zeros(Float32, n, n) + V = view(A, 1:(n ÷ 2), :) + region() do + for i in 1:4 + Dagger.@spawn scope=(isodd(i) ? g1 : g2) addarray!(InOut(V)) + end + end + @test all(==(4f0), A[1:(n ÷ 2), :]) + @test all(iszero, A[(n ÷ 2 + 1):end, :]) + + # Tiled GEMM across both GPUs, repeated on the same operands + DA = rand(Blocks(256, 256), Float32, 1024, 1024) + DB = rand(Blocks(256, 256), Float32, 1024, 1024) + DC = zeros(Blocks(256, 256), Float32, 1024, 1024) + ref = collect(DA) * collect(DB) + Dagger.with(Dagger.DATADEPS_SCHEDULER => sched) do + Dagger.with_options(; scope=local_scope) do + for _ in 1:3 + mul!(DC, DA, DB) + @test collect(DC) ≈ ref + end + end + end + end + @testset "Datadeps (GPU $gpu)" for gpu in gpu_configs local_scope = Dagger.scope(worker=1, cuda_gpus=(gpu == :all ? Colon() : [gpu])) @@ -342,6 +408,65 @@ end end end + # Two GPUs of one process: see the CUDA testset of the same name + ndevices > 1 && @testset "Cross-GPU copies, one process ($(nameof(typeof(sched))))" for sched in + (Dagger.RoundRobinScheduler(), Dagger.GreedyScheduler()) + g1 = Dagger.scope(worker=1, rocm_gpu=1) + g2 = Dagger.scope(worker=1, rocm_gpu=2) + local_scope = Dagger.scope(worker=1, rocm_gpus=[1, 2]) + region(f) = Dagger.with_options(; scope=local_scope) do + Dagger.spawn_datadeps(f; scheduler=sched) + end + n = 1024 + + # Writes alternating between the GPUs: each pulls the other's write + A = zeros(Float32, n, n) + region() do + for i in 1:8 + Dagger.@spawn scope=(isodd(i) ? g1 : g2) addarray!(InOut(A)) + end + end + @test all(==(8f0), A) + + # A writer on one GPU racing readers of its previous value on the + # other + A = zeros(Float32, n, n) + Bs = [zeros(Float32, n, n) for _ in 1:6] + region() do + for i in 1:6 + Dagger.@spawn scope=g1 addarray!(InOut(A)) + Dagger.@spawn scope=g2 copyto!(Out(Bs[i]), In(A)) + end + end + @test all(i->all(==(Float32(i)), Bs[i]), 1:6) + @test all(==(6f0), A) + + # Strided views: many short spans, packed on each side + A = zeros(Float32, n, n) + V = view(A, 1:(n ÷ 2), :) + region() do + for i in 1:4 + Dagger.@spawn scope=(isodd(i) ? g1 : g2) addarray!(InOut(V)) + end + end + @test all(==(4f0), A[1:(n ÷ 2), :]) + @test all(iszero, A[(n ÷ 2 + 1):end, :]) + + # Tiled GEMM across both GPUs, repeated on the same operands + DA = rand(Blocks(256, 256), Float32, 1024, 1024) + DB = rand(Blocks(256, 256), Float32, 1024, 1024) + DC = zeros(Blocks(256, 256), Float32, 1024, 1024) + ref = collect(DA) * collect(DB) + Dagger.with(Dagger.DATADEPS_SCHEDULER => sched) do + Dagger.with_options(; scope=local_scope) do + for _ in 1:3 + mul!(DC, DA, DB) + @test collect(DC) ≈ ref + end + end + end + end + @testset "Datadeps (GPU $gpu)" for gpu in gpu_configs local_scope = Dagger.scope(worker=1, rocm_gpus=(gpu == :all ? Colon() : [gpu])) diff --git a/test/gpu_distributed.jl b/test/gpu_distributed.jl new file mode 100644 index 000000000..f96742b0b --- /dev/null +++ b/test/gpu_distributed.jl @@ -0,0 +1,189 @@ +# Included by gpu.jl in every GPU CI job. Both workers select device 1, so +# this exercises inter-process transfers even on agents with just one GPU. +using Distributed, Dagger, LinearAlgebra, Test + +@everywhere begin + function distributed_gpu_add!(A) + A .+= 1f0 + return nothing + end + function distributed_gpu_scale!(A) + A .*= 2f0 + return nothing + end + # Slow on purpose, and exact in Float32: a read that doesn't wait for the + # kernel shows up as a wrong value rather than a rounding difference. + function distributed_gpu_slowinc(x) + for _ in 1:20000 + x += 1f0 + end + return x - 19999f0 + end +end + +function test_distributed_gpu(kind, gpu_key) + ws = filter(!=(myid()), workers()) + @test length(ws) >= 2 + length(ws) >= 2 || return + w1, w2 = ws[1:2] + @test Dagger.system_uuid(w1) == Dagger.system_uuid(w2) + gpu_kw = NamedTuple{(gpu_key,)}((1,)) + s1 = Dagger.scope(; worker=w1, gpu_kw...) + s2 = Dagger.scope(; worker=w2, gpu_kw...) + p1 = only(Dagger.compatible_processors(s1)) + p2 = only(Dagger.compatible_processors(s2)) + @test p1 isa Dagger.gpu_processor(kind) + @test p2 isa Dagger.gpu_processor(kind) + scope = Dagger.UnionScope(s1, s2) + + @testset "GPU chunk transfers" begin + # Include a size above the IPC threshold as well as small arrays. + # Use nonuniform data to detect lost, stale, or incorrectly offset copies. + for dims in ((8, 8), (256, 256)) + A = reshape(Float32.(1:prod(dims)), dims) + source_task = Dagger.@spawn scope=s1 identity(A) + source = fetch(source_task; raw=true) + @test Dagger.processor(source) == p1 + @test Dagger.chunktype(source) <: Dagger.GPUArraysCore.AbstractGPUArray + + # Repeat the import of the same source, and exercise both directions. + # Return a Chunk so the device array never crosses serialization. + for _ in 1:2 + received = remotecall_fetch(w2, p1, p2, source) do from, to, chunk + Dagger.with_context(to) do + value = Dagger.move(from, to, chunk) + Dagger.tochunk(value, to) + end + end + @test Dagger.processor(received) == p2 + @test Dagger.chunktype(received) <: Dagger.GPUArraysCore.AbstractGPUArray + @test collect(received) == A + remotecall_fetch(w2, p2, received) do proc, chunk + Dagger.with_context(proc) do + distributed_gpu_add!(Dagger.unwrap(chunk)) + Dagger.gpu_synchronize(proc) + end + nothing + end + @test collect(source) == A + @test collect(received) == A .+ 1f0 + returned = Dagger.@spawn scope=s1 identity(received) + @test fetch(returned) == A .+ 1f0 + @test Dagger.processor(fetch(returned; raw=true)) == p1 + end + end + end + + @testset "Cross-worker datadeps" begin + A = reshape(Float32.(1:64), 8, 8) + ref = copy(A) + Dagger.with_options(; scope) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=s1 distributed_gpu_add!(Dagger.InOut(A)) + Dagger.@spawn scope=s2 distributed_gpu_scale!(Dagger.InOut(A)) + Dagger.@spawn scope=s1 distributed_gpu_add!(Dagger.InOut(A)) + end + end + @test A == (ref .+ 1f0) .* 2f0 .+ 1f0 + + end + + @testset "collect waits for queued kernels" begin + # GPU tasks return with their kernels still queued on Dagger's stream; + # `collect` must wait on that stream, not the caller's. Covers tiles + # owned by the driver and by remote workers. + A = reshape(Float32.(1:4096), 64, 64) + local_gpu = any(p -> p isa Dagger.gpu_processor(kind), + Dagger.get_processors(Dagger.OSProc(myid()))) + scopes = local_gpu ? (Dagger.scope(; worker=myid(), gpu_kw...), s1, s2) : (s1, s2) + for s in scopes + for _ in 1:3 + B = Dagger.with_options(; scope=s) do + DB = map(distributed_gpu_slowinc, Dagger.distribute(A, Dagger.Blocks(32, 32))) + wait(DB) + collect(DB) + end + @test B == A .+ 1f0 + end + end + end + + @testset "Distributed GPU arrays" begin + A = reshape(Float32.(1:256) ./ 256f0, 16, 16) + B = reverse(A; dims=1) + Dagger.with_options(; scope) do + # Opposite placements guarantee that corresponding tiles must move. + DA = Dagger.distribute(A, Dagger.Blocks(8, 16), + reshape(Dagger.Processor[p1, p2], 2, 1)) + DB = Dagger.distribute(B, Dagger.Blocks(8, 16), + reshape(Dagger.Processor[p2, p1], 2, 1)) + wait(DA) + wait(DB) + @test collect(DA .+ DB) ≈ A .+ B + + # The original reproducer uses row tiles times column tiles. + # Metal/OpenCL currently lack the required dense matmul methods; + # their transfer, broadcast, and datadeps coverage above is shared. + if kind in (:CUDA, :ROC, :oneAPI) + DB_cols = Dagger.distribute(B, Dagger.Blocks(16, 8), + reshape(Dagger.Processor[p2, p1], 1, 2)) + wait(DB_cols) + @test collect(DA * DB_cols) ≈ A * B + end + end + end +end + +@testset "Distributed workers sharing one GPU ($kind)" for (kind, key) in + ((:CUDA, :cuda_gpu), (:ROC, :rocm_gpu), (:oneAPI, :intel_gpu), + (:Metal, :metal_gpu), (:OpenCL, :cl_device)) + if Dagger.gpu_can_compute(kind) + test_distributed_gpu(kind, key) + end +end + +# Two GPUs of one process need not be peer-accessible, so aliasing, halo +# copies, and `collect` must never touch one device's memory from the other's +# context. Needs a second device; single-GPU agents skip it. +@everywhere import Dagger: @stencil, Wrap + +function test_multi_gpu_process(kind) + ps = filter(p -> p isa Dagger.gpu_processor(kind), + collect(Dagger.get_processors(Dagger.OSProc(myid())))) + length(ps) >= 2 || return + p1, p2 = sort!(ps; by=string)[1:2] + procs = reshape(Dagger.Processor[p1, p2], 1, 2) + scope = Dagger.UnionScope(Dagger.ExactScope(p1), Dagger.ExactScope(p2)) + + @testset "Multiple GPUs in one process ($kind)" begin + A = reshape(Float32.(1:64), 8, 8) + ref = copy(A) + Dagger.with_options(; scope) do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=Dagger.ExactScope(p1) distributed_gpu_add!(Dagger.InOut(A)) + Dagger.@spawn scope=Dagger.ExactScope(p2) distributed_gpu_scale!(Dagger.InOut(A)) + Dagger.@spawn scope=Dagger.ExactScope(p1) distributed_gpu_add!(Dagger.InOut(A)) + end + end + @test A == (ref .+ 1f0) .* 2f0 .+ 1f0 + + # Column tiles on different devices: the halo exchange copies spans + # between devices, and `collect` gathers tiles from both. + n = size(A, 1) + box(i, j) = sum(A[mod1(i + di, n), mod1(j + dj, n)] for di in -1:1, dj in -1:1) + expected = [box(i, j) for i in 1:n, j in 1:n] + B = Dagger.with_options(; scope) do + DA = Dagger.distribute(A, Dagger.Blocks(8, 4), procs) + DB = Dagger.distribute(zeros(Float32, 8, 8), Dagger.Blocks(8, 4), procs) + @stencil DB[idx] = sum(@neighbors(DA[idx], 1, Wrap())) + collect(DB) + end + @test B ≈ expected + end +end + +for kind in (:CUDA, :ROC, :oneAPI, :Metal, :OpenCL) + if Dagger.gpu_can_compute(kind) + test_multi_gpu_process(kind) + end +end diff --git a/test/mpi/oneapi.jl b/test/mpi/oneapi.jl index 7321d4c02..78d446d2e 100644 --- a/test/mpi/oneapi.jl +++ b/test/mpi/oneapi.jl @@ -30,5 +30,6 @@ run_mpi_gpu_suite((; make_space = () -> IntelExt.IntelVRAMMemorySpace(1, 1), device_field = :device_id, kind = :oneAPI, + test_ipc = true, ), )) diff --git a/test/mpi/rocm.jl b/test/mpi/rocm.jl index 2376475e0..04b726163 100644 --- a/test/mpi/rocm.jl +++ b/test/mpi/rocm.jl @@ -4,6 +4,7 @@ # Run (env must provide Dagger, MPI, AMDGPU; see test/rocmenv): # mpiexec -n 2 julia --project=test/rocmenv --threads=2 test/mpi/rocm.jl +using AMDGPU_LLVM_Backend_jll using Dagger, MPI, AMDGPU, LinearAlgebra, Random, Test using Dagger: In, Out, InOut, Deps diff --git a/test/rocmenv/Project.toml b/test/rocmenv/Project.toml index 7b8ac67f3..a30850a60 100644 --- a/test/rocmenv/Project.toml +++ b/test/rocmenv/Project.toml @@ -1,5 +1,6 @@ [deps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +AMDGPU_LLVM_Backend_jll = "cc5c0156-bd05-5a77-8a68-bb0aafb29019" Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/test/runtests.jl b/test/runtests.jl index ce7c874d6..19fb41804 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -88,17 +88,26 @@ optin_test_names = USE_FINCH ? String[] : String[ ] additional_workers::Int = 3 -worker_threads::Int = 1 +# AMDGPU's persistent host-call tasks can monopolize a worker's only thread. +worker_threads::Int = USE_ROCM ? 2 : 1 -if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ - pushfirst!(LOAD_PATH, @__DIR__) - pushfirst!(LOAD_PATH, joinpath(@__DIR__, "..")) +direct_test_run = PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ +if direct_test_run using Pkg Pkg.activate(@__DIR__) + # Resolve direct invocations in one project, as Pkg.test does. + # A repository project ahead of it in LOAD_PATH can shadow backend versions. + Pkg.develop([PackageSpec(path=joinpath(@__DIR__, "..")), + PackageSpec(path=joinpath(@__DIR__, "..", "lib", "TimespanLogging"))]) try Pkg.instantiate() catch end +end + +# Pkg.test supplies test_args through ARGS without setting PROGRAM_FILE. +if direct_test_run || !isempty(ARGS) + using Pkg using ArgParse s = ArgParseSettings(description = "Dagger Testsuite") @eval begin @@ -122,7 +131,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ help = "How many additional workers to launch" "-t", "--threads" arg_type = Int - default = 1 + default = worker_threads help = "How many threads to give each additional worker" "-v", "--verbose" action = :store_true @@ -196,6 +205,8 @@ else to_test = filter(!in(optin_test_names), all_test_names) @info "Running all tests" end +flush(stdout) +flush(stderr) # Resolve optional dependencies before loading Dagger (and its dependencies) or # starting workers. Adding a GPU backend can downgrade LLVM/Atomix; a process @@ -204,7 +215,7 @@ if USE_GPU || USE_FINCH || "extlang/python" in to_test using Pkg optional_packages = String[] USE_CUDA && push!(optional_packages, "CUDA") - USE_ROCM && push!(optional_packages, "AMDGPU") + USE_ROCM && append!(optional_packages, ["AMDGPU", "AMDGPU_LLVM_Backend_jll"]) USE_ONEAPI && push!(optional_packages, "oneAPI") USE_METAL && push!(optional_packages, "Metal") USE_OPENCL && append!(optional_packages, ["OpenCL", "pocl_jll"]) @@ -246,7 +257,11 @@ try test_name = all_test_names[findfirst(x->x==test, all_test_names)] println() @info "Testing $test_title ($test_name)" + flush(stdout) + flush(stderr) @testset "$test_title" include(test * ".jl") + flush(stdout) + flush(stderr) end catch printstyled(stderr, "Tests Failed!\n"; color=:red) diff --git a/test/setup_gpu.jl b/test/setup_gpu.jl index eceef4ff5..4c016fc81 100644 --- a/test/setup_gpu.jl +++ b/test/setup_gpu.jl @@ -8,6 +8,8 @@ end if $USE_ROCM + # Use modern AMD codegen even when the environment resolves AMDGPU 2.1. + using AMDGPU_LLVM_Backend_jll using AMDGPU elseif !$IN_CI try using AMDGPU