Repository navigation
MPI: keep tasks with acceleration=DistributedAcceleration() rank-local - #779
Open
AnnaLena77 wants to merge 1 commit into
Open
AnnaLena77 wants to merge 1 commit into
AnnaLena77 wants to merge 1 commit into
Conversation
jpsamaroo
reviewed
Oct 15, 2026
jpsamaroo
left a comment
Member
There was a problem hiding this comment.
Thanks for this! I'm not approving or rejecting this PR - I need some time to figure out the design, and then we can come back to this.
| return UInt64(Threads.atomic_add!(LOCAL_ID_COUNTER, 1)) | ||
| end | ||
| return eager_next_id() | ||
| end |
Member
There was a problem hiding this comment.
This makes me quite uncomfortable - I want to think more on the design before committing to something like this.
| # 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) |
Member
There was a problem hiding this comment.
I think this will cause an increase in allocations, as it was a fast-path avoiding Set allocation.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #778.
Under
accelerate!(:mpi), tasks spawned withacceleration=DistributedAcceleration()(or insidewith_options(; acceleration=DistributedAcceleration())) are meant to run only on the rank that spawns them. Two things broke that when ranks spawn such tasks independently:accel_matches_proc(::DistributedAcceleration, proc)istruefor every processor, andcompatible_processors(::DistributedAcceleration, ...)skipped the filter completely. Some of these tasks were placed onMPIProcessors, sometimes on another rank, and then never ran. This PR excludes MPI processors forDistributedAccelerationinMPIExtand removes the shortcut, so the generic filter applies.LOCAL_ID_COUNTER, starting attypemax(Int) ÷ 2).Tests:
test/mpi/rank_local.jl(included intest/mpi/cpu.jl) spawns a different number of local tasks on every rank, both ways, checks that they run onThreadProcs of their own rank, then runs uniform work with cross-rank data, and checks that a local task can use a uniform result owned by its own rank. Without the fix it fails (tasks on MPI processors,Non-uniform value, hangs); with it, it passes. The fulltest/mpi/cpu.jlpasses on 2 ranks × 2 threads (MPICH_jll, macOS).Not covered: a local task using a uniform result owned by another rank fails with an error, and passing a local task's result back into uniform work isn't addressed here.
Written with help from Claude (Anthropic); I checked the results.