From 104b47eae57f6a2055851054b862bc2e0a231ccb Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:39:07 +0200 Subject: [PATCH 01/10] added a shared memory manager --- litebird_sim/shared_memory.py | 276 ++++++++++++++++++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 litebird_sim/shared_memory.py diff --git a/litebird_sim/shared_memory.py b/litebird_sim/shared_memory.py new file mode 100644 index 00000000..93a0f471 --- /dev/null +++ b/litebird_sim/shared_memory.py @@ -0,0 +1,276 @@ +from typing import cast +import numpy as np +import numpy.typing as npt + +try: + from mpi4py import MPI + from mpi4py.MPI import Intracomm +except ImportError: + pass + + +class SharedMemoryManager: + """Manages MPI shared-memory communicators, window allocations, and tree + group reductions. + + This manager splits a base MPI communicator into a node-level + shared-memory communicator and allocates MPI window-backed shared NumPy + arrays. It also splits the node-level communicator into a tree group + sub-communicators to orchestrate sequential accumulations and group-wise + reductions within each node. + + Parameters + ---------- + base_comm : Intracomm + The base MPI communicator (typically `MPI.COMM_WORLD`) + node_root : int, optional + The designated root rank within the node-level shared memory + communicator. By default `0` + + Attributes + ---------- + base_comm : Intracomm + The base MPI communicator + node_comm : Intracomm + The node-level shared-memory MPI communicator + node_rank : int + The process rank within the node-level communicator + node_size : int + The total number of processes on the current node + node_root : int + The root rank on the current node communicator + list_windows : dict[int, list[MPI.Win]] + Tracks allocated shared-memory MPI windows mapped by communicator handle + list_arrays : dict[int, list[npt.NDArray]] + Tracks allocated shared-memory NumPy array views mapped by + communicator handle + """ + + def __init__( + self, + base_comm: "Intracomm", + node_root: int = 0, + ) -> None: + if "MPI" not in globals(): + raise ImportError( + "mpi4py is not installed or enabled, but shared memory allocation was requested." + ) + + self._base_comm = base_comm + self._node_comm: Intracomm = cast( + Intracomm, self._base_comm.Split_type(MPI.COMM_TYPE_SHARED) + ) + self._node_rank = self._node_comm.rank + self._node_size = self._node_comm.size + self._node_root = node_root + + # List of MPI shared memory windows + self._list_windows: dict[int, list[MPI.Win]] = {} + self._list_arrays: dict[int, list[npt.NDArray]] = {} + + @property + def base_comm(self) -> "Intracomm": + """The base global MPI communicator""" + return self._base_comm + + @property + def node_comm(self) -> "Intracomm": + """The node-level shared-memory MPI communicator""" + return self._node_comm + + @property + def node_rank(self) -> int: + """The process rank within the node-level communicator""" + return self._node_rank + + @property + def node_size(self) -> int: + """The total number of processes on the current node""" + return self._node_size + + @property + def node_root(self) -> int: + """The root rank on the current node""" + return self._node_root + + @property + def list_windows(self) -> dict: + """The dictionary mapping communicators to list of allocated + shared-memory windows for that communicator + """ + return self._list_windows + + @property + def list_arrays(self) -> dict: + """The dictionary mapping communicators to list of allocated + shared-memory arrays for that communicator + """ + return self._list_arrays + + def alloc_shared_comm( + self, + size: int, + dtype: npt.DTypeLike, + comm: "Intracomm", + comm_root: int = 0, + ) -> tuple[npt.NDArray, "MPI.Win"]: + """Allocates a shared-memory MPI window-backed 1D NumPy array for a + communicator. + """ + dtype = np.dtype(dtype) + dtype_bytes = dtype.itemsize + arr_bytes = size * dtype_bytes if comm.rank == comm_root else 0 + + win = MPI.Win.Allocate_shared( + arr_bytes, + dtype_bytes, + comm=comm, + ) + buf, _ = win.Shared_query(rank=comm_root) + # np.ndarray provides the view, it doesn't owns the memory + array = np.ndarray(shape=size, dtype=dtype, buffer=buf) + + handle = comm.handle + if handle not in self._list_windows: + self._list_windows[handle] = [] + if handle not in self._list_arrays: + self._list_arrays[handle] = [] + + self._list_windows[handle].append(win) + self._list_arrays[handle].append(array) + return array, win + + def alloc_shared_node( + self, + size: int, + dtype: npt.DTypeLike, + ) -> tuple[npt.NDArray, "MPI.Win"]: + """Allocates a shared-memory MPI window-backed 1D NumPy array for the + node-level communicator + """ + return self.alloc_shared_comm( + size=size, + dtype=dtype, + comm=self.node_comm, + comm_root=self.node_root, + ) + + def alloc_shared_zeros_comm( + self, + size: int, + dtype: npt.DTypeLike, + comm: "Intracomm", + comm_root: int = 0, + ): + """Allocates a shared-memory MPI window-backed 1D NumPy array for a + communicator, initialized to zeros. + """ + array, win = self.alloc_shared_comm( + size=size, + dtype=dtype, + comm=comm, + comm_root=comm_root, + ) + + if comm.rank == 0: + array[:] = 0 + + return array, win + + def alloc_shared_zeros_node( + self, + size: int, + dtype: npt.DTypeLike, + ): + """Allocates a shared-memory MPI window-backed 1D NumPy array for the + node-level communicator, initialized to zeros. + """ + return self.alloc_shared_zeros_comm( + size=size, + dtype=dtype, + comm=self.node_comm, + comm_root=self.node_root, + ) + + def alloc_shared_ones_comm( + self, + size: int, + dtype: npt.DTypeLike, + comm: "Intracomm", + comm_root: int = 0, + ): + """Allocates a shared-memory MPI window-backed 1D NumPy array for a + communicator, initialized to ones. + """ + array, win = self.alloc_shared_comm( + size=size, + dtype=dtype, + comm=comm, + comm_root=comm_root, + ) + + if comm.rank == 0: + array[:] = 1 + + return array, win + + def alloc_shared_ones_node( + self, + size: int, + dtype: npt.DTypeLike, + ): + """Allocates a shared-memory MPI window-backed 1D NumPy array for the + node-level communicator, initialized to ones. + """ + return self.alloc_shared_ones_comm( + size=size, + dtype=dtype, + comm=self.node_comm, + comm_root=self.node_root, + ) + + def fence_comm_all(self, comm: "Intracomm", assertion: int = 0) -> None: + """Call MPI.Win.Fence on all windows allocated on the given + communicator. + """ + handle = comm.handle + if handle in self._list_windows: + for win in self._list_windows[handle]: + win.Fence(assertion) + + def free_shared_arrays_all(self) -> None: + """Frees all allocated shared-memory MPI windows and clears manager + state. + """ + for comm, wins in self._list_windows.items(): + for win in wins: + win.Free() + self._list_windows = {} + self._list_arrays = {} + + def free_shared_arrays_comm(self, comm: "Intracomm") -> None: + """Frees all shared-memory MPI windows allocated for a specific + communicator. + """ + handle = comm.handle + if handle in self._list_windows: + for win in self._list_windows[handle]: + win.Free() + del self._list_windows[handle] + if handle in self._list_arrays: + del self._list_arrays[handle] + + def free_shared_array(self, comm: "Intracomm", win: "MPI.Win") -> None: + """Frees a specific shared-memory MPI window and removes its associated + array view and window from the manager's tracking lists. + """ + handle = comm.handle + if handle in self._list_windows and win in self._list_windows[handle]: + idx = self._list_windows[handle].index(win) + win.Free() + self._list_windows[handle].pop(idx) + self._list_arrays[handle].pop(idx) + if not self._list_windows[handle]: + del self._list_windows[handle] + if not self._list_arrays[handle]: + del self._list_arrays[handle] From b8c68b6344e7e96b4d98b4d48c9e88078d484f8d Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:44:46 +0200 Subject: [PATCH 02/10] Added a SharedRotQuaternion class that is the shared memory analogue of RotQuaternion --- litebird_sim/scanning.py | 44 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/litebird_sim/scanning.py b/litebird_sim/scanning.py index 3d502d6a..e643a702 100644 --- a/litebird_sim/scanning.py +++ b/litebird_sim/scanning.py @@ -636,6 +636,50 @@ def is_close_to(self, other: "RotQuaternion") -> bool: return np.allclose(self.quats, other.quats) +class SharedRotQuaternion(RotQuaternion): + """A version of RotQuaternion that wraps an existing shared-memory NumPy array. + + This class is used to avoid copying or modifying (normalizing) the underlying + array, as that is expected to be handled explicitly by the + node root process of the shared memory communicator. + """ + + def __init__( + self, + quats: npt.NDArray, + start_time: float | astropy.time.Time | None = None, + sampling_rate_hz: float | None = None, + ): + """ + Create a new instance of a time-dependent quaternion + + If both `start_time` and `sampling_freq_hz` are ``None``, the quaternion + is assumed to be constant in time. + + :param quats: Either a 4-element NumPy array, or another instance + of :class:`TimeDependentQuaternion` + :param start_time: the start time, either a floating point number + or an ``astropy.time.Time`` object + :param sampling_rate_hz: the sampling frequency + """ + # Directly use the provided array without reshaping or copying + self.quats = quats + + if self.quats.shape[0] > 1: + assert start_time is not None, ( + "You must specify start_time if the quaternion is not constant" + ) + assert sampling_rate_hz is not None, ( + "You must specify sampling_rate_hz if the quaternion is not constant" + ) + + self.start_time = start_time + self.sampling_rate_hz = sampling_rate_hz + + # We DO NOT normalize quaternions here. + # The node root is expected to normalize them after writing to shared memory. + + # This is an Abstract Base Class (ABC) class ScanningStrategy(ABC): """A class that simulate a scanning strategy From 194c256b4433b7617653e150aba23d50737f109b Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:10:20 +0200 Subject: [PATCH 03/10] added mpi shared memory analogue set_scanning_strategy_shmem of Simulation.set_scanning_strategy --- litebird_sim/simulations.py | 101 +++++++++++++++++++++++++++++++++++- 1 file changed, 100 insertions(+), 1 deletion(-) diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py index 10ee60ff..9a22f3c3 100644 --- a/litebird_sim/simulations.py +++ b/litebird_sim/simulations.py @@ -71,7 +71,8 @@ ) from .profiler import TimeProfiler, profile_list_to_speedscope from .scan_map import scan_map_in_observations -from .scanning import ScanningStrategy, SpinningScanningStrategy +from .scanning import ScanningStrategy, SharedRotQuaternion, SpinningScanningStrategy +from .shared_memory import SharedMemoryManager from .seeding import RNGHierarchy from .spacecraft import SpacecraftOrbit, spacecraft_pos_and_vel from .units import Units @@ -1576,6 +1577,104 @@ def set_scanning_strategy( quat_memory_size_bytes=quat_memory_size_bytes, ) + def set_scanning_strategy_shmem( + self, + scanning_strategy: ScanningStrategy | None = None, + imo_url: str | None = None, + delta_time_s: float = 60.0, + append_to_report: bool = True, + ): + """ + Works identically to `set_scanning_strategy`, but uses MPI shared memory + to allocate the spin2ecliptic quaternions only once per physical node. + """ + assert not (scanning_strategy and imo_url), ( + "you must either specify scanning_strategy or imo_url (but not" + "the two together) when calling Simulation.set_shared_scanning_strategy" + ) + + if not scanning_strategy: + if not imo_url: + imo_url = "/releases/v1.0/satellite/scanning_parameters/" + scanning_strategy = SpinningScanningStrategy.from_imo( + imo=self.imo, url=imo_url + ) + + try: + import mpi4py # noqa + except ImportError: + raise RuntimeError( + "`mpi4py` is required to set MPI shared memory scanning strategy." + ) + + shared_manager = SharedMemoryManager(base_comm=self.mpi_comm) + + # Only the node root computes the array to avoid redundant allocations on other ranks + if shared_manager.node_rank == shared_manager.node_root: + local_spin2ecl = scanning_strategy.generate_spin2ecl_quaternions( + start_time=self.start_time, + time_span_s=self.duration_s, + delta_time_s=delta_time_s, + ) + n_samples = local_spin2ecl.quats.shape[0] + dtype = local_spin2ecl.quats.dtype + else: + n_samples = 0 + dtype = np.float64 + + n_samples = shared_manager.node_comm.bcast( + n_samples, root=shared_manager.node_root + ) + dtype = shared_manager.node_comm.bcast(dtype, root=shared_manager.node_root) + + flat_array, _ = shared_manager.alloc_shared_node( + size=n_samples * 4, + dtype=dtype, + ) + shared_quats_view = flat_array.reshape(n_samples, 4) + + if shared_manager.node_rank == shared_manager.node_root: + shared_quats_view[:] = local_spin2ecl.quats + start_time_val = local_spin2ecl.start_time + sampling_rate_hz = local_spin2ecl.sampling_rate_hz + else: + start_time_val = 0.0 + sampling_rate_hz = 0.0 + + shared_manager.fence_comm_all(shared_manager.node_comm) + + start_time_val = shared_manager.node_comm.bcast( + start_time_val, root=shared_manager.node_root + ) + sampling_rate_hz = shared_manager.node_comm.bcast( + sampling_rate_hz, root=shared_manager.node_root + ) + + self.spin2ecliptic_quats = SharedRotQuaternion( + quats=shared_quats_view, + start_time=start_time_val, + sampling_rate_hz=sampling_rate_hz, + ) + self.shmem_manager = shared_manager + + quat_memory_size_bytes = self.spin2ecliptic_quats.nbytes() + + num_of_obs = len(self.observations) + if append_to_report and MPI_ENABLED: + num_of_obs = self.mpi_comm.allreduce(num_of_obs) + + if append_to_report and MPI_COMM_WORLD.rank == 0: + template_file_path = get_template_file_path("report_quaternions.md") + with template_file_path.open("rt") as inpf: + markdown_template = "".join(inpf.readlines()) + self.append_to_report( + markdown_template, + num_of_obs=num_of_obs, + num_of_mpi_processes=MPI_COMM_WORLD.size, + delta_time_s=delta_time_s, + quat_memory_size_bytes=quat_memory_size_bytes, + ) + def set_instrument(self, instrument: InstrumentInfo): """Set the instrument to be used in the simulation. From 9429166e2c358695d436318db47158c71927f6f5 Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:25:49 +0200 Subject: [PATCH 04/10] added mpi shared memory analogue prepare_pointings_shmem of Observation.prepare_pointings --- litebird_sim/observations.py | 79 +++++++++++++++++++++++++++++++++++- 1 file changed, 78 insertions(+), 1 deletion(-) diff --git a/litebird_sim/observations.py b/litebird_sim/observations.py index 31f85c3a..8d976123 100644 --- a/litebird_sim/observations.py +++ b/litebird_sim/observations.py @@ -15,9 +15,11 @@ from .input_sky import SkyInput from .mpi import _SerialMpiCommunicator from .pointings import DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB, PointingProvider -from .scanning import RotQuaternion +from .scanning import RotQuaternion, SharedRotQuaternion from .units import Units from .utilities import resolve_nthreads +from .quaternions import normalize_quaternions +from .shared_memory import SharedMemoryManager @dataclass @@ -983,6 +985,81 @@ def prepare_pointings( self.hwp = hwp self.has_hwp = True + def prepare_pointings_shmem( + self, + instrument: InstrumentInfo, + spin2ecliptic_quats: RotQuaternion, + hwp: HWP | None = None, + maximum_internal_buffer_mem_mb: float = DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB, + ) -> None: + """Prepare quaternion-based pointing and HWP information using MPI shared memory. + + This functions similarly to `prepare_pointings` but uses the provided + `SharedMemoryManager` to allocate the underlying pointing quaternion arrays + in node-shared memory, optimizing the global memory usage. + """ + + assert (maximum_internal_buffer_mem_mb > 0) or ( + maximum_internal_buffer_mem_mb == -1 + ), ( + "Invalid value for maximum_internal_buffer_mem_mb ({val}), it must either be -1 or a positive number".format( + val=maximum_internal_buffer_mem_mb + ) + ) + + if not hasattr(self.comm_time_block, "Split_type"): + raise RuntimeError( + "MPI is required for shared memory pointings, but the current time block communicator is not a valid MPI communicator." + ) + + self.shared_memory_manager = SharedMemoryManager(base_comm=self.comm_time_block) + + n_quats = spin2ecliptic_quats.quats.shape[0] + dtype = spin2ecliptic_quats.quats.dtype + + # Allocate flat 1D array on node shared memory + flat_array, _ = self.shared_memory_manager.alloc_shared_node( + size=n_quats * 4, + dtype=dtype, + ) + + # Create a 2D view + shared_quats_view = flat_array.reshape(n_quats, 4) + + # Only the node root computes the multiplication and fills the array + if self.shared_memory_manager.node_rank == self.shared_memory_manager.node_root: + shared_quats_view[:] = spin2ecliptic_quats * instrument.bore2spin_quat.quats + normalize_quaternions(shared_quats_view) + + # Synchronize all ranks on this node + self.shared_memory_manager.fence_comm_all(self.shared_memory_manager.node_comm) + + # Wrap into SharedRotQuaternion (which does not normalize) + shared_rot_quat = SharedRotQuaternion( + quats=shared_quats_view, + start_time=spin2ecliptic_quats.start_time, + sampling_rate_hz=spin2ecliptic_quats.sampling_rate_hz, + ) + + pointing_provider = PointingProvider( + bore2ecliptic_quats=shared_rot_quat, + hwp=hwp, + maximum_internal_buffer_mem_mb=maximum_internal_buffer_mem_mb, + ) + + self.pointing_provider = pointing_provider + + # If the hwp object is passed and is not initialised in the observations, it gets applied to all detectors + if hwp is None: + assert self.no_mueller_hwp() or self.no_jones_hwp(), ( + "Some detectors have been initialized with a mueller_hwp or jones_hwp," + "but no HWP object has been passed to prepare_shared_pointings." + ) + self.has_hwp = False + else: + self.hwp = hwp + self.has_hwp = True + def get_pointings( self, detector_idx: int | list[int] | str = "all", From b0542c89c6c01eccae0421a0fb8887dc9a637620 Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:28:52 +0200 Subject: [PATCH 05/10] added a MPI shared memory analogue prepare_pointings_shmem of prepare_pointings in pointings_in_obs.py --- litebird_sim/pointings_in_obs.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/litebird_sim/pointings_in_obs.py b/litebird_sim/pointings_in_obs.py index edef0bc9..1215ee97 100644 --- a/litebird_sim/pointings_in_obs.py +++ b/litebird_sim/pointings_in_obs.py @@ -65,6 +65,29 @@ def prepare_pointings( ) +def prepare_pointings_shmem( + observations: Observation | list[Observation], + instrument: InstrumentInfo, + spin2ecliptic_quats: RotQuaternion, + hwp: HWP | None = None, +) -> None: + """Initialize pointing and HWP information using MPI shared memory. + + This acts just like `prepare_pointings`, but relies on `SharedMemoryManager` + to allocate the underlying pointing quaternion arrays into MPI node-shared + memory, thereby preventing redundant memory allocation on multi-core nodes. + """ + if isinstance(observations, Observation): + obs_list = [observations] + else: + obs_list = observations + + for cur_obs in obs_list: + cur_obs.prepare_pointings_shmem( + instrument=instrument, spin2ecliptic_quats=spin2ecliptic_quats, hwp=hwp + ) + + def precompute_pointings( observations: Observation | list[Observation], pointings_dtype=np.float64, From 4c112566a76e60e5a520d9f2084885362e7ebd10 Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:36:44 +0200 Subject: [PATCH 06/10] added a MPI shared memory analogue prepare_pointings_shmem of Simulation.prepare_pointings --- litebird_sim/simulations.py | 73 ++++++++++++++++++++++++++++--------- 1 file changed, 55 insertions(+), 18 deletions(-) diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py index 9a22f3c3..68eff859 100644 --- a/litebird_sim/simulations.py +++ b/litebird_sim/simulations.py @@ -68,6 +68,7 @@ from .pointings_in_obs import ( precompute_pointings, prepare_pointings, + prepare_pointings_shmem, ) from .profiler import TimeProfiler, profile_list_to_speedscope from .scan_map import scan_map_in_observations @@ -1709,21 +1710,11 @@ def set_hwp(self, hwp: HWP): obs.set_hwp(hwp) @_profile - def prepare_pointings( + def _prepare_pointings( self, + use_shmem: bool = False, append_to_report: bool = True, ): - """Trigger the computation of the quaternions needed to compute pointings. - - This method must be called after having set the scanning strategy, the - instrument, the HWP, and the list of detectors to simulate through calls to - :meth:`.set_instrument` and :meth:`.add_detector`. A set of observations must - have been created using the method :meth:`.create_observations`. - - It combines the quaternions of the spacecraft, of the instrument, and of the detectors - and prepares a number of data structures that will be used by the method - :meth:`.Observation.get_pointings` to determine the pointing angles and the HWP angle. - """ assert self.observations, ( "You must call Simulation.create_observations() " "before calling Simulation.prepare_pointings" @@ -1737,12 +1728,20 @@ def prepare_pointings( "before calling Simulation.prepare_pointings" ) - prepare_pointings( - observations=self.observations, - instrument=self.instrument, - spin2ecliptic_quats=self.spin2ecliptic_quats, - hwp=self.hwp, - ) + if not use_shmem: + prepare_pointings( + observations=self.observations, + instrument=self.instrument, + spin2ecliptic_quats=self.spin2ecliptic_quats, + hwp=self.hwp, + ) + else: + prepare_pointings_shmem( + observations=self.observations, + instrument=self.instrument, + spin2ecliptic_quats=self.spin2ecliptic_quats, + hwp=self.hwp, + ) pointing_provider = self.observations[0].pointing_provider @@ -1764,6 +1763,44 @@ def prepare_pointings( memory_occupation=int(memory_occupation), ) + @_profile + def prepare_pointings( + self, + append_to_report: bool = True, + ): + """Trigger the computation of the quaternions needed to compute pointings. + + This method must be called after having set the scanning strategy, the + instrument, the HWP, and the list of detectors to simulate through calls to + :meth:`.set_instrument` and :meth:`.add_detector`. A set of observations must + have been created using the method :meth:`.create_observations`. + + It combines the quaternions of the spacecraft, of the instrument, and of the detectors + and prepares a number of data structures that will be used by the method + :meth:`.Observation.get_pointings` to determine the pointing angles and the HWP angle. + """ + self._prepare_pointings( + use_shmem=False, + append_to_report=append_to_report, + ) + + @_profile + def prepare_pointings_shmem( + self, + append_to_report: bool = True, + ): + """Trigger the computation of the quaternions needed to compute pointings, + using MPI shared memory. + + This functions identically to `prepare_pointings`, but coordinates memory + allocation across MPI ranks on the same physical node, + optimizing the global memory usage. + """ + self._prepare_pointings( + use_shmem=True, + append_to_report=append_to_report, + ) + def precompute_pointings(self, pointings_dtype=np.float64) -> None: """Compute all the pointings for all observations and save them From b12515f9fdcb168d22912651c97394f2397cbe5c Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:00:53 +0200 Subject: [PATCH 07/10] added the tests for MPI shared memory computation of spin2ecliptic and bore2ecliptic quaternions --- test/test_shared_memory_pointings.py | 139 +++++++++++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 test/test_shared_memory_pointings.py diff --git a/test/test_shared_memory_pointings.py b/test/test_shared_memory_pointings.py new file mode 100644 index 00000000..1390c48e --- /dev/null +++ b/test/test_shared_memory_pointings.py @@ -0,0 +1,139 @@ +import numpy as np + +import litebird_sim as lbs +from litebird_sim.scanning import RotQuaternion, SharedRotQuaternion +import pytest + +pytest.importorskip( + modname="mpi4py", + reason="`mpi4py` is required to run MPI shared memory tests", +) +from mpi4py import MPI # noqa: E402 + + +def test_shared_scanning_strategy(tmp_path): + + # Create two simulations + sim_std = lbs.Simulation( + base_path=tmp_path / "simulation_std", + start_time=0.0, + duration_s=10.0, + random_seed=12345, + mpi_comm=MPI.COMM_WORLD, + ) + + sim_shared = lbs.Simulation( + base_path=tmp_path / "simulation_shared", + start_time=0.0, + duration_s=10.0, + random_seed=12345, + mpi_comm=MPI.COMM_WORLD, + ) + + # dummy scanning strategy class + class DummyScanningStrategy(lbs.ScanningStrategy): + def generate_spin2ecl_quaternions(self, start_time, time_span_s, delta_time_s): + n_samples = int(np.ceil(time_span_s / delta_time_s)) + quats = np.random.randn(n_samples, 4) + quats /= np.linalg.norm(quats, axis=1)[:, np.newaxis] + return RotQuaternion( + quats, start_time=start_time, sampling_rate_hz=1.0 / delta_time_s + ) + + scanning_strategy = DummyScanningStrategy() + + # We set a random seed here so the two instances generate the same quaternions + np.random.seed(123) + sim_std.set_scanning_strategy( + scanning_strategy=scanning_strategy, delta_time_s=1.0, append_to_report=False + ) + + np.random.seed(123) + sim_shared.set_scanning_strategy_shmem( + scanning_strategy=scanning_strategy, delta_time_s=1.0, append_to_report=False + ) + + q_std = sim_std.spin2ecliptic_quats.quats + q_shared = sim_shared.spin2ecliptic_quats.quats + + np.testing.assert_allclose(q_std, q_shared) + assert isinstance(sim_shared.spin2ecliptic_quats, SharedRotQuaternion) + + +def test_shmem_bore2ecliptic_quats(tmp_path): + comm_size = MPI.COMM_WORLD.size + + # Create two simulations + sim_std = lbs.Simulation( + base_path=tmp_path / "simulation_std", + start_time=0.0, + duration_s=10.0, + random_seed=12345, + mpi_comm=MPI.COMM_WORLD, + ) + + sim_shared = lbs.Simulation( + base_path=tmp_path / "simulation_shared", + start_time=0.0, + duration_s=10.0, + random_seed=12345, + mpi_comm=MPI.COMM_WORLD, + ) + + # Add an instrument + instrument = lbs.InstrumentInfo( + name="test_inst", + spin_boresight_angle_rad=np.deg2rad(50.0), + ) + sim_std.set_instrument(instrument) + sim_shared.set_instrument(instrument) + + # Create observations (this handles comm_time_block and comm_det_block properly) + det1 = lbs.DetectorInfo("det1", sampling_rate_hz=10.0) + det2 = lbs.DetectorInfo("det2", sampling_rate_hz=10.0) + + # n_blocks_det * n_blocks_time must equal comm.size + sim_std.create_observations( + detectors=[det1, det2], + n_blocks_time=comm_size // 2 if comm_size % 2 == 0 else comm_size, + n_blocks_det=2 if comm_size % 2 == 0 else 1, + split_list_over_processes=False, + ) + sim_shared.create_observations( + detectors=[det1, det2], + n_blocks_time=comm_size // 2 if comm_size % 2 == 0 else comm_size, + n_blocks_det=2 if comm_size % 2 == 0 else 1, + split_list_over_processes=False, + ) + + # Fake spin2ecliptic quaternions + n_samples = sim_std.observations[0].n_samples_global + if sim_std.mpi_comm.rank == 0: + quats = np.random.randn(n_samples, 4) + quats /= np.linalg.norm(quats, axis=1)[:, np.newaxis] + else: + quats = np.empty((n_samples, 4), dtype=np.float64) + sim_std.mpi_comm.Bcast(quats, root=0) + spin2ecliptic_quats = RotQuaternion(quats, start_time=0.0, sampling_rate_hz=10.0) + + # Standard pointings preparation + sim_std.observations[0].prepare_pointings( + instrument=sim_std.instrument, spin2ecliptic_quats=spin2ecliptic_quats + ) + + # Shared memory pointings preparation + sim_shared.observations[0].prepare_pointings_shmem( + instrument=sim_shared.instrument, spin2ecliptic_quats=spin2ecliptic_quats + ) + + # Check if the outputs are identical + q_std = sim_std.observations[0].pointing_provider.bore2ecliptic_quats.quats + q_shared = sim_shared.observations[0].pointing_provider.bore2ecliptic_quats.quats + + np.testing.assert_allclose(q_std, q_shared) + + # Ensure that it is indeed a SharedRotQuaternion + assert isinstance( + sim_shared.observations[0].pointing_provider.bore2ecliptic_quats, + SharedRotQuaternion, + ) From 1c115ee5e0dd428e4325b9655e4d94ef80da3208 Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:28:48 +0200 Subject: [PATCH 08/10] fixed the conditions before splitting the communicator in observation class; fixed the quaternion multiplication in prepare_pointings_shmem of observation class --- litebird_sim/observations.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litebird_sim/observations.py b/litebird_sim/observations.py index 8d976123..9bcd3a26 100644 --- a/litebird_sim/observations.py +++ b/litebird_sim/observations.py @@ -1028,7 +1028,13 @@ def prepare_pointings_shmem( # Only the node root computes the multiplication and fills the array if self.shared_memory_manager.node_rank == self.shared_memory_manager.node_root: - shared_quats_view[:] = spin2ecliptic_quats * instrument.bore2spin_quat.quats + from .quaternions import multiply_quaternions_list_x_one + + multiply_quaternions_list_x_one( + spin2ecliptic_quats.quats, + instrument.bore2spin_quat.quats[0], + shared_quats_view, + ) normalize_quaternions(shared_quats_view) # Synchronize all ranks on this node @@ -1347,7 +1353,7 @@ def _set_mpi_subcommunicators(self): self.comm_det_block = _SerialMpiCommunicator() self.comm_time_block = _SerialMpiCommunicator() - if self.comm and self.comm.size > 1: + if hasattr(self.comm, "Split"): det_color = self.comm.rank // self.n_blocks_time time_color = self.comm.rank % self.n_blocks_time self.comm_det_block = self.comm.Split(det_color) From eeb75d33ad9c991091c2cd96a27c6bb8a5dc9ba0 Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:31:12 +0200 Subject: [PATCH 09/10] added the tests for MPI shared memory pointing computations to github test action workflow --- .github/workflows/tests.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 00357388..850039c5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -116,6 +116,7 @@ jobs: echo "Running MPI test ($MPI) with $proc processes" PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_mpi.py PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_detector_blocks.py + PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_shared_memory_pointings.py done echo "Running MPI test ($MPI) with 4 processes" PYTHONPATH=. mpiexec --map-by :OVERSUBSCRIBE -n 4 uv run python -m pytest ./test/test_mpi_n4.py @@ -131,6 +132,7 @@ jobs: echo "Running MPI test ($MPI) with $proc processes" PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_mpi.py PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_detector_blocks.py + PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_shared_memory_pointings.py done echo "Running MPI test ($MPI) with 4 processes" PYTHONPATH=. mpiexec --map-by :OVERSUBSCRIBE -n 4 uv run python -m pytest ./test/test_mpi_n4.py From d824f3bee207ac0a97ddd03b2c72531219316280 Mon Sep 17 00:00:00 2001 From: Avinash Anand <36325275+anand-avinash@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:57:26 +0200 Subject: [PATCH 10/10] updated changelog --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 94fabcb0..80ffbc20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,7 @@ # HEAD +- Introduced MPI shared-memory allocator for pointing quaternions (`SharedMemoryManager`) to drastically reduce memory usage and prevent OOM issues on multi-core HPC nodes (like [#489](https://github.com/litebird/litebird_sim/issues/489)). Added `Simulation.set_scanning_strategy_shmem()` and `Simulation.prepare_pointings_shmem()` to leverage this functionality. + - Stop supporting unused MPI processes and drop `MPI_COMM_GRID` class and its attributes. Now throwing an error when `comm_size != n_blocks_det * n_blocks_time`. It fixes [#364](https://github.com/litebird/litebird_sim/issues/364), while incorporating selective changes from [#372](https://github.com/litebird/litebird_sim/pull/372) and [#427](https://github.com/litebird/litebird_sim/pull/427). Reference PR: [#539](https://github.com/litebird/litebird_sim/pull/539). # Version 0.17.0