paulistrings/pauli_sum/
hash.rs

1//! The GF(2)-linear bucket hash [`Gf2Hash`] and the partition rows [`PartitionRows`].
2
3use crate::pauli_string::word_mask;
4use crate::rng::{mix64, SPLITMIX_GAMMA};
5
6/// Maximum number of bucket bits; all rows are drawn up front so the active hash is a prefix of one fixed matrix.
7pub const B_MAX_BITS: u8 = 22;
8
9/// Maximum number of partition bits, i.e. `P ≤ 64` partitions.
10pub const P_MAX_BITS: u8 = 6;
11
12/// Salt keeping [`PartitionRows::from_seed`]'s rows independent of [`Gf2Hash::new`]'s under the same seed.
13const PARTITION_ROW_SALT: u64 = 0xD1B5_4A32_D192_ED03;
14
15/// Word `word` of the `half` (0 = x, 1 = z) of row `row`, draw `attempt`.
16// Not a GF(2)-linear generator such as xorshift: research/FINDINGS.md §`Gf2Hash` rows are splitmix64, not xorshift successors.
17fn row_word(seed: u64, row: usize, attempt: u32, word: usize, half: u64) -> u64 {
18    debug_assert!(row < 1 << 16 && word < 1 << 15);
19    let position = ((attempt as u64) << 32) | ((row as u64) << 16) | ((word as u64) << 1) | half;
20    mix64(seed.wrapping_add(SPLITMIX_GAMMA.wrapping_mul(position.wrapping_add(1))))
21}
22
23/// A GF(2) matrix over the key space `(x, z)`, each row masked to the live qubit columns; `Gf2Hash` and `PartitionRows` are its two uses.
24#[derive(Clone, Debug, PartialEq, Eq)]
25pub(crate) struct Gf2Matrix<const W: usize> {
26    rows_x: Vec<[u64; W]>,
27    rows_z: Vec<[u64; W]>,
28    num_qubits: usize,
29}
30
31impl<const W: usize> Gf2Matrix<W> {
32    /// `num_rows` rows drawn from `seed` with the `exclude` coordinates cleared, an all-zero row redrawn unless `num_qubits == 0`.
33    pub(crate) fn draw(
34        num_qubits: usize,
35        num_rows: usize,
36        seed: u64,
37        exclude_x: &[u64; W],
38        exclude_z: &[u64; W],
39    ) -> Self {
40        let mut rows_x: Vec<[u64; W]> = Vec::with_capacity(num_rows);
41        let mut rows_z: Vec<[u64; W]> = Vec::with_capacity(num_rows);
42        let has_live_columns = num_qubits > 0;
43        for row in 0..num_rows {
44            let mut attempt = 0u32;
45            let (row_x, row_z) = loop {
46                let mut row_x = [0u64; W];
47                let mut row_z = [0u64; W];
48                let mut any = false;
49                for w in 0..W {
50                    let mask = word_mask(num_qubits, w);
51                    row_x[w] = row_word(seed, row, attempt, w, 0) & mask & !exclude_x[w];
52                    row_z[w] = row_word(seed, row, attempt, w, 1) & mask & !exclude_z[w];
53                    any |= (row_x[w] | row_z[w]) != 0;
54                }
55                if any || !has_live_columns {
56                    break (row_x, row_z);
57                }
58                attempt += 1;
59            };
60            rows_x.push(row_x);
61            rows_z.push(row_z);
62        }
63        Self { rows_x, rows_z, num_qubits }
64    }
65
66    /// Explicit rows, masked to the live columns.
67    pub(crate) fn from_rows(
68        num_qubits: usize,
69        mut rows_x: Vec<[u64; W]>,
70        mut rows_z: Vec<[u64; W]>,
71    ) -> Self {
72        assert_eq!(rows_x.len(), rows_z.len(), "Gf2Matrix: row count mismatch");
73        for (row_x, row_z) in rows_x.iter_mut().zip(rows_z.iter_mut()) {
74            for w in 0..W {
75                let mask = word_mask(num_qubits, w);
76                row_x[w] &= mask;
77                row_z[w] &= mask;
78            }
79        }
80        Self { rows_x, rows_z, num_qubits }
81    }
82
83    pub(crate) fn num_rows(&self) -> usize {
84        self.rows_x.len()
85    }
86
87    pub(crate) fn num_qubits(&self) -> usize {
88        self.num_qubits
89    }
90
91    pub(crate) fn row(&self, i: usize) -> KeyRow<W> {
92        (self.rows_x[i], self.rows_z[i])
93    }
94
95    pub(crate) fn rows(&self) -> (&[[u64; W]], &[[u64; W]]) {
96        (&self.rows_x, &self.rows_z)
97    }
98
99    /// Bit `row` of `M·v`, as `0` or `1`.
100    pub(crate) fn row_parity(&self, x: &[u64; W], z: &[u64; W], row: usize) -> u32 {
101        let row_x = &self.rows_x[row];
102        let row_z = &self.rows_z[row];
103        let mut folded: u64 = 0;
104        for w in 0..W {
105            folded ^= (x[w] & row_x[w]) ^ (z[w] & row_z[w]);
106        }
107        folded.count_ones() & 1
108    }
109
110    /// The first `rows` bits of `M·v`, row `i` at bit `i`.
111    pub(crate) fn apply(&self, x: &[u64; W], z: &[u64; W], rows: usize) -> u32 {
112        let mut index: u32 = 0;
113        for i in 0..rows {
114            index |= self.row_parity(x, z, i) << i;
115        }
116        index
117    }
118
119    /// Whether this matrix's first `rows` rows and `other`'s first `other_rows` are jointly GF(2)-independent.
120    pub(crate) fn independent_with(&self, rows: usize, other: &Self, other_rows: usize) -> bool {
121        let stacked: Vec<KeyRow<W>> =
122            (0..rows).map(|i| self.row(i)).chain((0..other_rows).map(|i| other.row(i))).collect();
123        gf2_rank_wide(&stacked) == rows + other_rows
124    }
125}
126
127/// The GF(2)-linear hash `h(v) = H·v` from a Pauli key `v = (x, z)` to a bucket index, for a dense random `H` drawn from a seed.
128///
129/// Linearity gives `h(v ⊕ d) = h(v) ⊕ h(d)`, so a channel's output buckets are predictable from its input buckets.
130/// The active rows are a prefix of one fixed matrix, so a refine splits every bucket in two by one row parity with the within-bucket order inherited, and a coarsen merges bucket pairs; neither needs a sort to keep the [`PauliSum`](crate::PauliSum) invariant.
131#[derive(Clone, Debug)]
132pub struct Gf2Hash<const W: usize> {
133    /// All [`B_MAX_BITS`] rows of `H`.
134    matrix: Gf2Matrix<W>,
135    /// Active prefix length: `B = 1 << bits` buckets.
136    bits: u8,
137    /// Seed the rows were generated from.
138    seed: u64,
139}
140
141impl<const W: usize> Gf2Hash<W> {
142    /// Build a hash over `num_qubits` qubits with `bits` active bucket bits.
143    ///
144    /// The rows are a function of `(num_qubits, seed)` alone, independent of `W`.
145    ///
146    /// # Panics
147    ///
148    /// Panics if `bits > B_MAX_BITS`, or in debug builds if `num_qubits > 64 · W`.
149    pub fn new(num_qubits: usize, bits: u8, seed: u64) -> Self {
150        assert!(bits <= B_MAX_BITS, "Gf2Hash: bits {bits} exceeds B_MAX_BITS {B_MAX_BITS}",);
151        debug_assert!(num_qubits <= 64 * W);
152
153        Self {
154            matrix: Gf2Matrix::draw(num_qubits, B_MAX_BITS as usize, seed, &[0; W], &[0; W]),
155            bits,
156            seed,
157        }
158    }
159
160    /// Number of active bucket bits.
161    pub fn bits(&self) -> u8 {
162        self.bits
163    }
164
165    /// Number of buckets, `1 << bits()`.
166    pub fn num_buckets(&self) -> usize {
167        1usize << self.bits
168    }
169
170    /// The seed the rows were generated from.
171    pub fn seed(&self) -> u64 {
172        self.seed
173    }
174
175    /// The qubit count the rows were masked against.
176    pub fn num_qubits(&self) -> usize {
177        self.matrix.num_qubits()
178    }
179
180    /// `h(v)` for a key given as separate `x` and `z` words.
181    pub fn bucket_of(&self, x: &[u64; W], z: &[u64; W]) -> u32 {
182        self.matrix.apply(x, z, self.bits as usize)
183    }
184
185    /// Bit `row` of `H·v`, as `0` or `1`.
186    pub(super) fn row_parity(&self, x: &[u64; W], z: &[u64; W], row: u8) -> u32 {
187        self.matrix.row_parity(x, z, row as usize)
188    }
189
190    /// Double the bucket count: `B → 2B`.
191    ///
192    /// # Panics
193    ///
194    /// Panics if already at `B_MAX_BITS`.
195    pub fn refine(&mut self) {
196        assert!(self.bits < B_MAX_BITS, "Gf2Hash::refine: already at B_MAX_BITS {B_MAX_BITS}",);
197        self.bits += 1;
198    }
199
200    /// Halve the bucket count: `B → B/2`, merging buckets `b` and `b + B/2`.
201    ///
202    /// # Panics
203    ///
204    /// Panics if already at a single bucket.
205    pub fn coarsen(&mut self) {
206        assert!(self.bits > 0, "Gf2Hash::coarsen: already at a single bucket");
207        self.bits -= 1;
208    }
209
210    /// `true` if `other` was generated with the same rows, so sums partitioned by the two can be combined (after matching `bits`).
211    pub(crate) fn same_rows_as(&self, other: &Self) -> bool {
212        self.seed == other.seed && self.num_qubits() == other.num_qubits()
213    }
214
215    /// Row `i < B_MAX_BITS` of `H`, which may lie beyond the active prefix.
216    #[cfg(any(test, feature = "cuda"))]
217    pub(crate) fn row(&self, i: usize) -> ([u64; W], [u64; W]) {
218        self.matrix.row(i)
219    }
220}
221
222/// The `p` GF(2) rows `part(v) = P·v` that split a sum across `2^p` partitions: a partition holds exactly the terms with `part(v)` equal to its rank.
223///
224/// A matrix separate from [`Gf2Hash`]'s and fixed for a run, since a row that moved as the bucket hash refines would change a term's owner mid-run.
225/// Dependence on the bucket hash's rows costs load balance, never correctness.
226#[derive(Clone, Debug, PartialEq, Eq)]
227pub struct PartitionRows<const W: usize> {
228    /// One row per partition bit: `P = 1 << rows` partitions.
229    matrix: Gf2Matrix<W>,
230}
231
232impl<const W: usize> PartitionRows<W> {
233    /// Panics unless the rows split a sum over `num_qubits` qubits into `partitions` partitions; debug builds also check them independent of `hash`.
234    pub(crate) fn assert_splits(&self, hash: &Gf2Hash<W>, num_qubits: usize, partitions: usize) {
235        assert_eq!(
236            self.num_partitions(),
237            partitions,
238            "partition rows name {} partitions for a split into {partitions}",
239            self.num_partitions(),
240        );
241        assert_eq!(
242            self.num_qubits(),
243            num_qubits,
244            "partition rows are for {} qubits, the sum for {num_qubits}",
245            self.num_qubits(),
246        );
247        debug_assert!(
248            self.is_independent_of(hash),
249            "partition rows are dependent on the bucket hash rows — the split will correlate \
250             with the bucket partition and load-balance badly",
251        );
252    }
253
254    /// Draw `bits` partition rows from `seed`, unrelated to [`Gf2Hash::new`]'s rows under the same seed.
255    ///
256    /// # Panics
257    ///
258    /// Panics if `bits > P_MAX_BITS`, or in debug builds if `num_qubits > 64 · W`.
259    pub fn from_seed(num_qubits: usize, bits: u8, seed: u64) -> Self {
260        Self::from_seed_excluding(num_qubits, bits, seed, &[0; W], &[0; W])
261    }
262
263    /// [`Self::from_seed`] with the coordinates in `(exclude_x, exclude_z)` cleared from every row, so keys differing only there share a partition.
264    ///
265    /// # Panics
266    ///
267    /// As [`Self::from_seed`], and if `bits > 0` while the exclusion covers every live column.
268    pub fn from_seed_excluding(
269        num_qubits: usize,
270        bits: u8,
271        seed: u64,
272        exclude_x: &[u64; W],
273        exclude_z: &[u64; W],
274    ) -> Self {
275        assert!(bits <= P_MAX_BITS, "PartitionRows: bits {bits} exceeds P_MAX_BITS {P_MAX_BITS}",);
276        debug_assert!(num_qubits <= 64 * W);
277
278        let keeps_a_column = (0..W).any(|w| {
279            let mask = word_mask(num_qubits, w);
280            (mask & !exclude_x[w]) | (mask & !exclude_z[w]) != 0
281        });
282        assert!(
283            bits == 0 || num_qubits == 0 || keeps_a_column,
284            "PartitionRows::from_seed_excluding: the exclusion covers every column",
285        );
286        Self {
287            matrix: Gf2Matrix::draw(
288                num_qubits,
289                bits as usize,
290                mix64(seed) ^ PARTITION_ROW_SALT,
291                exclude_x,
292                exclude_z,
293            ),
294        }
295    }
296
297    /// Build from explicit rows, masked to the live qubit columns.
298    ///
299    /// # Panics
300    ///
301    /// Panics if `rows_x` and `rows_z` differ in length, if there are more than [`P_MAX_BITS`] rows, or if any row masks to all-zero while `num_qubits > 0`.
302    pub(crate) fn from_rows(
303        num_qubits: usize,
304        rows_x: Vec<[u64; W]>,
305        rows_z: Vec<[u64; W]>,
306    ) -> Self {
307        assert_eq!(
308            rows_x.len(),
309            rows_z.len(),
310            "PartitionRows::from_rows: row count mismatch, {} x-rows vs {} z-rows",
311            rows_x.len(),
312            rows_z.len(),
313        );
314        assert!(
315            rows_x.len() <= P_MAX_BITS as usize,
316            "PartitionRows: bits {} exceeds P_MAX_BITS {P_MAX_BITS}",
317            rows_x.len(),
318        );
319        debug_assert!(num_qubits <= 64 * W);
320
321        let matrix = Gf2Matrix::from_rows(num_qubits, rows_x, rows_z);
322        for i in 0..matrix.num_rows() {
323            let (row_x, row_z) = matrix.row(i);
324            assert!(
325                num_qubits == 0 || row_x.iter().chain(&row_z).any(|&w| w != 0),
326                "PartitionRows::from_rows: row {i} masks to zero",
327            );
328        }
329        Self { matrix }
330    }
331
332    /// Rows that label a qubit cut: `log2(blocks.len())` z-only rows giving block `b` the label `b`.
333    ///
334    /// A term's label is the XOR of the labels of the blocks it has odd z-weight in, so an x-only generator is local under any cut and a `ZZ(i, j)` bond is remote exactly when it crosses blocks.
335    /// Blocks must be disjoint; qubits in no block behave as if in block 0.
336    ///
337    /// # Panics
338    ///
339    /// Panics if `blocks.len()` is not a power of two or exceeds `2^P_MAX_BITS`, if a qubit is `>= num_qubits` or appears in two blocks, or if some row would be all-zero.
340    pub fn cut(num_qubits: usize, blocks: &[Vec<usize>]) -> Self {
341        assert!(
342            blocks.len().is_power_of_two(),
343            "PartitionRows::cut: block count {} is not a power of two",
344            blocks.len(),
345        );
346        let bits = blocks.len().trailing_zeros() as u8;
347        assert!(bits <= P_MAX_BITS, "PartitionRows: bits {bits} exceeds P_MAX_BITS {P_MAX_BITS}",);
348        debug_assert!(num_qubits <= 64 * W);
349
350        let mut seen = vec![false; num_qubits];
351        let mut rows_z = vec![[0u64; W]; bits as usize];
352        for (b, qubits) in blocks.iter().enumerate() {
353            for &q in qubits {
354                assert!(
355                    q < num_qubits,
356                    "PartitionRows::cut: qubit {q} in block {b} is outside 0..{num_qubits}",
357                );
358                assert!(
359                    !seen[q],
360                    "PartitionRows::cut: blocks must be disjoint, qubit {q} appears twice",
361                );
362                seen[q] = true;
363                for (i, row) in rows_z.iter_mut().enumerate() {
364                    if (b >> i) & 1 == 1 {
365                        row[q / 64] |= 1u64 << (q % 64);
366                    }
367                }
368            }
369        }
370        for (i, row) in rows_z.iter().enumerate() {
371            assert!(
372                row.iter().any(|w| *w != 0),
373                "PartitionRows::cut: row {i} is empty — no qubit lies in a block \
374                 whose label has bit {i} set",
375            );
376        }
377
378        Self::from_rows(num_qubits, vec![[0u64; W]; bits as usize], rows_z)
379    }
380
381    /// The trivial partitioning: one partition, no rows.
382    pub fn none(num_qubits: usize) -> Self {
383        Self { matrix: Gf2Matrix::from_rows(num_qubits, Vec::new(), Vec::new()) }
384    }
385
386    /// Number of partition bits.
387    pub fn bits(&self) -> u8 {
388        self.matrix.num_rows() as u8
389    }
390
391    /// Number of partitions, `1 << bits()`.
392    pub fn num_partitions(&self) -> usize {
393        1usize << self.bits()
394    }
395
396    /// The qubit count the rows were masked against.
397    pub fn num_qubits(&self) -> usize {
398        self.matrix.num_qubits()
399    }
400
401    /// `part(v)` for a key given as separate `x` and `z` words.
402    pub fn partition_of(&self, x: &[u64; W], z: &[u64; W]) -> u32 {
403        self.matrix.apply(x, z, self.matrix.num_rows())
404    }
405
406    /// The rows as `(x-masks, z-masks)`, already masked to the live columns.
407    pub fn rows(&self) -> (&[[u64; W]], &[[u64; W]]) {
408        self.matrix.rows()
409    }
410
411    /// `true` if no row reads a coordinate in `(mask_x, mask_z)`, i.e. keys differing only there always share a partition.
412    pub fn avoids(&self, mask_x: &[u64; W], mask_z: &[u64; W]) -> bool {
413        let (rows_x, rows_z) = self.matrix.rows();
414        rows_x.iter().zip(rows_z).all(|(row_x, row_z)| {
415            (0..W).all(|w| row_x[w] & mask_x[w] == 0 && row_z[w] & mask_z[w] == 0)
416        })
417    }
418
419    /// `true` if the partition rows and `hash`'s active rows are jointly GF(2)-independent.
420    pub fn is_independent_of(&self, hash: &Gf2Hash<W>) -> bool {
421        self.matrix.independent_with(self.matrix.num_rows(), &hash.matrix, hash.bits() as usize)
422    }
423}
424
425/// One row of the key space as `(x-masks, z-masks)`.
426type KeyRow<const W: usize> = ([u64; W], [u64; W]);
427
428/// GF(2) rank of rows over the key space, by Gaussian elimination.
429fn gf2_rank_wide<const W: usize>(rows: &[KeyRow<W>]) -> usize {
430    let mut pivots: Vec<(usize, KeyRow<W>)> = Vec::with_capacity(rows.len());
431    for row in rows {
432        let mut reduced = *row;
433        'reduce: loop {
434            let Some(lead) = leading_column(&reduced) else {
435                break;
436            };
437            for (pivot_column, pivot_row) in pivots.iter() {
438                if *pivot_column == lead {
439                    for w in 0..W {
440                        reduced.0[w] ^= pivot_row.0[w];
441                        reduced.1[w] ^= pivot_row.1[w];
442                    }
443                    continue 'reduce;
444                }
445            }
446            pivots.push((lead, reduced));
447            break;
448        }
449    }
450    pivots.len()
451}
452
453/// Index of the highest set column of a key-space row, or `None` if it is zero.
454fn leading_column<const W: usize>(row: &KeyRow<W>) -> Option<usize> {
455    for w in (0..W).rev() {
456        if row.1[w] != 0 {
457            return Some(W * 64 + w * 64 + (63 - row.1[w].leading_zeros() as usize));
458        }
459    }
460    for w in (0..W).rev() {
461        if row.0[w] != 0 {
462            return Some(w * 64 + (63 - row.0[w].leading_zeros() as usize));
463        }
464    }
465    None
466}
467
468#[cfg(test)]
469mod tests;
paulistrings/pauli_sum/hash.rs

paulistrings/pauli_sum/
hash.rs
