paulistrings/engine/
bucketed.rs

1//! One layer of the bucketed engine, applied coset by coset, and the reusable `LayerScratch`.
2//!
3//! A coset task swaps its bucket columns into worker scratch and writes its outputs back into the emptied columns, so the layer runs in place.
4//! Identity-delta rows keep their keys, so a run's identity stream inherits the source bucket's order and is never sorted; every other row, received rows included, joins the rest stream.
5
6use std::sync::Mutex;
7
8use num_complex::Complex64;
9use rayon::prelude::*;
10
11use super::coset::{CosetOrder, Gf2Span};
12use super::merge::{rest_rows_per_key, SortKernel};
13use super::trace::PartitionTrace;
14use crate::channel::prepared::{LocalPtm, Prepared, PreparedRotation};
15use crate::pauli_sum::storage::{BucketColumns, PauliSum};
16use crate::truncation::TruncationPolicy;
17
18#[cfg(feature = "phase-timing")]
19use super::stats::{PartitionPhaseStats, PhaseStats};
20use super::timing::{lap, stamp};
21
22mod coset_fill;
23
24pub(super) use coset_fill::{fill_coset, CosetScratch, MIN_COSETS_FOR_PARALLEL};
25
26const ZERO: Complex64 = Complex64::new(0.0, 0.0);
27
28/// Reusable layer working set for [`propagate_with`](crate::propagate_with), which keeps its capacity across layers and calls and carries the opt-in traces.
29///
30/// A task's output cannot depend on which worker slot it drew: every write destination is cleared before use.
31#[derive(Debug, Default)]
32pub struct LayerScratch<const W: usize> {
33    /// The serial path's coset working set.
34    pub(super) task: CosetScratch<W>,
35    /// The layer's coset order, with its inverse only under [`ExtraRows::NEEDS_BETA`]; a partitioned layer builds it before the export and passes it back in [`LayerKnobs::order`].
36    pub(super) order: CosetOrder,
37    /// Bucket handles in coset-contiguous order while a layer runs.
38    pub(super) staging: Vec<BucketColumns<W>>,
39    /// One coset working set per Rayon worker, indexed by `rayon::current_thread_index()`, so each mutex is uncontended.
40    pub(super) workers: Vec<Mutex<CosetScratch<W>>>,
41    /// Layer-level wall-clock counters; per-coset busy time lives in each `CosetScratch`.
42    #[cfg(feature = "phase-timing")]
43    pub(crate) stats: PhaseStats,
44    /// `None` unless [`Self::enable_trace`] was called.
45    pub(crate) trace: Option<PartitionTrace>,
46}
47
48impl<const W: usize> LayerScratch<W> {
49    /// An empty scratch.
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    /// Drain the phase counters since the last drain, zeroing them; `per_partition` has one entry.
55    #[cfg(feature = "phase-timing")]
56    pub fn take_stats(&mut self) -> PartitionPhaseStats {
57        let stats = self.take_phase_stats();
58        PartitionPhaseStats {
59            layers: stats.layers,
60            per_partition: vec![stats],
61            scatter_ns: 0,
62            gather_ns: 0,
63        }
64    }
65
66    /// Drain this scratch's own counters, one partition's entry of a driver's [`PartitionPhaseStats`].
67    #[cfg(feature = "phase-timing")]
68    pub(crate) fn take_phase_stats(&mut self) -> PhaseStats {
69        let mut total = std::mem::take(&mut self.stats);
70        total.absorb_coset(&std::mem::take(&mut self.task.stats));
71        for slot in &self.workers {
72            let mut worker_scratch = slot.lock().unwrap();
73            total.absorb_coset(&std::mem::take(&mut worker_scratch.stats));
74        }
75        total
76    }
77
78    /// Start recording a one-partition [`PartitionTrace`] on every later propagation driven by this scratch; idempotent.
79    pub fn enable_trace(&mut self) {
80        self.trace.get_or_insert_with(PartitionTrace::default);
81    }
82
83    /// Drain the per-layer records, or `None` if tracing was never enabled; tracing stays enabled.
84    pub fn take_trace(&mut self) -> Option<PartitionTrace> {
85        self.trace.as_mut().map(std::mem::take)
86    }
87}
88
89/// A prepared channel's delta set annotated with each entry's coset coordinate, computed once per layer.
90pub(super) enum DeltaPlan<'p, const W: usize> {
91    /// Tabulated deltas; `coordinates[e]` pairs with `ptm.deltas()[e]`.
92    Local {
93        ptm: &'p LocalPtm<W>,
94        coordinates: Vec<u32>,
95        /// Whether `deltas()[0]` is the identity delta, whose stream goes to the run's pre-sorted `id` columns.
96        has_identity: bool,
97        /// Whether the identity amplitude is nonzero on every active support pattern, so the identity stream is one-to-one with the source bucket and the merge borrows its key columns.
98        dense_identity: bool,
99        /// This layer's per-run sort kernel.
100        sort_kernel: SortKernel,
101    },
102    /// Wide rotation: two implicit entries, the identity pass and the generator pass.
103    Rotation {
104        rotation: &'p PreparedRotation<W>,
105        coordinate_identity: u32,
106        generator_coordinate: u32,
107        /// Whether the generator pass emits here ([`LayerKnobs::generator_local`]); `generator_coordinate` is 0 when not.
108        generator_local: bool,
109    },
110}
111
112/// Per-layer overrides the partitioned engine hands the coset loop; [`Default`] is the unpartitioned answer.
113#[derive(Clone, Copy, Debug)]
114pub(crate) struct LayerKnobs<'k> {
115    /// Bucket deltas to build the coset span from, instead of [`Prepared::bucket_deltas`].
116    pub bucket_deltas: Option<&'k [u32]>,
117    /// The channel's total rest-stream count, so every partition picks the sort kernel the unpartitioned run would.
118    pub rest_streams: Option<usize>,
119    /// The unrestricted channel's [`rest_rows_per_key`], since a PTM cut down to local entries can look more disjoint than the channel is.
120    pub rows_per_key: Option<f64>,
121    /// Whether a wide rotation's generator pass emits here; false when every generator row belongs to a partner.
122    pub generator_local: bool,
123    /// The layer's coset order, already built `with_inverse` over the span of `bucket_deltas`; built here when `None`.
124    pub order: Option<&'k CosetOrder>,
125}
126
127impl Default for LayerKnobs<'_> {
128    fn default() -> Self {
129        Self {
130            bucket_deltas: None,
131            rest_streams: None,
132            rows_per_key: None,
133            generator_local: true,
134            order: None,
135        }
136    }
137}
138
139impl<'p, const W: usize> DeltaPlan<'p, W> {
140    pub(super) fn new(prepared: &'p Prepared<W>, span: &Gf2Span, knobs: LayerKnobs<'_>) -> Self {
141        match prepared {
142            Prepared::Local(ptm) => {
143                let coordinates: Vec<u32> =
144                    ptm.deltas().iter().map(|d| span.coordinate_of(d.bucket_delta)).collect();
145                let has_identity = ptm.has_identity();
146                debug_assert!(!has_identity || coordinates[0] == 0);
147                let dense_identity = has_identity && ptm.entry_is_dense(0);
148                let rest_streams = knobs.rest_streams.unwrap_or(ptm.rest_streams());
149                let sort_kernel = SortKernel::for_layer(rest_streams, || {
150                    knobs.rows_per_key.unwrap_or_else(|| rest_rows_per_key(ptm))
151                });
152                DeltaPlan::Local { ptm, coordinates, has_identity, dense_identity, sort_kernel }
153            }
154            Prepared::Rotation(rotation) => DeltaPlan::Rotation {
155                rotation,
156                coordinate_identity: span.coordinate_of(rotation.bucket_delta_identity),
157                // `coordinate_of` demands its argument be in the span, and a remote generator's bucket delta is not.
158                generator_coordinate: if knobs.generator_local {
159                    span.coordinate_of(rotation.bucket_delta_generator)
160                } else {
161                    0
162                },
163                generator_local: knobs.generator_local,
164            },
165        }
166    }
167}
168
169/// Rows received from other partitions, appended to an output bucket's rest stream before the sort so `keep_term` sees the complete sum; never the identity stream, which the merge requires strictly ascending.
170/// An implementation that can return rows must set `NEEDS_BETA = true`; under `false` every call site compiles away.
171pub(crate) trait ExtraRows<const W: usize> {
172    /// Whether the engine recovers original bucket indices and calls this source at all.
173    const NEEDS_BETA: bool = false;
174
175    /// Rows destined for original bucket `beta`; must equal what [`append_into`](Self::append_into) pushes.
176    fn count(&self, beta: u32) -> usize {
177        let _ = beta;
178        0
179    }
180
181    /// Append original bucket `beta`'s rows onto the run's rest columns.
182    fn append_into(&self, beta: u32, rest: &mut BucketColumns<W>) {
183        let _ = (beta, rest);
184    }
185}
186
187/// The unpartitioned engine's [`ExtraRows`]: no rows.
188pub(crate) struct NoExtra;
189
190impl<const W: usize> ExtraRows<W> for NoExtra {}
191
192/// Apply one prepared channel to a bucketed sum, folding `keep_term` into the merge; the caller runs `finalize_layer`.
193pub fn apply_layer_bucketed<const W: usize, T>(
194    sum: &mut PauliSum<W>,
195    prepared: &Prepared<W>,
196    policy: &T,
197    scratch: &mut LayerScratch<W>,
198) where
199    T: TruncationPolicy<W> + ?Sized,
200{
201    apply_layer_bucketed_with(sum, prepared, policy, scratch, &NoExtra, LayerKnobs::default())
202}
203
204/// [`apply_layer_bucketed`] with the partitioned engine's [`ExtraRows`] source and [`LayerKnobs`].
205pub(crate) fn apply_layer_bucketed_with<const W: usize, T, X>(
206    sum: &mut PauliSum<W>,
207    prepared: &Prepared<W>,
208    policy: &T,
209    scratch: &mut LayerScratch<W>,
210    extra: &X,
211    knobs: LayerKnobs<'_>,
212) where
213    T: TruncationPolicy<W> + ?Sized,
214    X: ExtraRows<W> + Sync,
215{
216    stamp!(stamp);
217
218    // A partitioned layer's local table can look key-preserving while received rows still need merging.
219    if let Prepared::Local(ptm) = prepared {
220        if !X::NEEDS_BETA && ptm.is_key_preserving() {
221            rescale_in_place(sum, ptm, policy);
222            lap!(stamp, scratch.stats.rescale_ns);
223            return;
224        }
225    }
226
227    let own_deltas;
228    let deltas: &[u32] = match knobs.bucket_deltas {
229        Some(deltas) => deltas,
230        None => {
231            own_deltas = prepared.bucket_deltas();
232            &own_deltas
233        }
234    };
235    let span = Gf2Span::new(deltas, sum.hash().bits());
236    let plan = DeltaPlan::new(prepared, &span, knobs);
237    let coset_size = span.coset_size();
238    let num_cosets = span.num_cosets();
239    lap!(stamp, scratch.stats.span_plan_ns);
240
241    // Coset `c` owns `staging[c·2^r .. (c+1)·2^r]`; at `r = 0` the order is the identity and the handle passes are skipped.
242    let order: &CosetOrder = match knobs.order {
243        Some(order) => order,
244        None => {
245            scratch.order.rebuild(&span, X::NEEDS_BETA);
246            &scratch.order
247        }
248    };
249    debug_assert_eq!(order.positions(), sum.num_buckets());
250    debug_assert_eq!(order.dimension(), span.dimension());
251    let identity_order = order.is_identity();
252    if !identity_order {
253        let buckets = sum.buckets_mut();
254        scratch.staging.resize_with(buckets.len(), BucketColumns::default);
255        for (beta, columns) in buckets.iter_mut().enumerate() {
256            scratch.staging[order.position_of(beta as u32) as usize] = std::mem::take(columns);
257        }
258    }
259    lap!(stamp, scratch.stats.permute_ns);
260
261    // Each coset reads and writes only its own chunk (engine/coset.rs); research/FINDINGS.md §Static coset→worker placement.
262    {
263        if num_cosets >= MIN_COSETS_FOR_PARALLEL {
264            let pool = rayon::current_num_threads().max(1);
265            if scratch.workers.len() < pool {
266                scratch.workers.resize_with(pool, Mutex::default);
267            }
268        }
269        let workers = &scratch.workers;
270        let chunks: &mut [BucketColumns<W>] =
271            if identity_order { sum.buckets_mut() } else { scratch.staging.as_mut_slice() };
272        if num_cosets < MIN_COSETS_FOR_PARALLEL {
273            for (coset_index, chunk) in chunks.chunks_mut(coset_size).enumerate() {
274                let base = coset_index * coset_size;
275                fill_coset::<W, T, X>(chunk, &plan, policy, &mut scratch.task, extra, base, order);
276            }
277        } else {
278            chunks.par_chunks_mut(coset_size).enumerate().for_each(|(coset_index, chunk)| {
279                let base = coset_index * coset_size;
280                // `workers` was sized to `current_num_threads()` above, and this closure runs on that pool.
281                let worker = rayon::current_thread_index()
282                    .expect("the coset loop runs on a Rayon pool worker");
283                let mut worker_scratch = workers[worker].lock().unwrap();
284                fill_coset::<W, T, X>(
285                    chunk,
286                    &plan,
287                    policy,
288                    &mut worker_scratch,
289                    extra,
290                    base,
291                    order,
292                );
293            });
294        }
295    }
296    lap!(stamp, scratch.stats.coset_loop_ns);
297
298    if !identity_order {
299        let buckets = sum.buckets_mut();
300        for (beta, columns) in buckets.iter_mut().enumerate() {
301            *columns =
302                std::mem::take(&mut scratch.staging[order.position_of(beta as u32) as usize]);
303        }
304    }
305    lap!(stamp, scratch.stats.unpermute_ns);
306    sum.recount();
307    lap!(stamp, scratch.stats.recount_ns);
308
309    #[cfg(debug_assertions)]
310    sum.assert_invariants();
311}
312
313/// In-place coefficient rescale for a key-preserving channel, with the general path's zero drop and `keep_term`.
314fn rescale_in_place<const W: usize, T>(sum: &mut PauliSum<W>, ptm: &LocalPtm<W>, policy: &T)
315where
316    T: TruncationPolicy<W> + ?Sized,
317{
318    let amplitude = &ptm.deltas()[0].amplitude;
319    sum.update_per_bucket(|_| {
320        |x: &[u64; W], z: &[u64; W], c: Complex64| {
321            let coeff = c * amplitude[ptm.support_bits(x, z)];
322            (coeff != ZERO && policy.keep_term(x, z, coeff)).then_some(coeff)
323        }
324    });
325
326    #[cfg(debug_assertions)]
327    sum.assert_invariants();
328}
329
330#[cfg(test)]
331mod tests;
paulistrings/engine/bucketed.rs

paulistrings/engine/
bucketed.rs
