MPI: keep tasks with acceleration=DistributedAcceleration() rank-local

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

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))
 

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

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

test/mpi/cpu.jl

@@ -146,6 +146,7 @@ end
 end
 
 include("spawn.jl")
+include("rank_local.jl")
 
 @testset "Datadeps" begin
 

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