perf(ferret): ~5.7x faster COT extension (consistency check)

crates/core/benches/ggm.rs

@@ -8,8 +8,9 @@ fn criterion_benchmark(c: &mut Criterion) {
         let depth = 10;
         let seed = rand::random::<Block>();
         let mut leaves = vec![Block::ZERO; 1 << depth];
+        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
         bench.iter(|| {
-            GgmTree::new_from_seed(depth, seed, &mut leaves);
+            GgmTree::new_from_seed(depth, seed, &mut leaves, &mut buf);
             black_box(&leaves);
         });
     });
@@ -18,8 +19,9 @@ fn criterion_benchmark(c: &mut Criterion) {
         let depth = 10;
         let sums = vec![Block::ZERO; depth];
         let mut leaves = vec![Block::ZERO; 1 << depth];
+        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
         bench.iter(|| {
-            GgmTree::new_partial(depth, &sums, 420, &mut leaves);
+            GgmTree::new_partial(depth, &sums, 420, &mut leaves, &mut buf);
             black_box(&leaves);
         });
     });

crates/core/src/block.rs

@@ -100,21 +100,77 @@ impl Block {
 
     /// Compute the inner product of two block vectors, without reducing the
     /// polynomial.
+    ///
+    /// Uses 8 independent accumulators to break the carry-less multiply
+    /// latency-dependency chain (`clmul` is high-latency, fully-pipelined), so
+    /// the loop runs throughput-bound rather than latency-bound.
     #[inline]
     pub fn inn_prdt_no_red(a: &[Block], b: &[Block]) -> (Block, Block) {
         assert_eq!(a.len(), b.len());
-        a.iter()
-            .zip(b.iter())
-            .fold((Block::ZERO, Block::ZERO), |acc, (x, y)| {
-                let t = x.clmul(*y);
-                (t.0 ^ acc.0, t.1 ^ acc.1)
-            })
+
+        const LANES: usize = 8;
+        let mut hi = [Block::ZERO; LANES];
+        let mut lo = [Block::ZERO; LANES];
+
+        let mut a_chunks = a.chunks_exact(LANES);
+        let mut b_chunks = b.chunks_exact(LANES);
+        for (ac, bc) in a_chunks.by_ref().zip(b_chunks.by_ref()) {
+            for j in 0..LANES {
+                let (h, l) = ac[j].clmul(bc[j]);
+                hi[j] ^= h;
+                lo[j] ^= l;
+            }
+        }
+
+        let mut acc_hi = Block::ZERO;
+        let mut acc_lo = Block::ZERO;
+        for j in 0..LANES {
+            acc_hi ^= hi[j];
+            acc_lo ^= lo[j];
+        }
+
+        for (x, y) in a_chunks.remainder().iter().zip(b_chunks.remainder()) {
+            let (h, l) = x.clmul(*y);
+            acc_hi ^= h;
+            acc_lo ^= l;
+        }
+
+        (acc_hi, acc_lo)
     }
 
     /// Compute the inner product of two block vectors.
+    ///
+    /// With the `rayon` feature enabled, the (unreduced) inner product is
+    /// computed in parallel over chunks and combined before a single final
+    /// reduction.
     #[inline]
     pub fn inn_prdt_red(a: &[Block], b: &[Block]) -> Block {
-        let (x, y) = Block::inn_prdt_no_red(a, b);
+        assert_eq!(a.len(), b.len());
+
+        cfg_if::cfg_if! {
+            if #[cfg(feature = "rayon")] {
+                use rayon::prelude::*;
+
+                // Large enough that per-chunk overhead is negligible, small
+                // enough to keep all cores busy on production-sized vectors.
+                const CHUNK: usize = 1 << 16;
+
+                let (x, y) = if a.len() <= CHUNK {
+                    Block::inn_prdt_no_red(a, b)
+                } else {
+                    a.par_chunks(CHUNK)
+                        .zip(b.par_chunks(CHUNK))
+                        .map(|(ac, bc)| Block::inn_prdt_no_red(ac, bc))
+                        .reduce(
+                            || (Block::ZERO, Block::ZERO),
+                            |p, q| (p.0 ^ q.0, p.1 ^ q.1),
+                        )
+                };
+            } else {
+                let (x, y) = Block::inn_prdt_no_red(a, b);
+            }
+        }
+
         Block::reduce_gcm(x, y)
     }
 

crates/core/src/ggm.rs

