paulistrings/engine/partitioned/
sum.rs

1//! [`PartitionedSum`], one [`PauliSum`] split across the partitions of a [`PartitionRuntime`].
2
3use std::sync::Arc;
4use std::time::Instant;
5
6use num_complex::Complex64;
7
8use super::distributed::{PartitionRowPolicy, ScatterOptions, ScatterRows};
9#[cfg(feature = "phase-timing")]
10use super::driver::DriverLaps;
11use super::driver::{scatter_local, DriverFrame, LOG_TARGET};
12use super::runtime::PartitionRuntime;
13use super::topology::{PartitionConfig, TopologyError};
14use crate::circuit::Circuit;
15use crate::engine::backend::{HostPartition, PartitionBackend, PartitionStorage};
16#[cfg(feature = "phase-timing")]
17use crate::engine::stats::PartitionPhaseStats;
18use crate::engine::timing::timed;
19use crate::engine::trace::PartitionTrace;
20use crate::engine::{progress_scope, Direction, PropagateOptions};
21use crate::pauli_sum::hash::PartitionRows;
22use crate::pauli_sum::PauliSum;
23use crate::readout::product_state::ProductBasis;
24use crate::truncation::TruncationPolicy;
25
26/// A [`PauliSum`] split across the partitions of a [`PartitionRuntime`], held across calls so a driver scatters and gathers once.
27///
28/// Between calls every partition's sum holds only keys of its own partition and shares one hash and bucket count with its peers ([`PartitionedSum::assert_invariants`]).
29/// `B` is where a partition lives: host memory by default, a CUDA device for `gpu::GpuPartitionedSum`.
30pub struct PartitionedSum<const W: usize, B = HostPartition<W>> {
31    /// In rank order.
32    parts: Vec<B>,
33    rows: PartitionRows<W>,
34    runtime: Arc<PartitionRuntime>,
35    /// `None` unless [`enable_trace`](Self::enable_trace) was called.
36    trace: Option<PartitionTrace>,
37    #[cfg(feature = "phase-timing")]
38    laps: DriverLaps,
39}
40
41impl<const W: usize, B> PartitionedSum<W, B> {
42    /// A split around partitions another backend already scattered, with `scatter_ns` its scatter time.
43    #[cfg(feature = "cuda")]
44    pub(crate) fn from_backend(
45        parts: Vec<B>,
46        rows: PartitionRows<W>,
47        runtime: Arc<PartitionRuntime>,
48        scatter_ns: u64,
49    ) -> Self {
50        #[cfg(not(feature = "phase-timing"))]
51        let _ = scatter_ns;
52        Self {
53            parts,
54            rows,
55            runtime,
56            trace: None,
57            #[cfg(feature = "phase-timing")]
58            laps: DriverLaps::new(scatter_ns),
59        }
60    }
61
62    #[cfg(feature = "cuda")]
63    pub(crate) fn backends(&self) -> &[B] {
64        &self.parts
65    }
66
67    #[cfg(feature = "cuda")]
68    pub(crate) fn backends_mut(&mut self) -> &mut [B] {
69        &mut self.parts
70    }
71
72    /// Drain scatter time, gather time and layers driven.
73    #[cfg(all(feature = "cuda", feature = "phase-timing"))]
74    pub(crate) fn take_driver_laps(&mut self) -> (u64, u64, u64) {
75        self.laps.take()
76    }
77
78    #[cfg(all(feature = "cuda", feature = "phase-timing"))]
79    pub(crate) fn lap_gather(&self, ns: u64) {
80        self.laps.lap_gather(ns);
81    }
82
83    /// Partitions this sum is split across.
84    pub fn num_partitions(&self) -> usize {
85        self.parts.len()
86    }
87
88    /// The rows deciding which partition a key belongs to.
89    pub fn rows(&self) -> &PartitionRows<W> {
90        &self.rows
91    }
92
93    /// The runtime this sum runs on, for handing to another [`PartitionedSum`].
94    pub fn runtime(&self) -> &Arc<PartitionRuntime> {
95        &self.runtime
96    }
97
98    /// Start recording a [`PartitionTrace`] on every later propagation; idempotent.
99    pub fn enable_trace(&mut self) {
100        self.trace.get_or_insert_with(PartitionTrace::default);
101    }
102
103    /// Drain the per-layer records, or `None` if tracing was never enabled; tracing stays enabled.
104    pub fn take_trace(&mut self) -> Option<PartitionTrace> {
105        self.trace.as_mut().map(std::mem::take)
106    }
107
108    /// The host driver's `propagate_with` on any backend.
109    pub(crate) fn propagate_on_backend<T>(
110        &mut self,
111        circuit: &Circuit<W>,
112        policy: &T,
113        direction: Direction,
114        options: PropagateOptions,
115    ) where
116        T: TruncationPolicy<W> + ?Sized,
117        B: PartitionBackend<W, T>,
118    {
119        let runtime = &*self.runtime;
120        let size = self.parts.len();
121        let frame = DriverFrame {
122            parts: &mut self.parts,
123            rows: &self.rows,
124            runtime,
125            trace: &mut self.trace,
126            #[cfg(feature = "phase-timing")]
127            laps: &mut self.laps,
128            rank: None,
129            size,
130        };
131        frame.propagate(
132            circuit,
133            policy,
134            direction,
135            options,
136            || {},
137            |items, layers| {
138                runtime.map_partitions(items, |rank, mut work, transport| {
139                    layers.run(rank, &mut work, transport);
140                    work
141                })
142            },
143        );
144    }
145}
146
147impl<const W: usize, B: PartitionStorage<W>> PartitionedSum<W, B> {
148    /// Terms in the whole sum, summed over partitions.
149    pub fn len(&self) -> usize {
150        self.parts.iter().map(|p| p.len()).sum()
151    }
152
153    /// Whether every partition is empty.
154    pub fn is_empty(&self) -> bool {
155        self.parts.iter().all(|p| p.len() == 0)
156    }
157
158    /// The bucket bits every partition currently holds.
159    pub fn bits(&self) -> u8 {
160        self.parts[0].hash().bits()
161    }
162
163    /// Qubits the sum is over.
164    pub fn num_qubits(&self) -> usize {
165        self.rows.num_qubits()
166    }
167}
168
169impl<const W: usize> PartitionedSum<W> {
170    /// Splits `sum` across the partitions `config` places, by the default rows ([`PartitionRowPolicy::Seeded`] at the sum's own hash seed); other rows go through [`scatter_with`](Self::scatter_with).
171    ///
172    /// # Errors
173    ///
174    /// [`TopologyError`] if `config` cannot be resolved or a pool cannot be built.
175    ///
176    /// # Panics
177    ///
178    /// In debug builds, if the partition rows are not independent of the sum's hash rows.
179    pub fn scatter(sum: PauliSum<W>, config: &PartitionConfig) -> Result<Self, TopologyError> {
180        let options = ScatterOptions {
181            runtime: PartitionRuntime::new(config)?,
182            rows: ScatterRows::Policy(PartitionRowPolicy::Seeded(None)),
183        };
184        Ok(Self::scatter_with(sum, options))
185    }
186
187    /// [`scatter`](Self::scatter) onto a caller-built runtime, split by the rows `options` names; at `P = 1` the scatter changes nothing.
188    ///
189    /// # Panics
190    ///
191    /// If the rows do not name the runtime's partition count, or as [`PartitionRows::cut`].
192    /// In debug builds, if the rows are not independent of the sum's hash rows, which costs load balance but not correctness.
193    pub fn scatter_with(sum: PauliSum<W>, options: ScatterOptions<W>) -> Self {
194        let ScatterOptions { runtime, rows } = options;
195        let rows = rows.resolve(&sum, runtime.partition_bits());
196        let size = runtime.num_partitions();
197        rows.assert_splits(sum.hash(), sum.num_qubits(), size);
198
199        let started = Instant::now();
200        let locals = {
201            let sum = &sum;
202            let rows = &rows;
203            runtime.map_partitions((0..size).collect(), |rank, _, transport| {
204                scatter_local(sum, rows, rank as u32, transport)
205            })
206        };
207        log::info!(
208            target: LOG_TARGET,
209            "scatter: {} terms -> {} held, {} bucket bits, {:.3} s, {}",
210            sum.len(),
211            locals.iter().map(PauliSum::len).sum::<usize>(),
212            locals[0].hash().bits(),
213            started.elapsed().as_secs_f64(),
214            progress_scope(None, size, Some(&runtime.placement_summary())),
215        );
216
217        Self {
218            parts: locals.into_iter().map(HostPartition::new).collect(),
219            rows,
220            runtime,
221            trace: None,
222            #[cfg(feature = "phase-timing")]
223            laps: DriverLaps::new(started.elapsed().as_nanos() as u64),
224        }
225    }
226
227    /// Propagates through `circuit` under `policy`, in place; see [`propagate_with`](Self::propagate_with).
228    pub fn propagate<T>(&mut self, circuit: &Circuit<W>, policy: &T, direction: Direction)
229    where
230        T: TruncationPolicy<W> + ?Sized,
231    {
232        self.propagate_with(circuit, policy, direction, PropagateOptions::default())
233    }
234
235    /// Propagates through `circuit` under `policy` with explicit [`PropagateOptions`], as [`propagate_with`](crate::propagate_with) does unpartitioned.
236    ///
237    /// The per-layer `DEBUG` log line comes once per partition, tagged `partition r/P`, from that partition's own thread.
238    ///
239    /// # Panics
240    ///
241    /// If `options` fail [`PropagateOptions::validate`], or a channel's preparation declines, as the unpartitioned engine does.
242    /// Before the first layer, above one partition, unless [`policy.supports_partitioned()`](TruncationPolicy::supports_partitioned) — exact [`TopN`](crate::TopN) is not yet supported partitioned.
243    /// A panic in any partition is re-raised on the calling thread.
244    pub fn propagate_with<T>(
245        &mut self,
246        circuit: &Circuit<W>,
247        policy: &T,
248        direction: Direction,
249        options: PropagateOptions,
250    ) where
251        T: TruncationPolicy<W> + ?Sized,
252    {
253        self.propagate_on_backend(circuit, policy, direction, options);
254    }
255
256    /// Merges the partitions back into one sum on the calling thread, leaving `self` intact.
257    pub fn gather(&self) -> PauliSum<W> {
258        timed!(let started = Instant::now());
259        let out = PauliSum::merge_partitions(self.parts.iter().map(|p| p.sum.clone()).collect());
260        timed!(self.laps.lap_gather(started.elapsed().as_nanos() as u64));
261        out
262    }
263
264    /// [`gather`](Self::gather) by value — no clone of the parts.
265    pub fn into_gathered(self) -> PauliSum<W> {
266        PauliSum::merge_partitions(self.parts.into_iter().map(|p| p.sum).collect())
267    }
268
269    /// Partition `r`'s share of the sum.
270    ///
271    /// # Panics
272    ///
273    /// If `r` is not a partition of this sum.
274    pub fn partition(&self, r: usize) -> &PauliSum<W> {
275        &self.parts[r].sum
276    }
277
278    /// `⟨ψ|O|ψ⟩` in the product state `basis`, the partitions' own values added in rank order.
279    pub fn expectation_product(&self, basis: &ProductBasis<W>) -> Complex64 {
280        self.parts.iter().map(|part| part.sum.expectation_product(basis)).sum()
281    }
282
283    /// Drain the phase counters since the last drain, zeroing them; `per_partition` has one entry per partition.
284    #[cfg(feature = "phase-timing")]
285    pub fn take_stats(&mut self) -> PartitionPhaseStats {
286        let (scatter_ns, gather_ns, layers) = self.laps.take();
287        PartitionPhaseStats {
288            per_partition: self
289                .parts
290                .iter_mut()
291                .map(|part| part.state.layer.take_phase_stats())
292                .collect(),
293            scatter_ns,
294            gather_ns,
295            layers,
296        }
297    }
298
299    /// Checks every invariant a partitioned sum holds between calls, in `O(terms)`; release builds skip the per-partition structural check.
300    ///
301    /// # Panics
302    ///
303    /// If a partition's sum is internally inconsistent, holds a key of another partition, or disagrees with partition 0 about the hash rows, the bucket count or the qubit count.
304    pub fn assert_invariants(&self) {
305        assert_eq!(
306            self.parts.len(),
307            self.rows.num_partitions(),
308            "partition count disagrees with the rows",
309        );
310        let head = &self.parts[0].sum;
311        for (rank, local) in self.parts.iter().map(|p| &p.sum).enumerate() {
312            #[cfg(any(test, debug_assertions))]
313            local.assert_invariants();
314            assert_eq!(
315                local.num_qubits(),
316                head.num_qubits(),
317                "partition {rank}: qubit count differs from partition 0",
318            );
319            assert!(
320                local.hash().same_rows_as(head.hash()),
321                "partition {rank}: hash rows differ from partition 0",
322            );
323            assert_eq!(
324                local.hash().bits(),
325                head.hash().bits(),
326                "partition {rank}: bucket bits differ from partition 0",
327            );
328            let held = local.partition_rank_of_all(&self.rows);
329            assert!(
330                held.is_none() || held == Some(rank as u32),
331                "partition {rank} holds keys of partition {held:?}",
332            );
333        }
334    }
335}
paulistrings/engine/partitioned/sum.rs

paulistrings/engine/partitioned/
sum.rs
