diff --git a/Cargo.toml b/Cargo.toml index d900e8a..947722b 100644 --- a/Cargo.toml +++ b/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"] diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..5214b6a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/README.md b/README.md index c3a705e..8c1a3a7 100644 --- a/README.md +++ b/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] diff --git a/benches/speed_tests.rs b/benches/speed_tests.rs index 3e33f6b..901b4b8 100644 --- a/benches/speed_tests.rs +++ b/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), diff --git a/comparison.png b/comparison.png new file mode 100644 index 0000000..f0c8130 Binary files /dev/null and b/comparison.png differ diff --git a/graph.py b/graph.py new file mode 100644 index 0000000..98cc4a2 --- /dev/null +++ b/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 ") + +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) \ No newline at end of file diff --git a/src/field.rs b/src/field.rs index 215f092..b845db7 100644 --- a/src/field.rs +++ b/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); diff --git a/src/field/field255.rs b/src/field/field255.rs index 8a3f74b..a6c1949 100644 --- a/src/field/field255.rs +++ b/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 { diff --git a/src/flp.rs b/src/flp.rs index ed413b7..5a3dd59 100644 --- a/src/flp.rs +++ b/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( - ::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( + ::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. diff --git a/src/fp.rs b/src/fp.rs index b530656..6fd955e 100644 --- a/src/fp.rs +++ b/src/fp.rs @@ -30,6 +30,7 @@ impl FieldParameters 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 for FP64 { 10135969988448727187, 6815045114074884550, ]; + const HALF: u64 = 9223372036854775808; #[cfg(test)] const LOG2_BASE: usize = 64; #[cfg(test)] @@ -114,6 +116,7 @@ impl FieldParameters for FP128 { 258279638927684931537542082169183965856, 148221243758794364405224645520862378432, ]; + const HALF: u128 = 170141183460469231731687303715884105728; #[cfg(test)] const LOG2_BASE: usize = 64; #[cfg(test)] diff --git a/src/fp/ops.rs b/src/fp/ops.rs index 87aedc7..4c641b5 100644 --- a/src/fp/ops.rs +++ b/src/fp/ops.rs @@ -64,6 +64,8 @@ pub trait FieldParameters { /// `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. diff --git a/src/lib.rs b/src/lib.rs index c57c031..d3c94f9 100644 --- a/src/lib.rs +++ b/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"))] diff --git a/src/rhizomes/mod.rs b/src/rhizomes/mod.rs new file mode 100644 index 0000000..96f6909 --- /dev/null +++ b/src/rhizomes/mod.rs @@ -0,0 +1,264 @@ +// Copyright (c) 2025 Armando Faz Hernandez. +// SPDX-License-Identifier: MPL-2.0 + +//! Faster algorithms for polynomials in the Lagrange basis. +//! +//! Reference: +//! Faz-Hernandez, "Rhizomes and the Roots of Efficiency -- Improving Prio." +//! + +use crate::{ + field::{FieldElement, NttFriendlyFieldElement}, + fp::log2, +}; + +/// Returns the element `1/(2^n)` on `F`. +#[inline] +fn half_power(n: usize) -> F { + let half = F::half(); + let mut x = F::one(); + for _ in 0..n { + x *= half + } + x +} + +/// Returns the element `1/n` on `F`, where `n` must be a power of two. +#[inline] +fn inv_pow2(n: usize) -> F { + let log2_n = usize::try_from(log2(n as u128)).unwrap(); + assert_eq!(n, 1 << log2_n); + + half_power(log2_n) +} + +/// Evaluates a polynomial given in the Lagrange basis. +/// +/// This is the implementation of Algorithm 6. +#[allow(dead_code)] +pub fn poly_eval_rhizomes(poly: &[F], roots: &[F], x: &F) -> F { + let n = poly.len(); + let mut l = F::one(); + let mut u = poly[0]; + let mut d = roots[0] - *x; + for (yi, wn_i) in poly[1..n].iter().zip(&roots[1..n]) { + l *= d; + d = *wn_i - *x; + u = u * d + l * *wn_i * *yi; + } + + for wn_i in &roots[n..] { + u *= *wn_i - *x; + } + + if roots.len() > 1 { + let num_roots_inv = -inv_pow2::(roots.len()); + u *= num_roots_inv; + } + + u +} + +/// Evaluates all polynomials given in the Lagrange basis. +/// +/// This is the implementation of Algorithm 7. +pub fn poly_eval_rhizomes_batched( + polynomials: &[Vec], + roots: &[F], + x: F, +) -> Vec { + let mut l = F::one(); + let mut u = Vec::with_capacity(polynomials.len()); + u.extend(polynomials.iter().map(|poly| poly[0])); + let mut d = roots[0] - x; + for (i, wn_i) in (1..).zip(&roots[1..]) { + l *= d; + d = *wn_i - x; + let t = l * *wn_i; + for (u_j, poly) in u.iter_mut().zip(polynomials) { + *u_j *= d; + if let Some(yi) = poly.get(i) { + *u_j += t * *yi; + } + } + } + + if roots.len() > 1 { + let num_roots_inv = -inv_pow2::(roots.len()); + u.iter_mut().for_each(|u_j| *u_j *= num_roots_inv); + } + + u +} + +/// Generates the powers of the primitive n-th root of unity. +/// +/// Returns +/// roots\[i\] = w_n^i for 0 ≤ i < n, +/// where +/// w_n is the primitive n-th root of unity in `F`, and +/// n must be a power of two. +pub fn nth_root_powers(n: usize) -> Vec { + let log2_n = usize::try_from(log2(n as u128)).unwrap(); + assert_eq!(n, 1 << log2_n); + + let mut roots = vec![F::zero(); n]; + roots[0] = F::one(); + if n > 1 { + roots[1] = -F::one(); + for i in 2..=log2_n { + let mid = 1 << (i - 1); + // Due to w_{2n}^{2j} = w_{n}^j + for j in (1..mid).rev() { + roots[j << 1] = roots[j] + } + + let wn = F::root(i).unwrap(); + roots[1] = wn; + roots[1 + mid] = -wn; + + // Due to w_{n}^{j} = -w_{n}^{j+n/2} + for j in (3..mid).step_by(2) { + roots[j] = wn * roots[j - 1]; + roots[j + mid] = -roots[j] + } + } + } + + roots +} + +#[cfg(test)] +mod test_methods { + use crate::{ + field::NttFriendlyFieldElement, + fp::log2, + ntt::{ntt, ntt_inv_finish}, + polynomial::poly_eval, + }; + + /// Evaluates a polynomial given in Lagrange basis. + /// + /// Converts the polynomial from the Lagrange to monomial basis (with inverse NTT), + /// and performs evaluation using the Horner's method. + pub(crate) fn poly_eval_monomial( + points: &[F], + eval_at: F, + tmp_coeffs: &mut [F], + size_inv: F, + ) -> F { + ntt(tmp_coeffs, points, points.len()).unwrap(); + ntt_inv_finish(tmp_coeffs, points.len(), size_inv); + poly_eval(&tmp_coeffs[..points.len()], eval_at) + } + + /// Generates the powers of the primitive n-th root of unity. + /// + /// Returns + /// roots\[i\] = w_n^i for 0 ≤ i < n, + /// where + /// w_n is the primitive n-th root of unity in `F`, and + /// n must be a power of two. + /// + /// This is the iterative method. + pub(crate) fn nth_root_powers_slow(n: usize) -> Vec { + let log2_n = usize::try_from(log2(n as u128)).unwrap(); + let wn = F::root(log2_n).unwrap(); + core::iter::successors(Some(F::one()), |&x| Some(x * wn)) + .take(n) + .collect() + } +} + +#[cfg(test)] +mod tests { + use crate::{ + field::{Field64 as Fp, FieldElement, FieldElementWithInteger}, + rhizomes::test_methods::{nth_root_powers_slow, poly_eval_monomial}, + rhizomes::{nth_root_powers, poly_eval_rhizomes, poly_eval_rhizomes_batched}, + }; + + #[test] + fn test_nth_root_powers() { + for i in 0..8 { + assert_eq!( + nth_root_powers::(1 << i), + nth_root_powers_slow::(1 << i) + ); + } + } + + #[test] + fn test_poly_eval_rhizomes() { + for size in (1usize..100).step_by(10) { + let n: usize = size.next_power_of_two(); + let n_inv = + Fp::from(::Integer::try_from(n).unwrap()).inv(); + let values = Fp::random_vector(n); + let x = Fp::random_vector(1)[0]; + let mut ntt_mem = vec![Fp::zero(); n]; + + // Evaluates a polynomial converting to the monomial basis. + let want = poly_eval_monomial(&values, x, &mut ntt_mem, n_inv); + + // Evaluates a polynomial directly in the Lagrange basis. + let roots = nth_root_powers(n); + let got = poly_eval_rhizomes(&values, &roots, &x); + assert_eq!(got, want, "n: {n} x: {x} values: {values:?}"); + } + } + + #[test] + fn test_poly_eval_batched_ones() { + test_poly_eval_batched(&[1]); + test_poly_eval_batched(&[1, 1]); + } + + #[test] + fn test_poly_eval_batched_powers() { + test_poly_eval_batched(&[1, 2, 4, 16, 64]); + } + + #[test] + fn test_poly_eval_batched_arbitrary() { + test_poly_eval_batched(&[1, 6, 3, 9]); + } + + fn test_poly_eval_batched(lengths: &[usize]) { + let sizes = lengths + .iter() + .map(|s| s.next_power_of_two()) + .collect::>(); + + let polynomials = sizes + .iter() + .map(|&size| Fp::random_vector(size)) + .collect::>(); + let x = Fp::random_vector(1)[0]; + + let &n = sizes.iter().max().unwrap(); + let n_inv = Fp::from(::Integer::try_from(n).unwrap()).inv(); + let mut ntt_mem = vec![Fp::zero(); n]; + let roots = nth_root_powers(n); + + // Evaluates several polynomials converting them to the monomial basis (iteratively). + let want = polynomials + .iter() + .map(|poly| { + let extended_poly = [poly.clone(), vec![Fp::zero(); n - poly.len()]].concat(); + poly_eval_monomial(&extended_poly, x, &mut ntt_mem, n_inv) + }) + .collect::>(); + + // Evaluates several polynomials directly in the Lagrange basis (iteratively). + let got = polynomials + .iter() + .map(|poly| poly_eval_rhizomes(poly, &roots, &x)) + .collect::>(); + assert_eq!(got, want, "sizes: {sizes:?} x: {x} P: {polynomials:?}"); + + // Simultaneouly evaluates several polynomials directly in the Lagrange basis (batched). + let got = poly_eval_rhizomes_batched(&polynomials, &roots, x); + assert_eq!(got, want, "sizes: {sizes:?} x: {x} P: {polynomials:?}"); + } +}