@@ -24,7 +24,10 @@ fn width(n: usize) -> usize {
 /// GGM tree.
 pub struct GgmTree<'a> {
     depth: usize,
-    buf: Vec<Block>,
+    /// Internal (non-leaf) tree nodes. Caller-provided scratch of length
+    /// `(1 << depth) - 1`, so it can be reused across trees without
+    /// reallocating or re-zeroing.
+    buf: &'a mut [Block],
     leaves: &'a mut [Block],
 }
 
@@ -36,10 +39,17 @@ impl<'a> GgmTree<'a> {
     /// * `depth` - The depth of the tree.
     /// * `seed` - The seed of the tree.
     /// * `leaves` - The leaves of the tree.
-    pub fn new_from_seed(depth: usize, seed: Block, leaves: &'a mut [Block]) -> Self {
+    /// * `buf` - Scratch for the internal nodes, of length `(1 << depth) - 1`.
+    ///   Its contents are fully overwritten; reuse it across trees to avoid
+    ///   per-tree allocation.
+    pub fn new_from_seed(
+        depth: usize,
+        seed: Block,
+        leaves: &'a mut [Block],
+        buf: &'a mut [Block],
+    ) -> Self {
         assert_eq!(leaves.len(), 1 << depth, "invalid length of leaves");
-
-        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
+        assert_eq!(buf.len(), (1 << depth) - 1, "invalid length of buf");
 
         let tkprp = TwoKeyPrp::new([Block::ZERO, Block::ONE]);
 
@@ -76,12 +86,19 @@ impl<'a> GgmTree<'a> {
     ///   layer.
     /// * `idx` - Index of the missing leaf.
     /// * `leaves` - Leaves of the tree.
-    pub fn new_partial(depth: usize, sums: &[Block], idx: usize, leaves: &'a mut [Block]) -> Self {
+    /// * `buf` - Scratch for the internal nodes, of length `(1 << depth) - 1`.
+    ///   Reuse it across trees to avoid per-tree allocation.
+    pub fn new_partial(
+        depth: usize,
+        sums: &[Block],
+        idx: usize,
+        leaves: &'a mut [Block],
+        buf: &'a mut [Block],
+    ) -> Self {
         assert_eq!(leaves.len(), 1 << depth, "invalid length of leaves");
         assert!(idx < leaves.len(), "index out of bounds");
         assert_eq!(sums.len(), depth, "invalid length of sums");
-
-        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
+        assert_eq!(buf.len(), (1 << depth) - 1, "invalid length of buf");
 
         let tkprp = TwoKeyPrp::new([Block::ZERO, Block::ONE]);
 
@@ -189,8 +206,9 @@ mod tests {
         let depth = 4;
 
         let mut leaves = vec![Block::ZERO; 1 << depth];
+        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
 
-        GgmTree::new_from_seed(depth, seed, &mut leaves);
+        GgmTree::new_from_seed(depth, seed, &mut leaves, &mut buf);
 
         assert!(leaves.iter().all(|leaf| *leaf != Block::ZERO));
     }
@@ -201,8 +219,9 @@ mod tests {
         let depth = 4;
 
         let mut leaves = vec![Block::ZERO; 1 << depth];
+        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
 
-        let ggm = GgmTree::new_from_seed(depth, seed, &mut leaves);
+        let ggm = GgmTree::new_from_seed(depth, seed, &mut leaves, &mut buf);
 
         for i in 0..depth {
             let layer = ggm.layer(i).unwrap();
@@ -217,7 +236,8 @@ mod tests {
         let depth = 4;
 
         let mut full_leaves = vec![Block::ZERO; 1 << depth];
-        let ggm = GgmTree::new_from_seed(depth, seed, &mut full_leaves);
+        let mut buf = vec![Block::ZERO; (1 << depth) - 1];
+        let ggm = GgmTree::new_from_seed(depth, seed, &mut full_leaves, &mut buf);
 
         for i in 0..1 << depth {
             let path = i as u32;
@@ -228,7 +248,8 @@ mod tests {
                 .collect::<Vec<_>>();
 
             let mut leaves = vec![Block::ZERO; 1 << depth];
-            let ggm_partial = GgmTree::new_partial(depth, &sums, i, &mut leaves);
+            let mut partial_buf = vec![Block::ZERO; (1 << depth) - 1];
+            let ggm_partial = GgmTree::new_partial(depth, &sums, i, &mut leaves, &mut partial_buf);
             let mut full_leaves = ggm.leaves().to_vec();
 
             full_leaves[i] = Block::ZERO;