rhizomes/full: Both Prover and Verifier use polynomials in the Lagrange basis

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 -- --skip vec
+
+bench:
+	cargo bench --features experimental --bench speed_tests -- prio3 --quiet --save-baseline baseline
+	cargo bench --features rhizomes     --bench speed_tests -- prio3 --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

@@ -240,7 +240,14 @@ fn prio3(c: &mut Criterion) {
     group.finish();
 
     let mut group = c.benchmark_group("prio3sumvec_shard");
-    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),
@@ -281,7 +288,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),
@@ -355,6 +369,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));
@@ -409,6 +424,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));
@@ -484,7 +500,14 @@ fn prio3(c: &mut Criterion) {
     group.finish();
 
     let mut group = c.benchmark_group("prio3multihotcountvec_shard");
-    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),
@@ -529,7 +552,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

@@ -295,6 +295,7 @@ pub trait Flp: Sized + Eq + Clone + Debug {
             // Interpolate the wire polynomials `f[0], ..., f[g_arity-1]` from the input wires of each
             // evaluation of the gadget.
             let m = wire_poly_len(gadget.calls());
+            #[cfg(not(feature = "rhizomes"))]
             let m_inv = Self::Field::from(
                 <Self::Field as FieldElementWithInteger>::Integer::try_from(m).unwrap(),
             )
@@ -305,8 +306,15 @@ pub trait Flp: Sized + Eq + Clone + Debug {
                 .zip(gadget.f_vals[..gadget.arity()].iter())
                 .zip(proof[proof_len..proof_len + gadget.arity()].iter_mut())
             {
-                ntt(coefficients, values, m)?;
-                ntt_inv_finish(coefficients, m, m_inv);
+                #[cfg(not(feature = "rhizomes"))]
+                {
+                    ntt(coefficients, values, m)?;
+                    ntt_inv_finish(coefficients, m, m_inv);
+                }
+                #[cfg(feature = "rhizomes")]
+                // Returns the polynomial in the Lagrange basis (i.e., the values) directly.
+                // This avoids inverse NTT for recovering the coefficients.
+                coefficients[..values.len()].copy_from_slice(values);
 
                 // The first point on each wire polynomial is a random value chosen by the prover. This
                 // point is stored in the proof so that the verifier can reconstruct the wire
@@ -466,15 +474,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.
@@ -733,14 +754,48 @@ impl<F: NttFriendlyFieldElement> QueryShimGadget<F> {
         // Evaluate the gadget polynomial at roots of unity.
         let size = p.next_power_of_two();
         let mut p_vals = vec![F::zero(); size];
+        #[cfg(not(feature = "rhizomes"))]
         ntt(&mut p_vals, &proof_data[gadget_arity..], size)?;
 
         // The step is used to compute the element of `p_val` that will be returned by a call to
         // the gadget.
         let step = (1 << (log2(p as u128) - log2(m as u128))) as usize;
 
         // Evaluate the gadget polynomial `p` at query randomness `r`.
+        #[cfg(not(feature = "rhizomes"))]
         let p_at_r = poly_eval(&proof_data[gadget_arity..], r);
+        #[cfg(feature = "rhizomes")]
+        let p_at_r = {
+            use crate::rhizomes::{
+                extend_dimension_double, extend_dimension_one, nth_root_powers, poly_eval_rhizomes,
+            };
+            // The Prover sent the gadget_poly in the Lagrange basis.
+            // However, some missing coordinates must be recovered for polynomial evaluation.
+            let gadget_poly = &proof_data[gadget_arity..];
+            let l = gadget_poly.len();
+            let n = l.next_power_of_two();
+            assert!(l <= n);
+
+            // Extending the dimension (one by one) to the closest power of two.
+            // This keeps the degree of polynomial unchanged.
+            let roots = nth_root_powers(n);
+            p_vals[..l].copy_from_slice(gadget_poly);
+            for k in l..n {
+                p_vals[k] = extend_dimension_one(&p_vals[..k], &roots);
+            }
+
+            // Evaluating the polynomial in the Lagrange basis.
+            let p_at_r = poly_eval_rhizomes(&p_vals[..n], &roots, &r);
+
+            // Calculate all the 'size' evaluations at the roots of unity.
+            let mut k = n;
+            while k < size {
+                extend_dimension_double(&mut p_vals, k);
+                k *= 2;
+            }
+
+            p_at_r
+        };
 
         Ok(Self {
             inner,
@@ -962,6 +1017,10 @@ pub mod test_utils {
                 mutated_proof[i] *= T::Field::from(
                     <T::Field as FieldElementWithInteger>::Integer::try_from(23).unwrap(),
                 );
+                // Ensures one element of the proof was mutated.
+                if mutated_proof[i] == proof[i] {
+                    mutated_proof[i] += T::Field::one();
+                }
                 let verifier = self
                     .flp
                     .query(self.input, &mutated_proof, &query_rand, &joint_rand, 1)