prism_q/distributed/
mod.rs

1//! Distributed context, transport, and thresholds.
2//!
3//! [`DistributedContext`] wraps a shared [`RankComm`] transport. Backend modules
4//! define their own state partitioning and use this module for rank access.
5//!
6//! Available contexts:
7//! - [`DistributedContext::serial`]: one rank for tests and runs without MPI.
8//! - `DistributedContext::world`: the MPI world communicator (requires the
9//!   `distributed-mpi` feature and an MPI launcher).
10//!
11//! Tuning thresholds are cached after the first environment variable read.
12
13pub mod comm;
14#[cfg(any(test, feature = "bench-internal"))]
15pub mod loopback;
16
17use std::sync::Arc;
18
19pub use comm::{RankComm, SerialComm};
20
21#[cfg(feature = "distributed-mpi")]
22pub use comm::MpiComm;
23
24/// Default minimum local qubit count below which distribution is not worthwhile.
25///
26/// Small slices per rank spend more time in communication than computation.
27pub const MIN_LOCAL_QUBITS_DEFAULT: usize = 10;
28
29/// Minimum local qubits per rank, tunable via `PRISM_DIST_MIN_LOCAL_QUBITS`.
30///
31/// An unparseable or out-of-range value warns on stderr and uses the default.
32pub fn min_local_qubits() -> usize {
33    use std::sync::OnceLock;
34    static CACHED: OnceLock<usize> = OnceLock::new();
35    *CACHED.get_or_init(|| {
36        parse_usize_knob(
37            "PRISM_DIST_MIN_LOCAL_QUBITS",
38            std::env::var("PRISM_DIST_MIN_LOCAL_QUBITS").ok(),
39            MIN_LOCAL_QUBITS_DEFAULT,
40            1,
41        )
42    })
43}
44
45/// Maximum number of amplitudes exchanged per message on the direct exchange
46/// paths of the distributed backend. Chunking bounds the transfer buffers to
47/// this value.
48///
49/// Tunable via `PRISM_DIST_EXCHANGE_CHUNK`. The default (`usize::MAX`) keeps the
50/// original one message behavior, so there is no change unless set.
51pub const EXCHANGE_CHUNK_DEFAULT: usize = usize::MAX;
52
53/// Chunk size in amplitudes for the tiled rank exchanges.
54///
55/// An unparseable or out-of-range value warns on stderr and uses the default.
56pub fn exchange_chunk() -> usize {
57    use std::sync::OnceLock;
58    static CACHED: OnceLock<usize> = OnceLock::new();
59    *CACHED.get_or_init(|| {
60        parse_usize_knob(
61            "PRISM_DIST_EXCHANGE_CHUNK",
62            std::env::var("PRISM_DIST_EXCHANGE_CHUNK").ok(),
63            EXCHANGE_CHUNK_DEFAULT,
64            1,
65        )
66    })
67}
68
69/// Whether the distributed backend relabels qubits to keep busy qubits local.
70///
71/// On by default. Relabeling turns SWAP gates into zero-communication map
72/// updates and moves global qubits into local positions with a half-slice
73/// exchange before non-diagonal gates touch them. Set `PRISM_DIST_RELABEL=0`
74/// to disable and force direct per-gate exchange.
75///
76/// Accepts `0`/`false` and `1`/`true`; anything else warns on stderr and uses
77/// the default.
78pub fn relabel_enabled() -> bool {
79    use std::sync::OnceLock;
80    static CACHED: OnceLock<bool> = OnceLock::new();
81    *CACHED.get_or_init(|| {
82        parse_bool_knob(
83            "PRISM_DIST_RELABEL",
84            std::env::var("PRISM_DIST_RELABEL").ok(),
85            true,
86        )
87    })
88}
89
90/// Read a count knob, warning on stderr and falling back to `default` when the
91/// value does not parse or falls below `min`.
92///
93/// Warning rather than erroring is the knob contract: every reader is a cached
94/// initializer on an infallible path, so an invalid value must not take down a
95/// run that would otherwise be correct with the default.
96fn parse_usize_knob(var: &str, raw: Option<String>, default: usize, min: usize) -> usize {
97    let Some(raw) = raw else {
98        return default;
99    };
100    match raw.trim().parse::<usize>() {
101        Ok(n) if n >= min => n,
102        Ok(n) => {
103            eprintln!("warning: {var}={n} is below the minimum of {min}; using {default}.");
104            default
105        }
106        Err(_) => {
107            eprintln!("warning: {var}={raw:?} is not a count; using {default}.");
108            default
109        }
110    }
111}
112
113/// Read a flag knob. See [`parse_usize_knob`] for why an invalid value warns.
114fn parse_bool_knob(var: &str, raw: Option<String>, default: bool) -> bool {
115    let Some(raw) = raw else {
116        return default;
117    };
118    match raw.trim() {
119        "0" => false,
120        "1" => true,
121        other if other.eq_ignore_ascii_case("false") => false,
122        other if other.eq_ignore_ascii_case("true") => true,
123        other => {
124            eprintln!("warning: {var}={other:?} is not a flag; using {default}.");
125            default
126        }
127    }
128}
129
130/// Shared handle to a rank transport for distributed simulation.
131pub struct DistributedContext {
132    comm: Arc<dyn RankComm>,
133}
134
135impl std::fmt::Debug for DistributedContext {
136    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
137        f.debug_struct("DistributedContext")
138            .field("rank", &self.rank())
139            .field("size", &self.size())
140            .finish()
141    }
142}
143
144impl DistributedContext {
145    /// Build a context from any [`RankComm`] implementation.
146    pub fn from_comm(comm: Arc<dyn RankComm>) -> Arc<Self> {
147        Arc::new(Self { comm })
148    }
149
150    /// Single rank. Used by tests and runs without MPI.
151    pub fn serial() -> Arc<Self> {
152        Self::from_comm(Arc::new(SerialComm))
153    }
154
155    /// Initialize MPI and capture the world communicator. See [`MpiComm::world`]
156    /// for the thread level requested and the failure conditions.
157    #[cfg(feature = "distributed-mpi")]
158    pub fn world() -> crate::error::Result<Arc<Self>> {
159        MpiComm::world().map(|c| Self::from_comm(Arc::new(c)))
160    }
161
162    /// Attach to an MPI another component has already initialized, taking no
163    /// ownership of its lifetime. See [`MpiComm::attach_world`] for why an
164    /// embedding interpreter needs this rather than [`DistributedContext::world`].
165    ///
166    /// Returns `Ok(None)` when MPI is not initialized, and an error when it
167    /// runs below the thread level [`MpiComm`] requires.
168    #[cfg(feature = "distributed-mpi")]
169    pub fn attached_world() -> crate::error::Result<Option<Arc<Self>>> {
170        Ok(MpiComm::attach_world()?.map(|c| Self::from_comm(Arc::new(c))))
171    }
172
173    /// Index of the calling rank.
174    pub fn rank(&self) -> usize {
175        self.comm.rank()
176    }
177
178    /// Total number of ranks.
179    pub fn size(&self) -> usize {
180        self.comm.size()
181    }
182
183    pub(crate) fn comm(&self) -> &Arc<dyn RankComm> {
184        &self.comm
185    }
186}
187
188#[cfg(test)]
189mod knob_tests {
190    use super::*;
191
192    // One case per parse site. The knob functions cache per process, so the
193    // parsers are exercised directly rather than through the environment.
194    #[test]
195    fn min_local_qubits_knob_rejects_invalid_values() {
196        let d = MIN_LOCAL_QUBITS_DEFAULT;
197        let parse =
198            |raw: &str| parse_usize_knob("PRISM_DIST_MIN_LOCAL_QUBITS", Some(raw.into()), d, 1);
199        assert_eq!(parse("4"), 4);
200        assert_eq!(parse(" 4 "), 4);
201        assert_eq!(parse("abc"), d, "unparseable falls back");
202        assert_eq!(parse("-1"), d, "negative falls back");
203        assert_eq!(parse("0"), d, "below the minimum falls back");
204        assert_eq!(
205            parse_usize_knob("PRISM_DIST_MIN_LOCAL_QUBITS", None, d, 1),
206            d
207        );
208    }
209
210    #[test]
211    fn exchange_chunk_knob_rejects_invalid_values() {
212        let d = EXCHANGE_CHUNK_DEFAULT;
213        let parse =
214            |raw: &str| parse_usize_knob("PRISM_DIST_EXCHANGE_CHUNK", Some(raw.into()), d, 1);
215        assert_eq!(parse("4096"), 4096);
216        assert_eq!(parse("4k"), d, "unparseable falls back");
217        assert_eq!(parse("0"), d, "below the minimum falls back");
218        assert_eq!(parse_usize_knob("PRISM_DIST_EXCHANGE_CHUNK", None, d, 1), d);
219    }
220
221    #[test]
222    fn relabel_knob_rejects_invalid_values() {
223        let parse = |raw: &str| parse_bool_knob("PRISM_DIST_RELABEL", Some(raw.into()), true);
224        assert!(!parse("0"));
225        assert!(!parse("false"));
226        assert!(!parse("FALSE"));
227        assert!(parse("1"));
228        assert!(parse("true"));
229        assert!(parse("yes"), "unrecognized falls back to the default");
230        assert!(parse(""), "empty falls back to the default");
231        assert!(parse_bool_knob("PRISM_DIST_RELABEL", None, true));
232    }
233}
prism_q/distributed/mod.rs

prism_q/distributed/
mod.rs
