Repository navigation
MPI: reduction over task results limits datadeps concurrency to 2 tasks in total (hierarchical planner) #759
Description
Follow-up to #757 – thanks for the quick fix! With 0.22.6, independent tasks and the first example from the MPI docs now use all threads.
We still see low concurrency when tasks take other tasks' results as arguments, e.g. a reduction tree over partial results:
using Dagger, MPI
Dagger.accelerate!(:mpi)
work(n) = (s = 0.0; for i in 1:n; x = (i - 0.5) / n; s += 4.0 / (1.0 + x * x); end; s / n)
Dagger.spawn_datadeps() do
p = [Dagger.@spawn work(1 << 27) for _ in 1:16]
m = length(p)
while m > 1
for k in 1:m ÷ 2
p[k] = Dagger.@spawn +(p[2k - 1], p[2k])
end
m ÷= 2
end
p[1]
endOnly two work tasks run at the same time, also in total across ranks:
| 1 rank × 4 threads | 2 ranks × 2 threads | |
|---|---|---|
| independent tasks only | 0.30 s | 0.30 s |
| reduction tree | 0.60 s | 0.60 s |
reduction tree, hierarchical=false |
0.30 s | 0.31 s |
Without accelerate!(:mpi), the tree uses all four threads. Tested with Dagger 0.22.6, MPI.jl 0.20.27 (MPICH_jll), Julia 1.13.0 on macOS.
We are not sure about the cause, but it might be related to the hierarchical planner under MPI: it seems to plan in DFS order, where each + follows its two inputs, and planning a task with a DTask value argument may wait for that value (fetch in populate_task_info!). When we pass the partial results through In/InOut buffers instead of task values, all four threads are used.
Thanks a lot for looking into this, and for all your work on Dagger!
Full script with timing
using Dagger, MPI
const USE_MPI = get(ENV, "NO_MPI", "0") != "1"
USE_MPI ? Dagger.accelerate!(:mpi) : MPI.Init()
const TREE = get(ENV, "TREE", "1") == "1"
const HIERARCHICAL = get(ENV, "HIERARCHICAL", "1") == "1"
const rank = MPI.Comm_rank(MPI.COMM_WORLD)
const LOG = Tuple{Float64,Float64}[] # (start, end) of the work tasks on this rank
const LOG_LOCK = ReentrantLock()
function work(n)
t1 = time()
s = 0.0
for i in 1:n
x = (i - 0.5) / n
s += 4.0 / (1.0 + x * x)
end
lock(() -> push!(LOG, (t1, time())), LOG_LOCK)
return s / n
end
# maximum number of work tasks running at the same time
function max_concurrent(iv)
ev = sort!(vcat([(s, 1) for (s, _) in iv], [(e, -1) for (_, e) in iv]))
c = m = 0
for (_, d) in ev
c += d; m = max(m, c)
end
return m
end
for rep in 1:2
empty!(LOG)
t = @elapsed Dagger.spawn_datadeps(; hierarchical=HIERARCHICAL) do
p = [Dagger.@spawn work(1 << 27) for _ in 1:16]
if TREE
m = length(p)
while m > 1
for k in 1:m ÷ 2
p[k] = Dagger.@spawn +(p[2k - 1], p[2k])
end
m ÷= 2
end
end
p[1]
end
if rep == 2
USE_MPI && MPI.Barrier(MPI.COMM_WORLD)
print("rank $rank: $(round(t; digits=3)) s, $(length(LOG)) work tasks on this rank, " *
"max $(max_concurrent(LOG)) at a time\n")
end
endmpiexec -n 1 julia --project=. --threads=4 script.jl # 0.60 s, max 2 at a time
HIERARCHICAL=0 mpiexec -n 1 julia --project=. --threads=4 script.jl # 0.30 s, max 4 at a time
Activity
jpsamaroo commentedon Oct 7, 2026
There are problems with dealing with task results as arguments within a
spawn_datadepsregion (when producing and consuming task come from the same region), that forces us to serialize execution (specifically: we need to know the aliasing info of the task result before we can schedule and plan the task that consumes it), so I would try to avoid writing an algorithm like this with Datadeps for now. As you say, passing partial results through regular buffers (or allocating the buffers as Dagger tasks outside ofspawn_datadeps, or in a prior region) is also a good strategy that avoids this problem.AnnaLena77 commentedon Oct 8, 2026
Thanks a lot for the quick answer and the explanation, that makes sense!
One thing we noticed, though: with spawn_datadeps(; hierarchical=false), the same tree runs all 16 work tasks concurrently (0.30 s instead of 0.60 s), even though the flat planner also has to wait for the task results before planning the + tasks. Our guess is that the flat planner goes in submission order, so all work tasks are already submitted before the first wait. The hierarchical planner under MPI seems to go in DFS order (work, work, +, work, work, +, …), so it waits after every two tasks.
So maybe the serialization could be limited to the consumers, e.g. by planning in submission order, or by submitting all independent tasks before the first wait?
We'll use buffers for now. Since plain @Spawn outside datadeps currently also runs serially under MPI (#760 ), buffers seem to be the only option for this pattern at the moment.
Thanks again!
jpsamaroo commentedon Oct 13, 2026
That's a great point, and I'm working on that in #773 which optimizes the traversal order while maintaining correctness in the face of unknown aliasing of task results. It's showing some small improvements in benchmark CI but greater improvements locally (CI is quite noisy and not representative anyway), so I'll probably merge it soon.