prism_q/distributed/
mod.rs

1//! Distributed context, rank transport, and tuning thresholds.
2//!
3//! Each threshold reads its environment variable once and caches the result.
4
5pub mod comm;
6#[cfg(any(test, feature = "bench-internal"))]
7pub mod loopback;
8
9use std::sync::Arc;
10
11pub use comm::{RankComm, SerialComm};
12
13#[cfg(feature = "distributed-mpi")]
14pub use comm::MpiComm;
15
16use crate::env_knobs::parse_usize_knob;
17
18/// Default for `PRISM_DIST_MIN_LOCAL_QUBITS`. Smaller slices per rank spend more time
19/// in communication than computation.
20pub const MIN_LOCAL_QUBITS_DEFAULT: usize = 10;
21
22/// Minimum local qubits per rank, tunable via `PRISM_DIST_MIN_LOCAL_QUBITS`.
23///
24/// The distributed backend rejects a register split that leaves fewer. An unparseable
25/// or out-of-range value warns on stderr and uses the default.
26pub fn min_local_qubits() -> usize {
27    use std::sync::OnceLock;
28    static CACHED: OnceLock<usize> = OnceLock::new();
29    *CACHED.get_or_init(|| {
30        parse_usize_knob(
31            "PRISM_DIST_MIN_LOCAL_QUBITS",
32            std::env::var("PRISM_DIST_MIN_LOCAL_QUBITS").ok(),
33            MIN_LOCAL_QUBITS_DEFAULT,
34            1,
35        )
36    })
37}
38
39/// Default for `PRISM_DIST_EXCHANGE_CHUNK`: each rank exchange goes as one message.
40pub const EXCHANGE_CHUNK_DEFAULT: usize = usize::MAX;
41
42/// Largest rank exchange message in amplitudes, tunable via `PRISM_DIST_EXCHANGE_CHUNK`.
43///
44/// Tiling the direct and relabel exchanges by this bounds the transfer buffers.
45/// An unparseable or out-of-range value warns on stderr and uses the default.
46pub fn exchange_chunk() -> usize {
47    use std::sync::OnceLock;
48    static CACHED: OnceLock<usize> = OnceLock::new();
49    *CACHED.get_or_init(|| {
50        parse_usize_knob(
51            "PRISM_DIST_EXCHANGE_CHUNK",
52            std::env::var("PRISM_DIST_EXCHANGE_CHUNK").ok(),
53            EXCHANGE_CHUNK_DEFAULT,
54            1,
55        )
56    })
57}
58
59/// Whether the distributed backend relabels qubits to keep busy qubits local.
60///
61/// On by default. Relabeling turns SWAP gates into zero-communication map
62/// updates and moves global qubits into local positions with a half-slice
63/// exchange before non-diagonal gates touch them. Set `PRISM_DIST_RELABEL=0`
64/// to disable and force direct per-gate exchange.
65///
66/// Accepts `0`/`false` and `1`/`true`; anything else warns on stderr and uses
67/// the default.
68pub fn relabel_enabled() -> bool {
69    use std::sync::OnceLock;
70    static CACHED: OnceLock<bool> = OnceLock::new();
71    *CACHED.get_or_init(|| {
72        parse_bool_knob(
73            "PRISM_DIST_RELABEL",
74            std::env::var("PRISM_DIST_RELABEL").ok(),
75            true,
76        )
77    })
78}
79
80/// Read a flag knob: `0`/`false` and `1`/`true` in any case; anything else
81/// warns on stderr and yields `default`, for the reason `env_knobs` gives.
82fn parse_bool_knob(var: &str, raw: Option<String>, default: bool) -> bool {
83    let Some(raw) = raw else {
84        return default;
85    };
86    match raw.trim() {
87        "0" => false,
88        "1" => true,
89        other if other.eq_ignore_ascii_case("false") => false,
90        other if other.eq_ignore_ascii_case("true") => true,
91        other => {
92            eprintln!("warning: {var}={other:?} is not a flag; using {default}.");
93            default
94        }
95    }
96}
97
98/// Shared handle to a rank transport for distributed simulation.
99pub struct DistributedContext {
100    comm: Arc<dyn RankComm>,
101}
102
103impl std::fmt::Debug for DistributedContext {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        f.debug_struct("DistributedContext")
106            .field("rank", &self.rank())
107            .field("size", &self.size())
108            .finish()
109    }
110}
111
112impl DistributedContext {
113    pub fn from_comm(comm: Arc<dyn RankComm>) -> Arc<Self> {
114        Arc::new(Self { comm })
115    }
116
117    /// Single rank. Used by tests and runs without MPI.
118    pub fn serial() -> Arc<Self> {
119        Self::from_comm(Arc::new(SerialComm))
120    }
121
122    /// Initialize MPI and capture the world communicator. See [`MpiComm::world`]
123    /// for the thread level requested and the failure conditions.
124    #[cfg(feature = "distributed-mpi")]
125    pub fn world() -> crate::error::Result<Arc<Self>> {
126        MpiComm::world().map(|c| Self::from_comm(Arc::new(c)))
127    }
128
129    /// Attach to an MPI another component has already initialized, taking no
130    /// ownership of its lifetime. See [`MpiComm::attach_world`] for why an
131    /// embedding interpreter needs this rather than [`DistributedContext::world`].
132    ///
133    /// Returns `Ok(None)` when MPI is not initialized, and an error when it
134    /// runs below the thread level [`MpiComm`] requires.
135    #[cfg(feature = "distributed-mpi")]
136    pub fn attached_world() -> crate::error::Result<Option<Arc<Self>>> {
137        Ok(MpiComm::attach_world()?.map(|c| Self::from_comm(Arc::new(c))))
138    }
139
140    pub fn rank(&self) -> usize {
141        self.comm.rank()
142    }
143
144    pub fn size(&self) -> usize {
145        self.comm.size()
146    }
147
148    pub(crate) fn comm(&self) -> &Arc<dyn RankComm> {
149        &self.comm
150    }
151}
152
153#[cfg(test)]
154mod knob_tests {
155    use super::*;
156
157    // One case per parse site. The knob functions cache per process, so the
158    // parsers are exercised directly rather than through the environment.
159    #[test]
160    fn min_local_qubits_knob_rejects_invalid_values() {
161        let d = MIN_LOCAL_QUBITS_DEFAULT;
162        let parse =
163            |raw: &str| parse_usize_knob("PRISM_DIST_MIN_LOCAL_QUBITS", Some(raw.into()), d, 1);
164        assert_eq!(parse("4"), 4);
165        assert_eq!(parse(" 4 "), 4);
166        assert_eq!(parse("abc"), d, "unparseable falls back");
167        assert_eq!(parse("-1"), d, "negative falls back");
168        assert_eq!(parse("0"), d, "below the minimum falls back");
169        assert_eq!(
170            parse_usize_knob("PRISM_DIST_MIN_LOCAL_QUBITS", None, d, 1),
171            d
172        );
173    }
174
175    #[test]
176    fn exchange_chunk_knob_rejects_invalid_values() {
177        let d = EXCHANGE_CHUNK_DEFAULT;
178        let parse =
179            |raw: &str| parse_usize_knob("PRISM_DIST_EXCHANGE_CHUNK", Some(raw.into()), d, 1);
180        assert_eq!(parse("4096"), 4096);
181        assert_eq!(parse("4k"), d, "unparseable falls back");
182        assert_eq!(parse("0"), d, "below the minimum falls back");
183        assert_eq!(parse_usize_knob("PRISM_DIST_EXCHANGE_CHUNK", None, d, 1), d);
184    }
185
186    #[test]
187    fn relabel_knob_rejects_invalid_values() {
188        let parse = |raw: &str| parse_bool_knob("PRISM_DIST_RELABEL", Some(raw.into()), true);
189        assert!(!parse("0"));
190        assert!(!parse("false"));
191        assert!(!parse("FALSE"));
192        assert!(parse("1"));
193        assert!(parse("true"));
194        assert!(parse("yes"), "unrecognized falls back to the default");
195        assert!(parse(""), "empty falls back to the default");
196        assert!(parse_bool_knob("PRISM_DIST_RELABEL", None, true));
197    }
198}
prism_q/distributed/mod.rs

prism_q/distributed/
mod.rs
