rhizomes/verifier: Faster polynomial evaluation on flp.Query

Cargo.toml

@@ -55,6 +55,7 @@ experimental = ["bitvec", "fiat-crypto", "fixed", "num-bigint", "num-rational",
 multithreaded = ["rayon"]
 crypto-dependencies = ["aes", "ctr", "hmac", "sha2"]
 test-util = ["hex", "serde_json", "rand_distr"]
+rhizomes = ["experimental"]
 
 [workspace]
 members = [".", "binaries"]

Makefile

@@ -0,0 +1,16 @@
+all: test
+
+test:
+	cargo test --features experimental
+	cargo test --features rhizomes
+
+bench:
+	cargo bench --features experimental --bench speed_tests -- "prio3.*prepare_init" --quiet --save-baseline baseline
+	cargo bench --features rhizomes     --bench speed_tests -- "prio3.*prepare_init" --quiet --save-baseline rhizomes
+	critcmp baseline rhizomes
+
+graph:
+	critcmp --export baseline > baseline.json
+	critcmp --export rhizomes > rhizomes.json
+	python3 graph.py baseline.json rhizomes.json comparison.png
+	rm baseline.json rhizomes.json

README.md

@@ -1,3 +1,32 @@
+# Rhizomes and the Roots of Efficiency — Improving Prio
+
+This project is a fork of [divviup/libprio-rs](https://github.com/divviup/libprio-rs), a rust implementation of Prio.
+
+**Improvements**
+
+- Speeds up Polynomial Evaluation in the Lagrange basis.
+- Polynomials Basis Extension in the Lagrange basis.
+- Use of the Pólya polynomial basis.
+- Reduces the number of NTTs.
+
+## Benchmarks
+
+Invoke make for benchmarking the code:
+
+```sh
+make bench
+```
+
+Optionally, generate a graph to compare timings.
+
+```sh
+make graph
+```
+
+![graph comparison](comparison.png)
+
+---
+
 # libprio-rs
 [![Latest Version]][crates.io] [![Docs badge]][docs.rs]
 

benches/speed_tests.rs

@@ -281,7 +281,14 @@ fn prio3(c: &mut Criterion) {
     group.finish();
 
     let mut group = c.benchmark_group("prio3sumvec_prepare_init");
-    for (input_length, chunk_length) in [(10, 3), (100, 10), (1_000, 31)] {
+    for (input_length, chunk_length) in [
+        (10, 3),
+        (100, 10),
+        (1_000, 31),
+        (10_000, 100),
+        (100_000, 316),
+        (1_000_000, 1_000),
+    ] {
         group.bench_with_input(
             BenchmarkId::new("serial", input_length),
             &(input_length, chunk_length),
@@ -409,6 +416,7 @@ fn prio3(c: &mut Criterion) {
         (1_000, 31),
         (10_000, 100),
         (100_000, 316),
+        (1_000_000, 1_000),
     ] {
         if input_length >= 100_000 {
             group.measurement_time(Duration::from_secs(15));
@@ -529,7 +537,14 @@ fn prio3(c: &mut Criterion) {
     group.finish();
 
     let mut group = c.benchmark_group("prio3multihotcountvec_prepare_init");
-    for (input_length, chunk_length) in [(10, 3), (100, 10), (1_000, 31)] {
+    for (input_length, chunk_length) in [
+        (10, 3),
+        (100, 10),
+        (1_000, 31),
+        (10_000, 100),
+        (100_000, 316),
+        (1_000_000, 1_000),
+    ] {
         group.bench_with_input(
             BenchmarkId::new("serial", input_length),
             &(input_length, chunk_length),

graph.py

@@ -0,0 +1,43 @@
+import json
+import matplotlib.pyplot as plt
+import numpy as np
+import sys
+
+if len(sys.argv) == 4:
+    input0 = sys.argv[1]
+    input1 = sys.argv[2]
+    output = sys.argv[3]
+else:
+    print("Usage:\n\tpython3 plot.py <first.json> <second.json> <image.png>")
+
+with open(input0) as file:
+    dataset0 = json.load(file)
+with open(input1) as file:
+    dataset1 = json.load(file)
+
+name0 = dataset0["name"]
+name1 = dataset1["name"]
+cat0 = [n for n in dataset0["benchmarks"]]
+cat1 = [n for n in dataset1["benchmarks"]]
+categories = sorted(list(set(cat0) & set(cat1)))[::-1]
+
+values0 = []
+values1 = []
+for c in categories:
+    t0 = dataset0["benchmarks"][c]["criterion_estimates_v1"]["mean"]["point_estimate"]
+    t1 = dataset1["benchmarks"][c]["criterion_estimates_v1"]["mean"]["point_estimate"]
+    values0 += [t0 / t0]
+    values1 += [t0 / t1]
+
+y_pos = np.arange(len(categories))
+height = 0.3
+plt.barh(y_pos - height / 2, values0, height=height, label=name0)
+plt.barh(y_pos + height / 2, values1, height=height, label=name1)
+plt.xlabel("Speedup")
+plt.title(f"Speedup {name0} vs {name1}")
+plt.yticks(y_pos, categories)
+plt.legend()
+fig = plt.gcf()
+fig.set_size_inches(10, 12)
+fig.tight_layout()
+fig.savefig(output, dpi=200)

src/field.rs

@@ -109,6 +109,9 @@ pub trait FieldElement:
     /// Modular inversion, i.e., `self^-1 (mod p)`. If `self` is 0, then the output is undefined.
     fn inv(&self) -> Self;
 
+    /// Returns 1/2.
+    fn half() -> Self;
+
     /// Interprets the next [`Self::ENCODED_SIZE`] bytes from the input slice as an element of the
     /// field. Any of the most significant bits beyond the bit length of the modulus will be
     /// cleared, in order to minimize the amount of rejection sampling needed.
@@ -740,6 +743,10 @@ macro_rules! make_field {
             fn one() -> Self {
                 Self($fp::ROOTS[0])
             }
+
+            fn half() -> Self {
+                Self($fp::HALF)
+            }
         }
 
         impl FieldElementWithInteger for $elem {
@@ -1001,6 +1008,7 @@ pub(crate) mod test_utils {
         let int_one = F::TestInteger::try_from(1).unwrap();
         let zero = F::zero();
         let one = F::one();
+        let half = F::half();
         let two = F::from(F::TestInteger::try_from(2).unwrap());
         let four = F::from(F::TestInteger::try_from(4).unwrap());
 
@@ -1045,6 +1053,7 @@ pub(crate) mod test_utils {
         // mul
         assert_eq!(two * two, four);
         assert_eq!(two * one, two);
+        assert_eq!(two * half, one);
         assert_eq!(two * zero, zero);
         assert_eq!(one * F::from(int_modulus.clone()), zero);
 

src/field/field255.rs

@@ -309,6 +309,16 @@ impl FieldElement for Field255 {
     fn one() -> Self {
         Field255(fiat_25519_tight_field_element([1, 0, 0, 0, 0]))
     }
+
+    fn half() -> Self {
+        Field255(fiat_25519_tight_field_element([
+            2251799813685239,
+            2251799813685247,
+            2251799813685247,
+            2251799813685247,
+            1125899906842623,
+        ]))
+    }
 }
 
 impl Default for Field255 {

src/flp.rs

@@ -466,15 +466,28 @@ pub trait Flp: Sized + Eq + Clone + Debug {
             // Reconstruct the wire polynomials `f[0], ..., f[g_arity-1]` and evaluate each wire
             // polynomial at query randomness value.
             let m = (1 + gadget.calls()).next_power_of_two();
-            let m_inv = Self::Field::from(
-                <Self::Field as FieldElementWithInteger>::Integer::try_from(m).unwrap(),
-            )
-            .inv();
-            let mut f = vec![Self::Field::zero(); m];
-            for wire in 0..gadget.arity() {
-                ntt(&mut f, &gadget.f_vals[wire], m)?;
-                ntt_inv_finish(&mut f, m, m_inv);
-                verifier.push(poly_eval(&f, *query_rand_val));
+            #[cfg(not(feature = "rhizomes"))]
+            {
+                let m_inv = Self::Field::from(
+                    <Self::Field as FieldElementWithInteger>::Integer::try_from(m).unwrap(),
+                )
+                .inv();
+                let mut f = vec![Self::Field::zero(); m];
+                for wire in 0..gadget.arity() {
+                    ntt(&mut f, &gadget.f_vals[wire], m)?;
+                    ntt_inv_finish(&mut f, m, m_inv);
+                    verifier.push(poly_eval(&f, *query_rand_val));
+                }
+            }
+            #[cfg(feature = "rhizomes")]
+            {
+                // Evaluates a batch of polynomials in the Lagrange basis.
+                // This avoids using NTTs to convert them to the monomial basis.
+                use crate::rhizomes::{nth_root_powers, poly_eval_rhizomes_batched};
+                let roots = nth_root_powers(m);
+                let polynomials = &gadget.f_vals[..gadget.arity()];
+                let mut evals = poly_eval_rhizomes_batched(polynomials, &roots, *query_rand_val);
+                verifier.append(&mut evals);
             }
 
             // Add the value of the gadget polynomial evaluated at the query randomness value.

src/fp.rs

@@ -30,6 +30,7 @@ impl FieldParameters<u32> for FP32 {
         1534972560, 3732920810, 3229320047, 2836564014, 2170197442, 3760663902, 2144268387,
         3849278021, 1395394315, 574397626, 125025876, 3755041587, 2680072542, 3903828692,
     ];
+    const HALF: u32 = 2147483648;
     #[cfg(test)]
     const LOG2_BASE: usize = 32;
     #[cfg(test)]
@@ -72,6 +73,7 @@ impl FieldParameters<u64> for FP64 {
         10135969988448727187,
         6815045114074884550,
     ];
+    const HALF: u64 = 9223372036854775808;
     #[cfg(test)]
     const LOG2_BASE: usize = 64;
     #[cfg(test)]
@@ -114,6 +116,7 @@ impl FieldParameters<u128> for FP128 {
         258279638927684931537542082169183965856,
         148221243758794364405224645520862378432,
     ];
+    const HALF: u128 = 170141183460469231731687303715884105728;
     #[cfg(test)]
     const LOG2_BASE: usize = 64;
     #[cfg(test)]

src/fp/ops.rs

@@ -64,6 +64,8 @@ pub trait FieldParameters<W: Word> {
     /// `ROOTS[l]` has order `2^l` in the multiplicative group.
     /// `ROOTS[0]` is equal to one by definition.
     const ROOTS: [W; MAX_ROOTS + 1];
+    /// The multiplicative inverse of 2.
+    const HALF: W;
     /// The log2(base) for the base used for multiprecision arithmetic.
     /// So, `LOG2_BASE ≤ 64` as processors have at most a 64-bit
     /// integer multiplier.

src/lib.rs

@@ -33,6 +33,8 @@ pub mod idpf;
 mod ntt;
 mod polynomial;
 mod prng;
+#[cfg(feature = "rhizomes")]
+mod rhizomes;
 pub mod topology;
 pub mod vdaf;
 #[cfg(all(feature = "crypto-dependencies", feature = "experimental"))]