prism_q/distributed/
comm.rs

1//! Rank transport for distributed simulation.
2//!
3//! [`RankComm`] abstracts the collective and peer operations used by distributed
4//! backends. It is independent of state representation.
5//!
6//! [`SerialComm`] is the single rank implementation used by tests. `MpiComm`
7//! uses `rsmpi` behind the `distributed-mpi` feature and requires `mpiexec`.
8
9use num_complex::Complex64;
10
11#[cfg(feature = "distributed-mpi")]
12use crate::error::{PrismError, Result};
13#[cfg(feature = "distributed-mpi")]
14use mpi::environment::Threading;
15
16/// Collective and peer operations across a rank set.
17///
18/// The amplitude exchange routines treat `Complex64` as two contiguous `f64`
19/// values, matching the `#[repr(C)]` layout of `num_complex::Complex`.
20pub trait RankComm: std::fmt::Debug + Send + Sync {
21    /// Index of the calling rank, in `0..size()`.
22    fn rank(&self) -> usize;
23
24    /// Total number of ranks. Always a power of two for the distributed backend.
25    fn size(&self) -> usize;
26
27    /// Concatenate every rank's `local` block in ascending rank order.
28    ///
29    /// The returned vector has length `size() * local.len()` and is identical
30    /// on every rank.
31    fn allgather_c64(&self, local: &[Complex64]) -> Vec<Complex64>;
32
33    /// `f64` version of [`allgather_c64`](RankComm::allgather_c64).
34    fn allgather_f64(&self, local: &[f64]) -> Vec<f64>;
35
36    /// Concatenate blocks of differing length in ascending rank order.
37    ///
38    /// `counts[r]` is the length of rank `r`'s block, so `counts[rank()]`
39    /// equals `local.len()` and every rank passes the same `counts`. The
40    /// returned vector has length `counts.iter().sum()` and is identical on
41    /// every rank.
42    fn allgatherv_u64(&self, local: &[u64], counts: &[usize]) -> Vec<u64>;
43
44    /// `u64` version of [`allgather_c64`](RankComm::allgather_c64).
45    fn allgather_u64(&self, local: &[u64]) -> Vec<u64> {
46        self.allgatherv_u64(local, &vec![local.len(); self.size()])
47    }
48
49    /// Sum a scalar across all ranks; every rank receives the total.
50    fn allreduce_sum_f64(&self, value: f64) -> f64;
51
52    /// Exchange equal length amplitude blocks with `partner`.
53    ///
54    /// On return, `recv` holds `partner`'s `send` block. `send` and `recv` must
55    /// have the same length.
56    fn sendrecv_c64(&self, partner: usize, send: &[Complex64], recv: &mut [Complex64]);
57
58    /// Block until all ranks reach this point.
59    fn barrier(&self);
60}
61
62/// Single rank. All collectives are identity operations.
63#[derive(Debug, Default, Clone, Copy)]
64pub struct SerialComm;
65
66impl RankComm for SerialComm {
67    #[inline]
68    fn rank(&self) -> usize {
69        0
70    }
71
72    #[inline]
73    fn size(&self) -> usize {
74        1
75    }
76
77    #[inline]
78    fn allgather_c64(&self, local: &[Complex64]) -> Vec<Complex64> {
79        local.to_vec()
80    }
81
82    #[inline]
83    fn allgather_f64(&self, local: &[f64]) -> Vec<f64> {
84        local.to_vec()
85    }
86
87    #[inline]
88    fn allgatherv_u64(&self, local: &[u64], _counts: &[usize]) -> Vec<u64> {
89        local.to_vec()
90    }
91
92    #[inline]
93    fn allreduce_sum_f64(&self, value: f64) -> f64 {
94        value
95    }
96
97    #[inline]
98    fn sendrecv_c64(&self, _partner: usize, send: &[Complex64], recv: &mut [Complex64]) {
99        debug_assert_eq!(send.len(), recv.len());
100        recv.copy_from_slice(send);
101    }
102
103    #[inline]
104    fn barrier(&self) {}
105}
106
107/// `Complex64` is two adjacent `f64` values. Assert the layout used by MPI.
108#[cfg(feature = "distributed-mpi")]
109const _: () = assert!(std::mem::size_of::<Complex64>() == 2 * std::mem::size_of::<f64>());
110
111/// Reinterpret `Complex64` as flat `f64` values for MPI calls.
112#[cfg(feature = "distributed-mpi")]
113#[inline]
114fn as_f64(slice: &[Complex64]) -> &[f64] {
115    // SAFETY: Complex64 is repr(C) over two f64 values with no padding.
116    unsafe { std::slice::from_raw_parts(slice.as_ptr() as *const f64, slice.len() * 2) }
117}
118
119#[cfg(feature = "distributed-mpi")]
120#[inline]
121fn as_f64_mut(slice: &mut [Complex64]) -> &mut [f64] {
122    // SAFETY: same layout as `as_f64`; the mutable borrow is exclusive.
123    unsafe { std::slice::from_raw_parts_mut(slice.as_mut_ptr() as *mut f64, slice.len() * 2) }
124}
125
126/// Thread level requested from `MPI_Init_thread`. Rayon workers touch only
127/// memory; every MPI call is made from the thread that constructed the comm.
128#[cfg(feature = "distributed-mpi")]
129const REQUIRED_THREADING: Threading = Threading::Funneled;
130
131/// Reject an MPI whose provided thread level is below [`REQUIRED_THREADING`].
132#[cfg(feature = "distributed-mpi")]
133fn check_threading(provided: Threading) -> Result<()> {
134    if provided >= REQUIRED_THREADING {
135        return Ok(());
136    }
137    Err(PrismError::IncompatibleBackend {
138        backend: "distributed".into(),
139        reason: format!(
140            "MPI thread level {provided:?} is below {REQUIRED_THREADING:?}, the minimum for \
141             running Rayon workers beside MPI"
142        ),
143    })
144}
145
146/// MPI transport over `rsmpi`.
147///
148/// Requires the `distributed-mpi` feature, a system MPI install, and an MPI
149/// launcher. MPI must run at `MPI_THREAD_FUNNELED` or higher: [`MpiComm::world`]
150/// requests that level and both constructors return an error when the provided
151/// level is lower. Keep every MPI call, including dropping the comm, on the
152/// thread that constructed it.
153#[cfg(feature = "distributed-mpi")]
154pub struct MpiComm {
155    /// Held only by a comm that ran `MPI_Init` itself, because dropping a
156    /// `Universe` runs `MPI_Finalize`. A comm attached to an MPI another
157    /// component owns leaves that lifetime alone.
158    _universe: Option<mpi::environment::Universe>,
159    world: mpi::topology::SimpleCommunicator,
160    rank: usize,
161    size: usize,
162}
163
164#[cfg(feature = "distributed-mpi")]
165impl std::fmt::Debug for MpiComm {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        f.debug_struct("MpiComm")
168            .field("rank", &self.rank)
169            .field("size", &self.size)
170            .finish()
171    }
172}
173
174#[cfg(feature = "distributed-mpi")]
175impl MpiComm {
176    /// Initialize MPI at `MPI_THREAD_FUNNELED` and capture the world communicator.
177    ///
178    /// # Errors
179    ///
180    /// Fails when MPI is already initialized (`mpi::initialize_with_threading`
181    /// declines rather than attaching; use [`MpiComm::attach_world`] in a
182    /// process where something else owns MPI) and when the provided thread
183    /// level is below the requested one.
184    pub fn world() -> Result<Self> {
185        let (universe, provided) = mpi::environment::initialize_with_threading(REQUIRED_THREADING)
186            .ok_or_else(|| PrismError::IncompatibleBackend {
187                backend: "distributed".into(),
188                reason: "MPI is already initialized; attach to it with `MpiComm::attach_world`"
189                    .into(),
190            })?;
191        check_threading(provided)?;
192        let world = universe.world();
193        Ok(Self::from_parts(Some(universe), world))
194    }
195
196    /// Attach to an MPI another component has already initialized.
197    ///
198    /// The returned comm holds no `Universe`, so dropping it does not call
199    /// `MPI_Finalize`. That is the point of the constructor: in an interpreter
200    /// where mpi4py calls `MPI_Init_thread` at import and registers
201    /// `MPI_Finalize` at exit, finalizing from a dropped handle would make
202    /// every later MPI call in the process erroneous.
203    ///
204    /// Returns `Ok(None)` when MPI is not initialized, and an error when the
205    /// owner initialized it below `MPI_THREAD_FUNNELED`.
206    pub fn attach_world() -> Result<Option<Self>> {
207        if !mpi::environment::is_initialized() {
208            return Ok(None);
209        }
210        check_threading(mpi::environment::threading_support())?;
211        Ok(Some(Self::from_parts(
212            None,
213            mpi::topology::SimpleCommunicator::world(),
214        )))
215    }
216
217    fn from_parts(
218        universe: Option<mpi::environment::Universe>,
219        world: mpi::topology::SimpleCommunicator,
220    ) -> Self {
221        use mpi::traits::Communicator;
222        let rank = world.rank() as usize;
223        let size = world.size() as usize;
224        Self {
225            _universe: universe,
226            world,
227            rank,
228            size,
229        }
230    }
231}
232
233#[cfg(feature = "distributed-mpi")]
234impl RankComm for MpiComm {
235    fn rank(&self) -> usize {
236        self.rank
237    }
238
239    fn size(&self) -> usize {
240        self.size
241    }
242
243    fn allgather_c64(&self, local: &[Complex64]) -> Vec<Complex64> {
244        use mpi::traits::CommunicatorCollectives;
245        let mut out = vec![Complex64::new(0.0, 0.0); local.len() * self.size];
246        self.world
247            .all_gather_into(as_f64(local), as_f64_mut(&mut out));
248        out
249    }
250
251    fn allgather_f64(&self, local: &[f64]) -> Vec<f64> {
252        use mpi::traits::CommunicatorCollectives;
253        let mut out = vec![0.0_f64; local.len() * self.size];
254        self.world.all_gather_into(local, &mut out);
255        out
256    }
257
258    fn allgatherv_u64(&self, local: &[u64], counts: &[usize]) -> Vec<u64> {
259        use mpi::datatype::PartitionMut;
260        use mpi::traits::CommunicatorCollectives;
261        debug_assert_eq!(counts[self.rank], local.len());
262        let counts: Vec<mpi::Count> = counts.iter().map(|&c| c as mpi::Count).collect();
263        let mut displs = Vec::with_capacity(counts.len());
264        let mut total: mpi::Count = 0;
265        for &c in &counts {
266            displs.push(total);
267            total += c;
268        }
269        let mut out = vec![0_u64; total as usize];
270        self.world
271            .all_gather_varcount_into(local, &mut PartitionMut::new(&mut out[..], counts, displs));
272        out
273    }
274
275    fn allreduce_sum_f64(&self, value: f64) -> f64 {
276        use mpi::traits::CommunicatorCollectives;
277        let mut out = 0.0_f64;
278        self.world
279            .all_reduce_into(&value, &mut out, mpi::collective::SystemOperation::sum());
280        out
281    }
282
283    fn sendrecv_c64(&self, partner: usize, send: &[Complex64], recv: &mut [Complex64]) {
284        use mpi::point_to_point as p2p;
285        use mpi::traits::Communicator;
286        debug_assert_eq!(send.len(), recv.len());
287        let peer = self.world.process_at_rank(partner as i32);
288        p2p::send_receive_into(as_f64(send), &peer, as_f64_mut(recv), &peer);
289    }
290
291    fn barrier(&self) {
292        use mpi::traits::CommunicatorCollectives;
293        self.world.barrier();
294    }
295}
296
297#[cfg(all(test, feature = "distributed-mpi"))]
298mod threading_tests {
299    use super::*;
300
301    // The check leans on the crate's `Ord`, which compares the raw MPI constants,
302    // so the ordering the MPI standard guarantees is pinned alongside it.
303    #[test]
304    fn thread_level_check_rejects_only_single() {
305        assert!(Threading::Single < Threading::Funneled);
306        assert!(Threading::Funneled < Threading::Serialized);
307        assert!(Threading::Serialized < Threading::Multiple);
308        assert!(check_threading(Threading::Single).is_err());
309        for level in [
310            Threading::Funneled,
311            Threading::Serialized,
312            Threading::Multiple,
313        ] {
314            assert!(check_threading(level).is_ok(), "{level:?}");
315        }
316    }
317}
prism_q/distributed/comm.rs

prism_q/distributed/
comm.rs
