Skip to content

Implement MPI shared-memory allocation for pointing quaternions to mitigate out-of-memory issues - #551

Open
anand-avinash wants to merge 11 commits into
masterfrom
shared_quat
Open

Implement MPI shared-memory allocation for pointing quaternions to mitigate out-of-memory issues#551
anand-avinash wants to merge 11 commits into
masterfrom
shared_quat

Conversation

@anand-avinash

@anand-avinash anand-avinash commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

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_quats and bore2ecliptic_quats were 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:

  1. Simulation.set_scanning_strategy_shmem() (alternative to Simulation.set_scanning_strategy): It splits the global MPI communicator by physical nodes and allocates a single copy of the spin2ecliptic_quats array 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.
  2. Observation.prepare_pointings_shmem() (alternative to Observation.prepare_pointings): It splits the time-block communicator of a given observation object by physical nodes and allocates a single copy of the bore2ecliptic_quats array 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 SharedRotQuaternion class, which is a subclass of RotQuaternion. 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 n processes per node, this implementation reduces the memory occupation of quaternion arrays by a factor of n-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 Simulation object allocates the spin2ecliptic_quats array for the entire simulation duration, and any Observation object allocates bore2ecliptic_quats for the entire simulation duration as well. However, in practice, only a small portion of these arrays - corresponding strictly to the time span of the Observation instance - is actually used to compute the detector pointings.

In the next step of optimization, we can address the following:

  1. Currently, spin2ecliptic_quats is an attribute of the Simulation class. Since the Simulation class has no information about the time durations of the specific Observation objects it manages, it has no choice but to compute spin2ecliptic_quats for the entire simulation duration.
  2. To resolve this, we can make spin2ecliptic_quats an attribute of the Observation class instead. Since an Observation object contains the exact timing information for a specific observation, we can allocate both spin2ecliptic_quats and bore2ecliptic_quats precisely 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_quats are never intended to be used independently of an Observation instance 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant