Introducing MPI shared memory based weight pre-processing - #29
Merged
Merged
Conversation
…ons. The serial ones are transferred to compute_weights_commons.cpp
…sions of setuptools do not define it
…mons to compute_weights to simplify the workflow in prospective compute_weights_shared
…yManager class to facilitate allocating MPI shared memory from python side
… memory arrays on cpp side
…TimeSamples to make them the sub class of BaseProcessTimeSamples
…rocessTimeSamples
…ray to SharedMemoryManager
…emProcessTimeSamples
…ager from Comm object to Comm.handle, as Comm is not hashable - and can't be used as a dict key
….cpp to ensure that the fence calls are made globally in order to avoid deadlocks
…o SharedMemProcessTimeSamples to enable freeing shared memory window allocated by it
… memory PTS to free shared memory arrays appropriately
…o avoid resource exhaustion
…) methods of the shared memory manager
…that frees shared memory window, arrays, and the communicators created by the manager; updated the relevant tests to use the new method
… and iterations through benchmarks.pedantic() function to prevent segmentation fault in running benchmark on mpi related functions; added cli args to supply number of rounds, iterations, and warmup rounds
…sary with appropriate setup call
…amples, pointing operator, and gls wrapper
… after each round whenever a shared memory time samples container is created
… the command line
anand-avinash
marked this pull request as ready for review
August 6, 2026 11:09
… level arrays are allocated unconditionally irrespective of whether nproc_reduce is one or greater; this prevents race condition led by multiple grp comm level arrays writing to same node root arrays at the same time whenever nproc_reduce was one
…ension; now grp comm level arrays are allocated unconditionally irrespective of whether nproc_reduce is one or greater
…ved grp_reduce arguments from benchmark tests
…o tests and benchmarks
…hile running the tests as github actions
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
1 participant
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.
This PR introduces
SharedMemProcessTimeSamples, an MPI shared-memory optimized alternative to the standardProcessTimeSamplesdata container. It utilizes MPI shared-memory features (viampi4pyand custom C++ bindings) to significantly reduce the memory footprint when preprocessing pointing information and map-making weights across multi-process, multi-node environments. Consequently, it also reduces the amount of inter-node MPI communication significantly.Instead of replicating the pixel-space data structures (such as hit counts and trigonometric weights) on every MPI process on a node,
SharedMemProcessTimeSamplesallocates node-level shared-memory windows, aggregates weights using sub-communicators, and makes the results visible across all node ranks.The implementation details are as follows:
Shared memory allocator/manager on Python side: The changes include the implementation of a shared memory allocator/manager (
SharedMemoryManagerin mpi.py). It provides an interface to split a given MPI communicator into communicators with shared memory access. It also implements a robust communicator split hierarchy (_node_comm,_node_root_comm,_tree_grp_comm,_tree_grp_root_comm) for shared memory topologies and correctly allocates NumPy arrays backed by shared MPI windows withMPI.Win.Allocate_shared.Shared memory allocator/manager on C++ side: We have implemented an analogous shared memory allocator/manager on the C++ side in
mpi_utils.hppto allow direct allocation and querying of MPI shared windows in C++ routines.Implementation of
SharedMemProcessTimeSamples: It is a shared memory analogue of the classicProcessTimeSamplescontainer. While it exposes the same attributes asProcessTimeSamples, they point to shared-memory buffers that are allocated only once for each shared memory domain (usually a single compute node). The corresponding C++ backends are defined incompute_weights_shared.cpp.Additionally, this PR updates documentation and refactors multiple tests to use fixtures, which improves efficiency and reduces code repetition.