Skip to content

Convert 'anyv' parallelisation to 'anysv' #361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
johnomotani merged 2 commits into from
Jul 13, 2025
Merged

Convert 'anyv' parallelisation to 'anysv' #361

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
928353c
Convert 'anyv' parallelisation to 'anysv'
johnomotani May 2, 2025
0c07882
Convert Fokker-Planck collision operator to 'anysv' region
johnomotani May 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 29 additions & 28 deletions docs/src/developing.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,43 +195,44 @@ Internally, when the `@begin_*_region()` macros need to change the region type (

If for some reason it is necessary to synchronize explicitly, not by using an `@begin_*_region()` call, use the `@_block_synchronize()` macro. This calls `_block_synchronize()`, but when debugging can also pass in some information (a hash of the file an line number) about the calling site that is used for consistency checking.

### Collision operator and `anyv` region
### Collision operator and `anysv` region

The Fokker-Planck collision operator requires a special approach to
shared-memory parallelisation. There is an outer loop over spatial points (and
potentially over species). Inside that outer loop there are operations that can
benefit from parallelisation over $v_{\perp}$, or over $v_{\parallel}$, or over
both $v_{\perp}$ and $v_{\parallel}$, as well as some that do not parallelise
over velocity space at all. To deal with this, it is beneficial to parallelise
the outer loop over species and spatial dimensions as much as possible, and
then within that allow changes between different ways of parallelizing over
velocity space.
shared-memory parallelisation. There is an outer loop over spatial points.
Inside that outer loop there are operations that can benefit from
parallelisation over (possibly species and) $v_{\perp}$, or over
$v_{\parallel}$, or over both $v_{\perp}$ and $v_{\parallel}$, as well as some
that do not parallelise over velocity space at all. To deal with this, it is
beneficial to parallelise the outer loop over species and spatial dimensions as
much as possible, and then within that allow changes between different ways of
parallelizing over velocity space.

The mechanism introduced to allow the type of parallelization just described is
the 'anyv' (read any-$v$) region. Before the outer loop of the collision
operator `@begin_s_r_z_anyv_region()` is used to start the 'anyv'
parallelization. Then within the `@loop is ir iz begin...` the functions
`@begin_anyv_region()` (for no parallelization over velocity space),
`@begin_anyv_vperp_region()`, `@begin_anyv_vpa_region()` and
`@begin_anyv_vperp_vpa_region()` can be used to parallelize over neither
velocity space dimension, either velocity space dimension individually, or over
both velocity space dimensions together. This is possible because 'subblocks'
of processes are defined. Each subblock shares the same range of species and
spatial indices, which stay the same throughout the `@begin_s_r_z_anyv_region()`
section, and are not shared with any other subblock of processes. Because the
subblock has an independent set of species- and spatial-indices, when changing
the velocity-space parallelization only the processes in the sub-block need to
be synchronized which is done by
[`moment_kinetics.communication._anyv_subblock_synchronize`](@ref), which is
called when necessary within the `@begin_anyv*_region()` functions (the whole
the 'anysv' (read any-$s$-$v$) region. Before the outer loop of the collision
operator `@begin_r_z_anysv_region()` is used to start the 'anysv'
parallelization. Then within the `@loop_r_z ir iz begin...` the functions
`@begin_anysv_region()` (for no parallelization over velocity space),
`@begin_anysv_s_vperp_region()`, `@begin_anysv_s_vpa_region()`,
`@begin_anysv_s_vperp_vpa_region()`, `@begin_anysv_vperp_region()`,
`@begin_anysv_vpa_region()` and `@begin_anysv_vperp_vpa_region()` can be used
to parallelize over species and neither velocity space dimension, either
velocity space dimension individually, or over both velocity space dimensions
together. This is possible because 'subblocks' of processes are defined. Each
subblock shares the same range of spatial indices, which stay the same
throughout the `@begin_r_z_anysv_region()` section, and are not shared with any
other subblock of processes. Because the subblock has an independent set of
spatial-indices, when changing the species- and velocity-space parallelization
only the processes in the sub-block need to be synchronized which is done by
[`moment_kinetics.communication._anysv_subblock_synchronize`](@ref), which is
called when necessary within the `@begin_anysv*_region()` functions (the whole
shared-memory block does not need to be synchronized at once, as would be done
by [`moment_kinetics.communication._block_synchronize`](@ref)). The processes
that share an anyv subblock are all part of the `comm_anyv_subblock[]`
that share an anysv subblock are all part of the `comm_anysv_subblock[]`
communicator (which is a subset of the processes in the full block, whose
communicator is `comm_block[]`).

See also notes on debugging the 'anyv' parallelisation: [Collision operator and
'anyv' region](@ref).
See also notes on debugging the 'anysv' parallelisation: [Collision operator
and 'anysv' region](@ref).

## Bounds checking

Expand Down
22 changes: 11 additions & 11 deletions moment_kinetics/debug_test/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,20 +23,20 @@ something like
julia --project --check-bounds=yes --compiled-modules=no debug_test/runtests.jl --debug 99
```

Collision operator and 'anyv' region
Collision operator and 'anysv' region
------------------------------------

The collision operator uses a slightly hacky special set of functions for
shared memory parallelism, to allow the outer loop over species and spatial
dimensions to be parallelised, but also inner loops over `vperp`, `vpa` or
`vperp` and `vpa` to be parallelised - changing the type of inner-loop
parallelism within the outer loop. This happens within an 'anyv' region, which
is started with the `begin_s_r_z_anyv_region()` function. The debug checks
within an 'anyv' region only check for correctness on the sub-block
communicator that parallelises over velocity space, so if there were errors due
to incorrect species or spatial parallelism they would not (might not?) be
detected. These errors should be unlikely as the collision operator only writes
to a single species at a single spatial point.
shared memory parallelism, to allow the outer loop over spatial dimensions to
be parallelised, but also inner loops over `s`, `vperp`, `vpa` or combinations
to be parallelised - changing the type of inner-loop parallelism within the
outer loop. This happens within an 'anysv' region, which is started with the
`begin_r_z_anysv_region()` function. The debug checks within an 'anysv' region
only check for correctness on the sub-block communicator that parallelises over
velocity space, so if there were errors due to incorrect species or spatial
parallelism they would not (might not?) be detected. These errors should be
unlikely as the collision operator only writes to a single species at a single
spatial point.

Finding race conditions
-----------------------
Expand Down
6 changes: 3 additions & 3 deletions moment_kinetics/debug_test/runtest_template.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ using moment_kinetics: setup_moment_kinetics, cleanup_moment_kinetics!
using moment_kinetics.time_advance: time_advance!
using moment_kinetics.communication
using moment_kinetics.looping: all_dimensions, dimension_combinations,
anyv_dimension_combinations
anysv_dimension_combinations
using moment_kinetics.type_definitions: OptionsDict
using moment_kinetics.Glob
using moment_kinetics.Primes
Expand Down Expand Up @@ -52,7 +52,7 @@ function runtests(; restart=false)
# Only need to test dimension combinations that are actually used for parallel loops
# in some part of the code
dimension_combinations_to_test = [c for c in tuple(dimension_combinations...,
anyv_dimension_combinations...)
anysv_dimension_combinations...)
if dimension_combination_is_used(c)]

@testset "$test_type" begin
Expand All @@ -69,7 +69,7 @@ function runtests(; restart=false)
continue
end

if :anyv ∈ debug_loop_type
if :anysv ∈ debug_loop_type
dims_to_test = debug_loop_type[2:end]
else
dims_to_test = debug_loop_type
Expand Down
86 changes: 43 additions & 43 deletions moment_kinetics/src/communication.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,12 @@ module communication

export allocate_shared, block_rank, block_size, n_blocks, comm_block, comm_inter_block,
iblock_index, comm_world, finalize_comms!, initialize_comms!, global_rank,
MPISharedArray, global_size, comm_anyv_subblock, anyv_subblock_rank,
anyv_subblock_size, anyv_isubblock_index, anyv_nsubblocks_per_block
MPISharedArray, global_size, comm_anysv_subblock, anysv_subblock_rank,
anysv_subblock_size, anysv_isubblock_index, anysv_nsubblocks_per_block
export setup_distributed_memory_MPI
export setup_distributed_memory_MPI_for_weights_precomputation
export setup_serial_MPI
export @_block_synchronize, @_anyv_subblock_synchronize
export @_block_synchronize, @_anysv_subblock_synchronize

using LinearAlgebra
using MPI
Expand Down Expand Up @@ -58,18 +58,18 @@ MPI.jl delete the communicator.
const comm_inter_block = Ref(MPI.COMM_NULL)

"""
Communicator for the local velocity-space subset of a shared-memory block in a 'anyv'
Communicator for the local velocity-space subset of a shared-memory block in a 'anysv'
region

The 'anyv' region is used to parallelise the collision operator. See
[`moment_kinetics.looping.get_best_anyv_split`](@ref).
The 'anysv' region is used to parallelise the collision operator. See
[`moment_kinetics.looping.get_best_anysv_split`](@ref).

Must use a `Ref{MPI.Comm}` to allow a non-const `MPI.Comm` to be stored. Need to actually
assign to this and not just copy a pointer into the `.val` member because otherwise the
`MPI.Comm` object created by `MPI.Comm_split()` would be deleted, which probably makes
MPI.jl delete the communicator.
"""
const comm_anyv_subblock = Ref(MPI.COMM_NULL)
const comm_anysv_subblock = Ref(MPI.COMM_NULL)

# Use Ref for these variables so that they can be made `const` (so have a definite
# type), but contain a value assigned at run-time.
Expand All @@ -95,19 +95,19 @@ const block_size = Ref{mk_int}()

"""
"""
const anyv_subblock_rank = Ref{mk_int}()
const anysv_subblock_rank = Ref{mk_int}()

"""
"""
const anyv_subblock_size = Ref{mk_int}()
const anysv_subblock_size = Ref{mk_int}()

"""
"""
const anyv_isubblock_index = Ref{Union{mk_int,Nothing}}()
const anysv_isubblock_index = Ref{Union{mk_int,Nothing}}()

"""
"""
const anyv_nsubblocks_per_block = Ref{mk_int}()
const anysv_nsubblocks_per_block = Ref{mk_int}()

"""
"""
Expand Down Expand Up @@ -561,11 +561,11 @@ end
# instances, so their is_read and is_written members can be checked and
# reset by _block_synchronize()
const global_debugmpisharedarray_store = Vector{DebugMPISharedArray}(undef, 0)
# 'anyv' regions require a separate array store, because within an anyv region,
# 'anysv' regions require a separate array store, because within an anysv region,
# processes in the same shared memory block may still not be synchronized if they are
# in different anyv sub-blocks, so debug checks within an anyv region should only
# consider the anyv-specific arrays.
const global_anyv_debugmpisharedarray_store = Vector{DebugMPISharedArray}(undef, 0)
# in different anysv sub-blocks, so debug checks within an anysv region should only
# consider the anysv-specific arrays.
const global_anysv_debugmpisharedarray_store = Vector{DebugMPISharedArray}(undef, 0)
end

"""
Expand Down Expand Up @@ -684,8 +684,8 @@ function allocate_shared(T, dims; comm=nothing, maybe_debug=true)
# If @debug_shared_array is active, create DebugMPISharedArray instead of Array
if maybe_debug
debug_array = DebugMPISharedArray(array, comm)
if comm == comm_anyv_subblock[]
push!(global_anyv_debugmpisharedarray_store, debug_array)
if comm == comm_anysv_subblock[]
push!(global_anysv_debugmpisharedarray_store, debug_array)
else
push!(global_debugmpisharedarray_store, debug_array)
end
Expand Down Expand Up @@ -820,8 +820,8 @@ end
Can be added when debugging to help pin down where an error occurs.
"""
function debug_check_shared_memory(; comm=comm_block[], kwargs...)
if comm == comm_anyv_subblock[]
for array ∈ global_anyv_debugmpisharedarray_store
if comm == comm_anysv_subblock[]
for array ∈ global_anysv_debugmpisharedarray_store
debug_check_shared_array(array; comm=comm, kwargs...)
end
else
Expand Down Expand Up @@ -973,49 +973,49 @@ debugging routines need to be updated.
end
end

# Also check 'anyv' arrays, as these are synchronized by this call.
# Also check 'anysv' arrays, as these are synchronized by this call.
# `missing` passed as the call_site argument here indicates that the check of
# call_site has already been done.
_anyv_subblock_synchronize(missing)
_anysv_subblock_synchronize(missing)

MPI.Barrier(comm_block[])
end
end

"""
Call an MPI Barrier for all processors in an 'anyv' sub-block.
Call an MPI Barrier for all processors in an 'anysv' sub-block.

The 'anyv' region is used to parallelise the collision operator. See
[`moment_kinetics.looping.get_best_anyv_split`](@ref).
The 'anysv' region is used to parallelise the collision operator. See
[`moment_kinetics.looping.get_best_anysv_split`](@ref).

Used to synchronise processors that are working on the same shared-memory array(s)
between operations, to avoid race conditions. Should be even cheaper than
[`@_block_synchronize`](@ref) because it only requires communication on a smaller
communicator.

Note: `_anyv_subblock_synchronize()` may be called different numbers of times on different
Note: `_anysv_subblock_synchronize()` may be called different numbers of times on different
sub-blocks, depending on how the species and spatial dimensions are split up.
`@debug_detect_redundant_block_synchronize` is not implemented (yet?) for
`_anyv_subblock_synchronize()`.
`_anysv_subblock_synchronize()`.
"""
macro _anyv_subblock_synchronize()
macro _anysv_subblock_synchronize()
id_hash = @debug_block_synchronize_quick_ifelse(
hash(string(@__FILE__, @__LINE__)),
nothing
)
return :( _anyv_subblock_synchronize($id_hash) )
return :( _anysv_subblock_synchronize($id_hash) )
end

"""
Internal function called by `anyv` synchronization macros.
Internal function called by `anysv` synchronization macros.
"""
function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
if comm_anyv_subblock[] == MPI.COMM_NULL
function _anysv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
if comm_anysv_subblock[] == MPI.COMM_NULL
# No synchronization to do for a null communicator
return nothing
end

MPI.Barrier(comm_anyv_subblock[])
MPI.Barrier(comm_anysv_subblock[])

@debug_block_synchronize_backtrace begin
st = stacktrace()
Expand All @@ -1028,12 +1028,12 @@ function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
signaturestring = string([string(s.file, s.line) for s ∈ st]...)

hash = sha256(signaturestring)
all_hashes = reshape(MPI.Allgather(hash, comm_anyv_subblock[]), length(hash),
MPI.Comm_size(comm_anyv_subblock[]))
all_hashes = reshape(MPI.Allgather(hash, comm_anysv_subblock[]), length(hash),
MPI.Comm_size(comm_anysv_subblock[]))
for i ∈ 1:block_size[]
h = all_hashes[:, i]
if h != hash
error("_anyv_subblock_synchronize() called inconsistently\n",
error("_anysv_subblock_synchronize() called inconsistently\n",
"rank $(block_rank[]) called from:\n",
stackstring)
end
Expand All @@ -1048,9 +1048,9 @@ function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
# If call_site===missing, then this function was called from inside
# _block_synchronize(), and the call site was already checked there.
if call_site !== missing
all_hashes = MPI.Allgather(call_site, comm_anyv_subblock[])
all_hashes = MPI.Allgather(call_site, comm_anysv_subblock[])
if !all(h -> h == all_hashes[1], all_hashes)
error("_anyv_subblock_synchronize() called inconsistently")
error("_anysv_subblock_synchronize() called inconsistently")
end
end
end
Expand All @@ -1063,9 +1063,9 @@ function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
# * If an element is written to, only the rank that writes to it should read it.
#
@debug_detect_redundant_block_synchronize previous_was_unnecessary = true
for array ∈ global_anyv_debugmpisharedarray_store
for array ∈ global_anysv_debugmpisharedarray_store

debug_check_shared_array(array; comm=comm_anyv_subblock[])
debug_check_shared_array(array; comm=comm_anysv_subblock[])

@debug_detect_redundant_block_synchronize begin
# debug_detect_redundant_is_active[] is set to true at the beginning of
Expand All @@ -1077,7 +1077,7 @@ function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
if debug_detect_redundant_is_active[]

if !debug_check_shared_array(array; check_redundant=true,
comm_anyv_subblock)
comm_anysv_subblock)
# If there was a failure for at least one array, the previous
# _block_synchronize was necessary - if the previous call was not
# there, for this array array.is_read and array.is_written would
Expand Down Expand Up @@ -1107,7 +1107,7 @@ function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
# Check the previous call was unnecessary on all processes, not just
# this one
previous_was_unnecessary = MPI.Allreduce(previous_was_unnecessary,
MPI.Op(&, Bool), comm_anyv_subblock[])
MPI.Op(&, Bool), comm_anysv_subblock[])

if (previous_was_unnecessary && global_size[] > 1)
# The intention of this debug block is to detect when calls to
Expand All @@ -1134,7 +1134,7 @@ function _anyv_subblock_synchronize(call_site::Union{Nothing,Missing,UInt64})
end
end

MPI.Barrier(comm_anyv_subblock[])
MPI.Barrier(comm_anysv_subblock[])
end
end

Expand Down Expand Up @@ -1177,7 +1177,7 @@ end
"""
function free_shared_arrays()
@debug_shared_array resize!(global_debugmpisharedarray_store, 0)
@debug_shared_array resize!(global_anyv_debugmpisharedarray_store, 0)
@debug_shared_array resize!(global_anysv_debugmpisharedarray_store, 0)

for w ∈ global_Win_store
MPI.free(w)
Expand Down
2 changes: 1 addition & 1 deletion moment_kinetics/src/electron_kinetic_equation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1885,7 +1885,7 @@ The r-dimension is not parallelised. For 1D runs this makes no difference. In 2D
or might not be necessary. If r-dimension parallelisation is needed, it would need some
work. The simplest option would be a non-parallelised outer loop over r, with each
nonlinear solve being parallelised over {z,vperp,vpa}. More efficient might be to add an
equivalent to the 'anyv' parallelisation used for the collision operator (e.g. 'anyzv'?)
equivalent to the 'anysv' parallelisation used for the collision operator (e.g. 'anyzv'?)
to allow the outer r-loop to be parallelised.
"""
@timeit global_timer implicit_electron_advance!(
Expand Down
Loading
Loading