paulistrings/engine/partitioned/
driver.rs

1//! The call frame around `run_layers` shared by every partitioned driver, the scatter, and the [`propagate_partitioned`] front door.
2
3#[cfg(feature = "phase-timing")]
4use std::sync::atomic::Ordering;
5use std::time::Instant;
6
7use super::runtime::PartitionRuntime;
8use super::topology::{PartitionConfig, TopologyError};
9use crate::circuit::Circuit;
10use crate::collectives::Collectives;
11use crate::engine::backend::{PartitionBackend, PartitionStorage};
12use crate::engine::driver::{run_layers, PartitionContext, PartitionWork};
13use crate::engine::timing::tally;
14use crate::engine::trace::{assemble, PartitionTrace};
15use crate::engine::transport::Transport;
16use crate::engine::{progress_scope, Direction, PropagateOptions};
17use crate::pauli_sum::hash::PartitionRows;
18use crate::pauli_sum::storage::{desired_bits, DEFAULT_MIN_BUCKETS, DEFAULT_TARGET_BUCKET_LEN};
19use crate::pauli_sum::PauliSum;
20use crate::truncation::{assert_supports_partitions, TruncationPolicy};
21
22use super::sum::PartitionedSum;
23
24pub(super) use crate::engine::LOG_TARGET;
25
26impl<B> PartitionWork<B> {
27    pub(super) fn take<const W: usize>(local: &mut B, layers: usize, tracing: bool) -> Self
28    where
29        B: PartitionStorage<W>,
30    {
31        Self { local: local.detach(), rows: Vec::with_capacity(if tracing { layers } else { 0 }) }
32    }
33}
34
35/// One partition's share of `sum` at the agreed bucket count; collective, and run on the partition's own pool for first touch.
36pub(crate) fn scatter_local<const W: usize>(
37    sum: &PauliSum<W>,
38    rows: &PartitionRows<W>,
39    rank: u32,
40    collectives: &dyn Collectives,
41) -> PauliSum<W> {
42    let mut local = sum.filter_partition(rows, rank);
43    let want = desired_bits(local.len(), DEFAULT_TARGET_BUCKET_LEN, DEFAULT_MIN_BUCKETS);
44    let want = collectives.allreduce_max_u8(want);
45    local.coarsen_to(scatter_bits(local.hash().bits(), rows.bits(), want));
46    local
47}
48
49/// The scatter's agreed bucket bits, with `pbits = log2(P)` and `want` the all-reduced [`desired_bits`]: a partition sheds at most `pbits` of the bits the sum arrived with and never drops below what its share wants.
50fn scatter_bits(bits: u8, pbits: u8, want: u8) -> u8 {
51    want.max(bits.saturating_sub(pbits)).min(bits)
52}
53
54/// A driver's own laps since they were last drained.
55#[cfg(feature = "phase-timing")]
56pub(super) struct DriverLaps {
57    scatter_ns: u64,
58    /// Atomic because `gather` takes `&self`.
59    gather_ns: std::sync::atomic::AtomicU64,
60    layers: u64,
61}
62
63#[cfg(feature = "phase-timing")]
64impl DriverLaps {
65    pub(super) fn new(scatter_ns: u64) -> Self {
66        Self { scatter_ns, gather_ns: std::sync::atomic::AtomicU64::new(0), layers: 0 }
67    }
68
69    pub(super) fn lap_gather(&self, ns: u64) {
70        self.gather_ns.fetch_add(ns, Ordering::Relaxed);
71    }
72
73    /// Drain scatter time, gather time and layers driven.
74    pub(super) fn take(&mut self) -> (u64, u64, u64) {
75        (
76            std::mem::take(&mut self.scatter_ns),
77            self.gather_ns.swap(0, Ordering::Relaxed),
78            std::mem::take(&mut self.layers),
79        )
80    }
81}
82
83/// What one `propagate_on_backend` call reads besides a partition's own work, rank and endpoint.
84pub(super) struct LayerRun<'a, const W: usize, T: ?Sized> {
85    circuit: &'a Circuit<W>,
86    policy: &'a T,
87    direction: Direction,
88    options: PropagateOptions,
89    rows: &'a PartitionRows<W>,
90    size: usize,
91    tracing: bool,
92}
93
94impl<const W: usize, T: TruncationPolicy<W> + ?Sized> LayerRun<'_, W, T> {
95    /// Partition `rank`'s [`run_layers`].
96    pub(super) fn run<X: Transport, B: PartitionBackend<W, T>>(
97        &self,
98        rank: usize,
99        work: &mut PartitionWork<B>,
100        transport: &X,
101    ) {
102        let context =
103            PartitionContext { rows: self.rows, rank, size: self.size, tracing: self.tracing };
104        run_layers(
105            self.circuit,
106            self.policy,
107            self.direction,
108            self.options,
109            context,
110            work,
111            transport,
112        );
113    }
114}
115
116/// The partitions a driver holds and its records, lent to [`propagate`](Self::propagate) for one call.
117pub(super) struct DriverFrame<'a, const W: usize, B> {
118    /// In rank order: every partition in process, or this rank's one.
119    pub(super) parts: &'a mut [B],
120    pub(super) rows: &'a PartitionRows<W>,
121    pub(super) runtime: &'a PartitionRuntime,
122    pub(super) trace: &'a mut Option<PartitionTrace>,
123    #[cfg(feature = "phase-timing")]
124    pub(super) laps: &'a mut DriverLaps,
125    /// This process's rank in a distributed group, `None` when it holds every partition.
126    pub(super) rank: Option<usize>,
127    /// Partitions in the group.
128    pub(super) size: usize,
129}
130
131impl<const W: usize, B: PartitionStorage<W>> DriverFrame<'_, W, B> {
132    /// The prologue and epilogue every partitioned driver shares around `run`, which drives the partitions' layers and returns their work in rank order; `agree` runs after the entry line, before any layer.
133    pub(super) fn propagate<T>(
134        self,
135        circuit: &Circuit<W>,
136        policy: &T,
137        direction: Direction,
138        options: PropagateOptions,
139        agree: impl FnOnce(),
140        run: impl FnOnce(Vec<PartitionWork<B>>, &LayerRun<'_, W, T>) -> Vec<PartitionWork<B>>,
141    ) where
142        T: TruncationPolicy<W> + ?Sized,
143        B: PartitionBackend<W, T>,
144    {
145        options.assert_valid();
146        let num_layers = circuit.channels.len();
147        let size = self.size;
148        assert_supports_partitions(policy, size);
149        let terms_in: usize = self.parts.iter().map(PartitionStorage::len).sum();
150        let started = Instant::now();
151        log::info!(
152            target: LOG_TARGET,
153            "propagate: {terms_in} terms through {num_layers} channels ({direction:?}), {}",
154            progress_scope(self.rank, size, Some(&self.runtime.placement_summary())),
155        );
156
157        agree();
158
159        if num_layers > 0 {
160            let tracing = self.trace.is_some();
161            let items: Vec<PartitionWork<B>> = self
162                .parts
163                .iter_mut()
164                .map(|part| PartitionWork::take(part, num_layers, tracing))
165                .collect();
166            let layers =
167                LayerRun { circuit, policy, direction, options, rows: self.rows, size, tracing };
168            let done = run(items, &layers);
169            assert_eq!(done.len(), self.parts.len(), "one work item per partition");
170            let mut traced = Vec::with_capacity(if tracing { done.len() } else { 0 });
171            for (part, work) in self.parts.iter_mut().zip(done) {
172                *part = work.local;
173                if tracing {
174                    traced.push(work.rows);
175                }
176            }
177            if let Some(trace) = self.trace.as_mut() {
178                assemble(trace, traced, size);
179            }
180            tally!(self.laps.layers, num_layers);
181        }
182
183        log::info!(
184            target: LOG_TARGET,
185            "propagate: {num_layers} layers applied, {terms_in} -> {} terms, {:.3} s, {}",
186            self.parts.iter().map(PartitionStorage::len).sum::<usize>(),
187            started.elapsed().as_secs_f64(),
188            progress_scope(self.rank, size, Some(&self.runtime.placement_summary())),
189        );
190    }
191}
192
193/// Propagates `sum` through `circuit` on a partitioned engine built from `config`, and gathers the result.
194///
195/// A caller propagating repeatedly should hold a [`PartitionRuntime`] and a [`PartitionedSum`] instead, so pools, split and scratch survive between calls.
196///
197/// # Errors
198///
199/// [`TopologyError`] if `config` cannot be resolved into slots or a pool cannot be built.
200/// Everything else is a panic, as in [`propagate`](crate::propagate) and [`PartitionedSum::propagate_with`].
201pub fn propagate_partitioned<const W: usize, T>(
202    circuit: &Circuit<W>,
203    sum: PauliSum<W>,
204    policy: &T,
205    direction: Direction,
206    config: &PartitionConfig,
207    options: PropagateOptions,
208) -> Result<PauliSum<W>, TopologyError>
209where
210    T: TruncationPolicy<W> + ?Sized,
211{
212    let mut split = PartitionedSum::scatter(sum, config)?;
213    split.propagate_with(circuit, policy, direction, options);
214    Ok(split.into_gathered())
215}
216
217#[cfg(test)]
218mod tests;
paulistrings/engine/partitioned/driver.rs

paulistrings/engine/partitioned/
driver.rs
