paulistrings/engine/partitioned/
distributed.rs

1//! [`DistributedSum`], the partitioned driver with one partition per process, and its scatter options.
2
3use std::sync::Arc;
4use std::time::Instant;
5
6use num_complex::Complex64;
7
8#[cfg(feature = "phase-timing")]
9use super::driver::DriverLaps;
10use super::driver::{scatter_local, DriverFrame, LOG_TARGET};
11use super::runtime::PartitionRuntime;
12use super::topology::{PartitionConfig, TopologyError};
13use crate::circuit::Circuit;
14use crate::engine::backend::{HostPartition, PartitionBackend, PartitionStorage};
15use crate::engine::timing::timed;
16use crate::engine::trace::PartitionTrace;
17use crate::engine::transport::Transport;
18use crate::engine::{progress_scope, Direction, PropagateOptions};
19use crate::pauli_string::qubit_mask;
20use crate::pauli_sum::hash::PartitionRows;
21use crate::pauli_sum::PauliSum;
22use crate::truncation::TruncationPolicy;
23
24#[cfg(feature = "phase-timing")]
25use crate::engine::stats::PartitionPhaseStats;
26
27/// Parts the gather ships per rank: bucket lengths, then the three columns.
28const GATHER_PARTS: usize = 4;
29
30/// Which partition rows a scatter splits by: a seeded GF(2)-random draw, or a qubit cut that keeps a channel local when its qubits share a block.
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub enum PartitionRowPolicy {
33    /// [`PartitionRows::from_seed`] with this seed, or the sum's own hash seed at `None`; the default.
34    Seeded(Option<u64>),
35    /// [`PartitionRows::cut`]: one disjoint qubit block per rank, in rank order; the rows are z-only.
36    Cut(Vec<Vec<usize>>),
37    /// [`PartitionRows::from_seed_excluding`]: a seeded draw whose rows read neither the x-bits of `exclude_x` nor the z-bits of `exclude_z`, so keys differing only there share a partition.
38    SeededExcluding {
39        /// As in [`Seeded`](Self::Seeded).
40        seed: Option<u64>,
41        /// Qubits whose x-bit no row reads.
42        exclude_x: Vec<usize>,
43        /// Qubits whose z-bit no row reads.
44        exclude_z: Vec<usize>,
45    },
46}
47
48impl PartitionRowPolicy {
49    /// The rows for `2^bits` partitions of a `num_qubits` register, with `default_seed` standing in for a `None` seed.
50    ///
51    /// # Panics
52    ///
53    /// As [`PartitionRows::cut`] or [`PartitionRows::from_seed_excluding`], and if an excluded qubit is out of range.
54    pub fn rows<const W: usize>(
55        &self,
56        num_qubits: usize,
57        bits: u8,
58        default_seed: u64,
59    ) -> PartitionRows<W> {
60        let mask = |qubits: &[usize]| qubit_mask(qubits.iter().copied(), num_qubits);
61        match self {
62            Self::Seeded(seed) => {
63                PartitionRows::from_seed(num_qubits, bits, seed.unwrap_or(default_seed))
64            }
65            Self::Cut(blocks) => PartitionRows::cut(num_qubits, blocks),
66            Self::SeededExcluding { seed, exclude_x, exclude_z } => {
67                PartitionRows::from_seed_excluding(
68                    num_qubits,
69                    bits,
70                    seed.unwrap_or(default_seed),
71                    &mask(exclude_x),
72                    &mask(exclude_z),
73                )
74            }
75        }
76    }
77}
78
79/// How a `scatter_with` runs its partitions and picks the rows the sum splits by.
80#[derive(Clone)]
81pub struct ScatterOptions<const W: usize> {
82    /// The runtime the partitions run on, reusable across sums; one partition for a [`DistributedSum`].
83    pub runtime: Arc<PartitionRuntime>,
84    /// The partition rows the group splits by.
85    pub rows: ScatterRows<W>,
86}
87
88/// The partition rows a `scatter_with` splits by.
89#[derive(Clone, Debug, PartialEq, Eq)]
90pub enum ScatterRows<const W: usize> {
91    /// The rows a policy draws for the partition count, falling back to the sum's own hash seed.
92    Policy(PartitionRowPolicy),
93    /// Rows the caller built; in a distributed group every rank must pass the same ones.
94    Explicit(PartitionRows<W>),
95}
96
97impl<const W: usize> ScatterRows<W> {
98    /// The rows for `2^bits` partitions of `sum`.
99    pub(crate) fn resolve(self, sum: &PauliSum<W>, bits: u8) -> PartitionRows<W> {
100        match self {
101            Self::Policy(policy) => policy.rows(sum.num_qubits(), bits, sum.hash().seed()),
102            Self::Explicit(rows) => rows,
103        }
104    }
105}
106
107/// `log2(size)`, the partition bits of a group; panics unless `size` is a power of two.
108pub(crate) fn group_bits(size: u32) -> u8 {
109    assert!(
110        size.is_power_of_two(),
111        "a group of {size} ranks cannot be a partitioning: a partition is named by log2(P) \
112         GF(2) rows, so the rank count must be a power of two",
113    );
114    size.trailing_zeros() as u8
115}
116
117/// One process's partition of a sum split across a [`Transport`]'s group, held across calls so a driver scatters once, steps many times and gathers once.
118///
119/// The input is replicated: every rank scatters the same sum and keeps its own partition.
120/// [`gather`](Self::gather) returns `Some` on rank 0 only, and a [`PartitionTrace`] holds this rank's counts only.
121/// `B` is where the partition lives: host memory, or a CUDA device under the `cuda` feature (`gpu::GpuDistributedSum`).
122pub struct DistributedSum<const W: usize, X: Transport, B = HostPartition<W>> {
123    local: B,
124    rows: PartitionRows<W>,
125    runtime: Arc<PartitionRuntime>,
126    transport: X,
127    trace: Option<PartitionTrace>,
128    #[cfg(feature = "phase-timing")]
129    laps: DriverLaps,
130}
131
132impl<const W: usize, X: Transport, B> DistributedSum<W, X, B> {
133    /// This rank's index in the group.
134    pub fn rank(&self) -> u32 {
135        self.transport.rank()
136    }
137
138    /// Ranks in the group, which is also the partition count.
139    pub fn size(&self) -> u32 {
140        self.transport.size()
141    }
142
143    /// This rank's endpoint, for a collective of the caller's own.
144    pub fn transport(&self) -> &X {
145        &self.transport
146    }
147
148    /// The rows deciding which rank a key belongs to.
149    pub fn rows(&self) -> &PartitionRows<W> {
150        &self.rows
151    }
152
153    /// The runtime this rank's work runs on, for handing to another [`DistributedSum`].
154    pub fn runtime(&self) -> &Arc<PartitionRuntime> {
155        &self.runtime
156    }
157
158    /// Start recording a [`PartitionTrace`] of this rank's layers on every later propagation; idempotent.
159    pub fn enable_trace(&mut self) {
160        self.trace.get_or_insert_with(PartitionTrace::default);
161    }
162
163    /// Drain this rank's per-layer records, or `None` if tracing was never enabled; tracing stays enabled.
164    ///
165    /// `terms_in`, `terms_out` and `rows_received` have one entry, while `rows_sent[0]` and `bytes_sent[0]` are indexed by destination rank.
166    pub fn take_trace(&mut self) -> Option<PartitionTrace> {
167        self.trace.as_mut().map(std::mem::take)
168    }
169
170    /// A split around a partition another backend already scattered, with `scatter_ns` its scatter time.
171    #[cfg(feature = "cuda")]
172    pub(crate) fn from_backend(
173        local: B,
174        rows: PartitionRows<W>,
175        runtime: Arc<PartitionRuntime>,
176        transport: X,
177        scatter_ns: u64,
178    ) -> Self {
179        #[cfg(not(feature = "phase-timing"))]
180        let _ = scatter_ns;
181        Self {
182            local,
183            rows,
184            runtime,
185            transport,
186            trace: None,
187            #[cfg(feature = "phase-timing")]
188            laps: DriverLaps::new(scatter_ns),
189        }
190    }
191
192    #[cfg(feature = "cuda")]
193    pub(crate) fn backend(&self) -> &B {
194        &self.local
195    }
196
197    #[cfg(feature = "cuda")]
198    pub(crate) fn backend_mut(&mut self) -> &mut B {
199        &mut self.local
200    }
201
202    /// Drain scatter time, gather time and layers driven.
203    #[cfg(all(feature = "cuda", feature = "phase-timing"))]
204    pub(crate) fn take_driver_laps(&mut self) -> (u64, u64, u64) {
205        self.laps.take()
206    }
207
208    #[cfg(feature = "phase-timing")]
209    pub(crate) fn lap_gather(&self, ns: u64) {
210        self.laps.lap_gather(ns);
211    }
212
213    /// [`DistributedSum::propagate_with`] on any backend.
214    pub(crate) fn propagate_on_backend<T>(
215        &mut self,
216        circuit: &Circuit<W>,
217        policy: &T,
218        direction: Direction,
219        options: PropagateOptions,
220    ) where
221        T: TruncationPolicy<W> + ?Sized,
222        B: PartitionBackend<W, T>,
223    {
224        let runtime = &*self.runtime;
225        let transport = &self.transport;
226        let rank = transport.rank() as usize;
227        let fingerprint =
228            run_fingerprint(circuit.channels.len(), direction, options, self.rows.num_qubits(), W);
229        let frame = DriverFrame {
230            parts: std::slice::from_mut(&mut self.local),
231            rows: &self.rows,
232            runtime,
233            trace: &mut self.trace,
234            #[cfg(feature = "phase-timing")]
235            laps: &mut self.laps,
236            rank: Some(rank),
237            size: transport.size() as usize,
238        };
239        frame.propagate(
240            circuit,
241            policy,
242            direction,
243            options,
244            || transport.check_consistency(fingerprint),
245            |mut items, layers| {
246                let work = &mut items[0];
247                runtime.install(move || layers.run(rank, work, transport));
248                items
249            },
250        );
251    }
252}
253
254impl<const W: usize, X: Transport, B: PartitionStorage<W>> DistributedSum<W, X, B> {
255    /// Terms in the whole sum. **Collective.**
256    pub fn len(&self) -> usize {
257        let mut total = [self.local.len() as u64];
258        self.transport.allreduce_sum_u64(&mut total);
259        total[0] as usize
260    }
261
262    /// Whether the whole sum is empty. **Collective**, via [`len`](Self::len).
263    pub fn is_empty(&self) -> bool {
264        self.len() == 0
265    }
266
267    /// The bucket bits, equal on every rank.
268    pub fn bits(&self) -> u8 {
269        self.local.hash().bits()
270    }
271
272    /// Qubits the sum is over.
273    pub fn num_qubits(&self) -> usize {
274        self.rows.num_qubits()
275    }
276}
277
278impl<const W: usize, X: Transport> DistributedSum<W, X> {
279    /// Split the replicated `sum` across `transport`'s group by the default rows ([`PartitionRowPolicy::Seeded`] at the sum's own hash seed), keeping this rank's share on the one-partition runtime `config` describes.
280    ///
281    /// # Errors
282    ///
283    /// [`TopologyError`] if `config` cannot be resolved or the pool cannot be built.
284    ///
285    /// # Panics
286    ///
287    /// If `config` asks for more than one partition, or if the group size is not a power of two.
288    pub fn scatter(
289        sum: PauliSum<W>,
290        transport: X,
291        config: &PartitionConfig,
292    ) -> Result<Self, TopologyError> {
293        let options = ScatterOptions {
294            runtime: PartitionRuntime::new(config)?,
295            rows: ScatterRows::Policy(PartitionRowPolicy::Seeded(None)),
296        };
297        Ok(Self::scatter_with(sum, transport, options))
298    }
299
300    /// [`scatter`](Self::scatter) onto a caller-built runtime, split by the rows `options` names.
301    ///
302    /// [`ScatterRows::Explicit`] rows must be the same on every rank; nothing checks it, and a disagreement misroutes an exchange.
303    ///
304    /// # Panics
305    ///
306    /// If the runtime has more than one partition, if the group size is not a power of two, if the rows do not name one partition per rank or are for a different qubit count than `sum`, or as [`PartitionRows::cut`].
307    /// In debug builds, if the rows are not independent of the sum's hash rows.
308    pub fn scatter_with(sum: PauliSum<W>, transport: X, options: ScatterOptions<W>) -> Self {
309        let ScatterOptions { runtime, rows } = options;
310        let rows = rows.resolve(&sum, group_bits(transport.size()));
311        assert_eq!(
312            runtime.num_partitions(),
313            1,
314            "a distributed run holds one partition per process, but the runtime has {}; \
315             domains-per-rank hybrids are not implemented",
316            runtime.num_partitions(),
317        );
318        rows.assert_splits(sum.hash(), sum.num_qubits(), 1usize << group_bits(transport.size()));
319
320        let started = Instant::now();
321        let local = {
322            let sum = &sum;
323            let rows = &rows;
324            let transport = &transport;
325            runtime.install(move || scatter_local(sum, rows, transport.rank(), transport))
326        };
327        log::info!(
328            target: LOG_TARGET,
329            "scatter: {} terms -> {} held, {} bucket bits, {:.3} s, {}",
330            sum.len(),
331            local.len(),
332            local.hash().bits(),
333            started.elapsed().as_secs_f64(),
334            progress_scope(
335                Some(transport.rank() as usize),
336                transport.size() as usize,
337                Some(&runtime.placement_summary()),
338            ),
339        );
340
341        Self {
342            local: HostPartition::new(local),
343            rows,
344            runtime,
345            transport,
346            trace: None,
347            #[cfg(feature = "phase-timing")]
348            laps: DriverLaps::new(started.elapsed().as_nanos() as u64),
349        }
350    }
351
352    /// [`propagate_with`](Self::propagate_with) under [`PropagateOptions::default()`].
353    pub fn propagate<T>(&mut self, circuit: &Circuit<W>, policy: &T, direction: Direction)
354    where
355        T: TruncationPolicy<W> + ?Sized,
356    {
357        self.propagate_with(circuit, policy, direction, PropagateOptions::default())
358    }
359
360    /// Propagate through `circuit` under `policy`, in place. **Collective**: every rank passes the same circuit, direction and options.
361    ///
362    /// A shape fingerprint of the run (channel count, direction, options, qubit count, `W`) is checked across ranks first; channel contents are not.
363    ///
364    /// # Panics
365    ///
366    /// If `options` fail [`PropagateOptions::validate`], if a channel's preparation declines, or if the ranks disagree about the run's shape.
367    /// Before the first layer, above one rank, unless [`policy.supports_partitioned()`](TruncationPolicy::supports_partitioned) — exact [`TopN`](crate::TopN) is not yet supported partitioned.
368    pub fn propagate_with<T>(
369        &mut self,
370        circuit: &Circuit<W>,
371        policy: &T,
372        direction: Direction,
373        options: PropagateOptions,
374    ) where
375        T: TruncationPolicy<W> + ?Sized,
376    {
377        self.propagate_on_backend(circuit, policy, direction, options);
378    }
379
380    /// Collect the whole sum on rank 0, `None` elsewhere, leaving `self` intact. **Collective.**
381    pub fn gather(&self) -> Option<PauliSum<W>> {
382        timed!(let started = Instant::now());
383        let out = gather_share(&self.local.sum, &self.transport);
384        timed!(self.lap_gather(started.elapsed().as_nanos() as u64));
385        out
386    }
387
388    /// This rank's share, a valid [`PauliSum`] under the group's shared hash.
389    pub fn local(&self) -> &PauliSum<W> {
390        &self.local.sum
391    }
392
393    /// Drain this rank's phase counters since the last drain, zeroing them; `per_partition` has one entry.
394    #[cfg(feature = "phase-timing")]
395    pub fn take_stats(&mut self) -> PartitionPhaseStats {
396        let (scatter_ns, gather_ns, layers) = self.laps.take();
397        PartitionPhaseStats {
398            per_partition: vec![self.local.state.layer.take_phase_stats()],
399            scatter_ns,
400            gather_ns,
401            layers,
402        }
403    }
404
405    /// Panic unless this rank's share is well-formed and holds only its own keys; `O(terms)`, local only.
406    pub fn assert_invariants(&self) {
407        #[cfg(any(test, debug_assertions))]
408        self.local.sum.assert_invariants();
409        let held = self.local.sum.partition_rank_of_all(&self.rows);
410        assert!(
411            held.is_none() || held == Some(self.rank()),
412            "rank {} holds keys of partition {held:?}",
413            self.rank(),
414        );
415    }
416}
417
418/// [`DistributedSum::gather`] over any rank's host share. **Collective.**
419pub(crate) fn gather_share<const W: usize, X: Transport>(
420    local: &PauliSum<W>,
421    transport: &X,
422) -> Option<PauliSum<W>> {
423    let lens: Vec<u64> = (0..local.num_buckets()).map(|b| local.bucket_len(b) as u64).collect();
424    let (x, z, coeff) = local.to_arrays();
425    let parts: Vec<&[u8]> = vec![
426        bytemuck::cast_slice(&lens),
427        bytemuck::cast_slice(x.as_flattened()),
428        bytemuck::cast_slice(z.as_flattened()),
429        bytemuck::cast_slice(&coeff),
430    ];
431    let all = transport.gather_to_root(parts);
432    drop((x, z, coeff));
433
434    all.map(|all| {
435        let hash = local.hash().clone();
436        let num_qubits = local.num_qubits();
437        let parts: Vec<PauliSum<W>> = all
438            .into_iter()
439            .enumerate()
440            .map(|(rank, parts)| decode_rank(&parts, hash.clone(), num_qubits, rank))
441            .collect();
442        PauliSum::merge_partitions(parts)
443    })
444}
445
446/// A shape-only fingerprint of the run, since a `Circuit` of trait objects has no canonical encoding.
447fn run_fingerprint(
448    channels: usize,
449    direction: Direction,
450    options: PropagateOptions,
451    num_qubits: usize,
452    w: usize,
453) -> u64 {
454    // FNV-1a over the fields, in a fixed order.
455    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
456    let mut mix = |v: u64| {
457        hash ^= v;
458        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
459    };
460    mix(channels as u64);
461    mix(matches!(direction, Direction::Heisenberg) as u64);
462    mix(options.target_bucket_len as u64);
463    mix(options.min_buckets as u64);
464    mix(num_qubits as u64);
465    mix(w as u64);
466    hash
467}
468
469/// Copy `len` values of `T` out of a possibly unaligned byte view.
470fn decode_column<T: bytemuck::Pod>(bytes: &[u8], len: usize, what: &str) -> Vec<T> {
471    let stride = std::mem::size_of::<T>();
472    assert_eq!(
473        bytes.len(),
474        len * stride,
475        "gather {what}: expected {} bytes for {len} entries, got {}",
476        len * stride,
477        bytes.len(),
478    );
479    bytes.chunks_exact(stride).map(bytemuck::pod_read_unaligned).collect()
480}
481
482/// [`decode_column`] for `[u64; W]` rows, which are not `Pod` at a generic `W`.
483fn decode_rows<const W: usize>(bytes: &[u8], rows: usize, what: &str) -> Vec<[u64; W]> {
484    let stride = W * std::mem::size_of::<u64>();
485    assert_eq!(
486        bytes.len(),
487        rows * stride,
488        "gather {what}: expected {} bytes for {rows} rows, got {}",
489        rows * stride,
490        bytes.len(),
491    );
492    bytes
493        .chunks_exact(stride)
494        .map(|row| std::array::from_fn(|i| bytemuck::pod_read_unaligned(&row[i * 8..i * 8 + 8])))
495        .collect()
496}
497
498/// Rebuild one rank's share from the [`GATHER_PARTS`] parts it shipped; panics on contradictory lengths.
499fn decode_rank<const W: usize>(
500    parts: &[Vec<u8>],
501    hash: crate::pauli_sum::hash::Gf2Hash<W>,
502    num_qubits: usize,
503    rank: usize,
504) -> PauliSum<W> {
505    assert_eq!(
506        parts.len(),
507        GATHER_PARTS,
508        "gather: rank {rank} sent {} parts, expected {GATHER_PARTS}",
509        parts.len(),
510    );
511    let num_buckets = hash.num_buckets();
512    let lens: Vec<u64> = decode_column(&parts[0], num_buckets, "bucket lengths");
513    let lens: Vec<usize> = lens.iter().map(|&l| l as usize).collect();
514    let terms: usize = lens.iter().sum();
515    let x = decode_rows::<W>(&parts[1], terms, "x column");
516    let z = decode_rows::<W>(&parts[2], terms, "z column");
517    let coeff = decode_column::<Complex64>(&parts[3], terms, "coeff column");
518    PauliSum::from_bucket_columns(&lens, x, z, coeff, hash, num_qubits)
519}
520
521#[cfg(test)]
522mod tests;
paulistrings/engine/partitioned/distributed.rs

paulistrings/engine/partitioned/
distributed.rs
