Implement MPI shared-memory allocation for pointing quaternions to mitigate out-of-memory issues - #551
Open
anand-avinash wants to merge 11 commits into
Open
Implement MPI shared-memory allocation for pointing quaternions to mitigate out-of-memory issues#551anand-avinash wants to merge 11 commits into
anand-avinash wants to merge 11 commits into
Conversation
…ation.set_scanning_strategy
…on.prepare_pointings
…_pointings in pointings_in_obs.py
…ion.prepare_pointings
…d bore2ecliptic quaternions
… class; fixed the quaternion multiplication in prepare_pointings_shmem of observation class
… test action workflow
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.
This PR provides a workaround for the out-of-memory issues that occur on multi-core HPC nodes during large simulations (#489) by introducing an MPI shared-memory allocator for pointing quaternions. Previously, the large quaternion arrays
spin2ecliptic_quatsandbore2ecliptic_quatswere independently allocated by every single MPI process, duplicating exactly the same contents. On a typical node with multiple ranks, this resulted in a massive footprint of redundant memory.To solve this, I have implemented a shared memory manager,
SharedMemoryManager, which splits an MPI communicator so that every process in the resulting sub-communicator has access to the same shared memory context (typically spanning a single physical node). This manager is then used to allocate the MPI shared-memory-backed NumPy arrays, ensuring only a single physical copy exists per sub-communicator.Rather than modifying the existing quaternion computation workflow and risking backwards incompatibility, I have introduced two new parallel interfaces:
Simulation.set_scanning_strategy_shmem()(alternative toSimulation.set_scanning_strategy): It splits the global MPI communicator by physical nodes and allocates a single copy of thespin2ecliptic_quatsarray that is shared among all MPI processes within a node. The values within the array are computed only by the node's root rank to prevent data race conditions.Observation.prepare_pointings_shmem()(alternative toObservation.prepare_pointings): It splits the time-block communicator of a given observation object by physical nodes and allocates a single copy of thebore2ecliptic_quatsarray that is shared among all MPI processes of that time-block communicator on any given node. Again, the values are computed strictly by the node's root rank.We wrap these shared quaternion arrays in a new
SharedRotQuaternionclass, which is a subclass ofRotQuaternion. It provides all the same functionality as the latter, except it skips the redundant internal normalization of quaternion samples (since the root rank handles normalization before sharing).After the quaternion arrays are computed, we use MPI fence functions to synchronize them across the shared memory context before they are accessed by other ranks. For a typical workflow using
nprocesses per node, this implementation reduces the memory occupation of quaternion arrays by a factor ofn-which can translate to freeing up to 100 GB of memory per node, effectively mitigating (#489) with zero computational overhead.The shared memory manager,
SharedMemoryManager, is a general-purpose class that can be used for other memory-related optimizations - for example, allocating a single copy of the large HEALPix maps only once per node.Future work and optimizations
Nonetheless, the current implementations (both the original and the new shared-memory workflows) can be optimized further. As noted by @paganol in this comment, any
Simulationobject allocates thespin2ecliptic_quatsarray for the entire simulation duration, and anyObservationobject allocatesbore2ecliptic_quatsfor the entire simulation duration as well. However, in practice, only a small portion of these arrays - corresponding strictly to the time span of theObservationinstance - is actually used to compute the detector pointings.In the next step of optimization, we can address the following:
spin2ecliptic_quatsis an attribute of theSimulationclass. Since theSimulationclass has no information about the time durations of the specificObservationobjects it manages, it has no choice but to computespin2ecliptic_quatsfor the entire simulation duration.spin2ecliptic_quatsan attribute of theObservationclass instead. Since anObservationobject contains the exact timing information for a specific observation, we can allocate bothspin2ecliptic_quatsandbore2ecliptic_quatsprecisely to cover only that required time duration.Since the quaternion computation will be limited to a fraction of the total simulation duration, it will become significantly cheaper computationally. Because
spin2ecliptic_quatsare never intended to be used independently of anObservationinstance in the existing codebase, this refactor will not require massive architectural changes.Once this extra optimization is finalized, in conjunction with the new MPI shared-memory workflow (
Observation.prepare_pointings_shmem()), the pointing computations will become exceptionally cheap and require a tiny fraction of the memory footprint, which so far, has been the limiting factor. In turn, this will allow us to pack larger simulations into far fewer computational resources in the future.