diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 05d4f5eac..e722a5a30 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -2661,5 +2661,9 @@ accel_matches_proc(accel::MPIAcceleration, proc::MPIOSProc) = true accel_matches_proc(accel::MPIAcceleration, proc::MPIClusterProc) = true accel_matches_proc(accel::MPIAcceleration, proc::MPIProcessor) = true accel_matches_proc(accel::MPIAcceleration, proc) = false +# A task with `acceleration=DistributedAcceleration()` runs only on this rank, outside +# SPMD execution: it must not be placed on (or stolen by) an MPI processor, whose +# execution protocol expects every rank to run the same task. +Dagger.accel_matches_proc(::Dagger.DistributedAcceleration, ::Union{MPIProcessor,MPIOSProc,MPIClusterProc}) = false end # module MPIExt diff --git a/src/acceleration.jl b/src/acceleration.jl index 409bd2c5a..e25022818 100644 --- a/src/acceleration.jl +++ b/src/acceleration.jl @@ -65,10 +65,6 @@ function compatible_processors(accel::Union{Acceleration,Nothing}, scope::Abstra # function's contract is `-> Set{Processor}`. return Set{Processor}(p for p in comp if accel_matches_proc(accel, p)) end -# accel_matches_proc is identically true for DistributedAcceleration; skip the -# second Set construction -compatible_processors(accel::DistributedAcceleration, scope::AbstractScope, procs::Vector{<:Processor}) = - compatible_processors(scope, procs) compatible_processors(accel::Union{Acceleration,Nothing}, scope::AbstractScope=get_compute_scope(), ctx::Context=Sch.eager_context()) = compatible_processors(accel, scope, procs(ctx)) diff --git a/src/dtask.jl b/src/dtask.jl index 29fecce51..929c41e09 100644 --- a/src/dtask.jl +++ b/src/dtask.jl @@ -152,3 +152,17 @@ function eager_next_id() return remotecall_fetch(eager_next_id, 1)::UInt64 end end + +# Under uniform (SPMD) execution, task IDs feed rank-uniform MPI tags, so every rank +# must allocate the same IDs in the same order. A task that opts out of uniform +# execution (e.g. `acceleration=DistributedAcceleration()` inside `accelerate!(:mpi)`) +# is spawned by one rank only; it draws from a separate, disjoint range so that it +# does not shift the IDs of later uniform tasks. +const LOCAL_ID_COUNTER = Threads.Atomic{Int}(typemax(Int) รท 2) +function eager_next_id(spec) + accel = spec.options.acceleration + if accel !== nothing && uniform_execution(current_acceleration()) && !uniform_execution(accel) + return UInt64(Threads.atomic_add!(LOCAL_ID_COUNTER, 1)) + end + return eager_next_id() +end diff --git a/src/submission.jl b/src/submission.jl index 2f8102a4d..2bd83835f 100644 --- a/src/submission.jl +++ b/src/submission.jl @@ -526,7 +526,7 @@ arg_chunktypes(fargs::Vector) = function eager_spawn(spec::DTaskSpec) # Generate new unlaunched DTask - uid = eager_next_id() + uid = eager_next_id(spec) future = ThunkFuture() metadata = DTaskMetadata(spec) # Propagate inferred return type to options diff --git a/test/mpi/cpu.jl b/test/mpi/cpu.jl index 65c2697ea..dc7cc6358 100644 --- a/test/mpi/cpu.jl +++ b/test/mpi/cpu.jl @@ -146,6 +146,7 @@ end end include("spawn.jl") +include("rank_local.jl") @testset "Datadeps" begin diff --git a/test/mpi/rank_local.jl b/test/mpi/rank_local.jl new file mode 100644 index 000000000..f5d1e1bb7 --- /dev/null +++ b/test/mpi/rank_local.jl @@ -0,0 +1,33 @@ +# Tasks spawned with `acceleration=DistributedAcceleration()` run only on the +# rank that spawns them, outside SPMD execution. Ranks may spawn different +# numbers of them; they must neither land on MPI processors nor shift the task +# IDs (and so the MPI tags) of the uniform tasks that follow. +mpi_local_location(i) = (MPI.Comm_rank(MPI.COMM_WORLD), i) + +@testset "Rank-local tasks" begin + local_accel = Dagger.DistributedAcceleration() + for round in 1:2 + # A different number of local tasks on every rank + n = 2 + 3rank + tasks = [Dagger.@spawn(acceleration=local_accel, mpi_local_location(i)) for i in 1:n] + for (i, task) in enumerate(tasks) + @test fetch(task) == (rank, i) + proc = fetch(task; raw=true).processor + @test proc isa Dagger.ThreadProc + end + tasks = Dagger.with_options(; acceleration=local_accel) do + [Dagger.@spawn(mpi_local_location(i)) for i in 1:n] + end + @test map(fetch, tasks) == [(rank, i) for i in 1:n] + + # Uniform work afterwards still agrees across ranks + x = Dagger.@spawn scope=Dagger.scope(mpi_rank=nranks-1) fill(Float64(round), 100) + @test fetch(Dagger.@spawn scope=Dagger.scope(mpi_rank=0) sum(x)) == 100round + @test sum(fetch.([Dagger.@spawn(mpi_local_location(i)) for i in 1:4]) .|> last) == 10 + + # A local task can use a uniform result owned by its own rank + mine = [Dagger.@spawn(scope=Dagger.scope(mpi_rank=r), fill(Float64(r + 1), 10)) for r in 0:nranks-1] + @test fetch(Dagger.@spawn(acceleration=local_accel, sum(mine[rank+1]))) == 10(rank + 1) + MPI.Barrier(comm) + end +end