Merge pull request #69 from mir-protocol/optimize_reductions

Optimize reductions by `alpha`

src/fri/verifier.rs

@@ -7,8 +7,9 @@ use crate::fri::FriConfig;
 use crate::hash::hash_n_to_1;
 use crate::merkle_proofs::verify_merkle_proof;
 use crate::plonk_challenger::Challenger;
-use crate::plonk_common::{reduce_with_iter, PlonkPolynomials};
+use crate::plonk_common::PlonkPolynomials;
 use crate::proof::{FriInitialTreeProof, FriProof, FriQueryRound, Hash, OpeningSet};
+use crate::util::scaling::ReducingFactor;
 use crate::util::{log2_strict, reverse_bits, reverse_index_bits_in_place};
 
 /// Computes P'(x^arity) from {P(x*g^i)}_(i=0..arity), where g is a `arity`-th root of unity
@@ -151,7 +152,7 @@ fn fri_combine_initial<F: Field + Extendable<D>, const D: usize>(
     assert!(D > 1, "Not implemented for D=1.");
     let degree_log = proof.evals_proofs[0].1.siblings.len() - config.rate_bits;
     let subgroup_x = F::Extension::from_basefield(subgroup_x);
-    let mut alpha_powers = alpha.powers();
+    let mut alpha = ReducingFactor::new(alpha);
     let mut sum = F::Extension::ZERO;
 
     // We will add three terms to `sum`:
@@ -173,30 +174,33 @@ fn fri_combine_initial<F: Field + Extendable<D>, const D: usize>(
         .iter()
         .chain(&os.plonk_s_sigmas)
         .chain(&os.quotient_polys);
-    let single_diffs = single_evals.zip(single_openings).map(|(e, &o)| e - o);
-    let single_numerator = reduce_with_iter(single_diffs, &mut alpha_powers);
+    let single_diffs = single_evals
+        .into_iter()
+        .zip(single_openings)
+        .map(|(e, &o)| e - o)
+        .collect::<Vec<_>>();
+    let single_numerator = alpha.reduce(single_diffs.iter());
     let single_denominator = subgroup_x - zeta;
     sum += single_numerator / single_denominator;
+    alpha.reset();
 
     // Polynomials opened at `x` and `g x`, i.e., the Zs polynomials.
     let zs_evals = proof
         .unsalted_evals(PlonkPolynomials::ZS)
         .iter()
         .map(|&e| F::Extension::from_basefield(e));
-    let zs_composition_eval = reduce_with_iter(zs_evals, alpha_powers.clone());
+    let zs_composition_eval = alpha.clone().reduce(zs_evals);
     let zeta_right = F::Extension::primitive_root_of_unity(degree_log) * zeta;
     let zs_interpol = interpolate2(
         [
-            (zeta, reduce_with_iter(&os.plonk_zs, alpha_powers.clone())),
-            (
-                zeta_right,
-                reduce_with_iter(&os.plonk_zs_right, &mut alpha_powers),
-            ),
+            (zeta, alpha.clone().reduce(os.plonk_zs.iter())),
+            (zeta_right, alpha.reduce(os.plonk_zs_right.iter())),
         ],
         subgroup_x,
     );
     let zs_numerator = zs_composition_eval - zs_interpol;
     let zs_denominator = (subgroup_x - zeta) * (subgroup_x - zeta_right);
+    sum = alpha.shift(sum);
     sum += zs_numerator / zs_denominator;
 
     // Polynomials opened at `x` and `x.frobenius()`, i.e., the wires polynomials.
@@ -204,19 +208,20 @@ fn fri_combine_initial<F: Field + Extendable<D>, const D: usize>(
         .unsalted_evals(PlonkPolynomials::WIRES)
         .iter()
         .map(|&e| F::Extension::from_basefield(e));
-    let wire_composition_eval = reduce_with_iter(wire_evals, alpha_powers.clone());
+    let wire_composition_eval = alpha.clone().reduce(wire_evals);
     let zeta_frob = zeta.frobenius();
-    let wire_eval = reduce_with_iter(&os.wires, alpha_powers.clone());
+    let mut alpha_frob = alpha.repeated_frobenius(D - 1);
+    let wire_eval = alpha.reduce(os.wires.iter());
     // We want to compute `sum a^i*phi(w_i)`, where `phi` denotes the Frobenius automorphism.
     // Since `phi^D=id` and `phi` is a field automorphism, we have the following equalities:
     // `sum a^i*phi(w_i) = sum phi(phi^(D-1)(a^i)*w_i) = phi(sum phi^(D-1)(a)^i*w_i)`
     // So we can compute the original sum using only one call to the `D-1`-repeated Frobenius of alpha,
     // and one call at the end of the sum.
-    let alpha_powers_frob = alpha_powers.repeated_frobenius(D - 1);
-    let wire_eval_frob = reduce_with_iter(&os.wires, alpha_powers_frob).frobenius();
+    let wire_eval_frob = alpha_frob.reduce(os.wires.iter()).frobenius();
     let wire_interpol = interpolate2([(zeta, wire_eval), (zeta_frob, wire_eval_frob)], subgroup_x);
     let wire_numerator = wire_composition_eval - wire_interpol;
     let wire_denominator = (subgroup_x - zeta) * (subgroup_x - zeta_frob);
+    sum = alpha.shift(sum);
     sum += wire_numerator / wire_denominator;
 
     sum

src/polynomial/commitment.rs

@@ -7,10 +7,11 @@ use crate::field::field::Field;
 use crate::fri::{prover::fri_proof, verifier::verify_fri_proof, FriConfig};
 use crate::merkle_tree::MerkleTree;
 use crate::plonk_challenger::Challenger;
-use crate::plonk_common::reduce_polys_with_iter;
+use crate::plonk_common::PlonkPolynomials;
 use crate::polynomial::polynomial::PolynomialCoeffs;
 use crate::proof::{FriProof, FriProofTarget, Hash, OpeningSet};
 use crate::timed;
+use crate::util::scaling::ReducingFactor;
 use crate::util::{log2_strict, reverse_index_bits_in_place, transpose};
 
 pub const SALT_SIZE: usize = 2;
@@ -109,36 +110,49 @@ impl<F: Field> ListPolynomialCommitment<F> {
         challenger.observe_opening_set(&os);
 
         let alpha = challenger.get_extension_challenge();
-        let mut alpha_powers = alpha.powers();
+        let mut alpha = ReducingFactor::new(alpha);
 
         // Final low-degree polynomial that goes into FRI.
         let mut final_poly = PolynomialCoeffs::empty();
 
         // Polynomials opened at a single point.
-        let single_polys = [0, 1, 4]
-            .iter()
-            .flat_map(|&i| &commitments[i].polynomials)
-            .map(|p| p.to_extension());
-        let single_composition_poly = reduce_polys_with_iter(single_polys, &mut alpha_powers);
+        let single_polys = [
+            PlonkPolynomials::CONSTANTS,
+            PlonkPolynomials::SIGMAS,
+            PlonkPolynomials::QUOTIENT,
+        ]
+        .iter()
+        .flat_map(|&p| &commitments[p.index].polynomials)
+        .map(|p| p.to_extension());
+        let single_composition_poly = alpha.reduce_polys(single_polys);
 
         let single_quotient = Self::compute_quotient([zeta], single_composition_poly);
         final_poly += single_quotient;
+        alpha.reset();
 
         // Zs polynomials are opened at `zeta` and `g*zeta`.
-        let zs_polys = commitments[3].polynomials.iter().map(|p| p.to_extension());
-        let zs_composition_poly = reduce_polys_with_iter(zs_polys, &mut alpha_powers);
+        let zs_polys = commitments[PlonkPolynomials::ZS.index]
+            .polynomials
+            .iter()
+            .map(|p| p.to_extension());
+        let zs_composition_poly = alpha.reduce_polys(zs_polys);
 
         let zs_quotient = Self::compute_quotient([zeta, g * zeta], zs_composition_poly);
+        alpha.shift_poly(&mut final_poly);
         final_poly += zs_quotient;
 
         // When working in an extension field, need to check that wires are in the base field.
         // Check this by opening the wires polynomials at `zeta` and `zeta.frobenius()` and using the fact that
         // a polynomial `f` is over the base field iff `f(z).frobenius()=f(z.frobenius())` with high probability.
-        let wire_polys = commitments[2].polynomials.iter().map(|p| p.to_extension());
-        let wire_composition_poly = reduce_polys_with_iter(wire_polys, &mut alpha_powers);
+        let wire_polys = commitments[PlonkPolynomials::WIRES.index]
+            .polynomials
+            .iter()
+            .map(|p| p.to_extension());
+        let wire_composition_poly = alpha.reduce_polys(wire_polys);
 
         let wires_quotient =
             Self::compute_quotient([zeta, zeta.frobenius()], wire_composition_poly);
+        alpha.shift_poly(&mut final_poly);
         final_poly += wires_quotient;
 
         let lde_final_poly = final_poly.lde(config.rate_bits);

src/polynomial/polynomial.rs

@@ -1,6 +1,6 @@
 use std::cmp::max;
 use std::iter::Sum;
-use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
+use std::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};
 
 use crate::field::extension_field::Extendable;
 use crate::field::fft::{fft, ifft};
@@ -253,6 +253,16 @@ impl<F: Field> AddAssign for PolynomialCoeffs<F> {
     }
 }
 
+impl<F: Field> AddAssign<&Self> for PolynomialCoeffs<F> {
+    fn add_assign(&mut self, rhs: &Self) {
+        let len = max(self.len(), rhs.len());
+        self.coeffs.resize(len, F::ZERO);
+        for (l, &r) in self.coeffs.iter_mut().zip(&rhs.coeffs) {
+            *l += r;
+        }
+    }
+}
+
 impl<F: Field> SubAssign for PolynomialCoeffs<F> {
     fn sub_assign(&mut self, rhs: Self) {
         let len = max(self.len(), rhs.len());
@@ -263,6 +273,16 @@ impl<F: Field> SubAssign for PolynomialCoeffs<F> {
     }
 }
 
+impl<F: Field> SubAssign<&Self> for PolynomialCoeffs<F> {
+    fn sub_assign(&mut self, rhs: &Self) {
+        let len = max(self.len(), rhs.len());
+        self.coeffs.resize(len, F::ZERO);
+        for (l, &r) in self.coeffs.iter_mut().zip(&rhs.coeffs) {
+            *l -= r;
+        }
+    }
+}
+
 impl<F: Field> Mul<F> for &PolynomialCoeffs<F> {
     type Output = PolynomialCoeffs<F>;
 
@@ -272,6 +292,12 @@ impl<F: Field> Mul<F> for &PolynomialCoeffs<F> {
     }
 }
 
+impl<F: Field> MulAssign<F> for PolynomialCoeffs<F> {
+    fn mul_assign(&mut self, rhs: F) {
+        self.coeffs.iter_mut().for_each(|x| *x *= rhs);
+    }
+}
+
 impl<F: Field> Mul for &PolynomialCoeffs<F> {
     type Output = PolynomialCoeffs<F>;
 

src/util/mod.rs

@@ -1,3 +1,4 @@
+pub mod scaling;
 pub(crate) mod timing;
 
 use crate::field::field::Field;

src/util/scaling.rs

@@ -0,0 +1,75 @@
+use std::borrow::Borrow;
+
+use crate::field::extension_field::Frobenius;
+use crate::field::field::Field;
+use crate::polynomial::polynomial::PolynomialCoeffs;
+
+/// When verifying the composition polynomial in FRI we have to compute sums of the form
+/// `(sum_0^k a^i * x_i)/d_0 + (sum_k^r a^i * y_i)/d_1`
+/// The most efficient way to do this is to compute both quotient separately using Horner's method,
+/// scale the second one by `a^(r-1-k)`, and add them up.
+/// This struct abstract away these operations by implementing Horner's method and keeping track
+/// of the number of multiplications by `a` to compute the scaling factor.
+/// See https://github.com/mir-protocol/plonky2/pull/69 for more details and discussions.
+#[derive(Debug, Copy, Clone)]
+pub struct ReducingFactor<F: Field> {
+    base: F,
+    count: u64,
+}
+
+impl<F: Field> ReducingFactor<F> {
+    pub fn new(base: F) -> Self {
+        Self { base, count: 0 }
+    }
+
+    fn mul(&mut self, x: F) -> F {
+        self.count += 1;
+        self.base * x
+    }
+
+    fn mul_poly(&mut self, p: &mut PolynomialCoeffs<F>) {
+        self.count += 1;
+        *p *= self.base;
+    }
+
+    pub fn reduce(&mut self, iter: impl DoubleEndedIterator<Item = impl Borrow<F>>) -> F {
+        iter.rev()
+            .fold(F::ZERO, |acc, x| self.mul(acc) + *x.borrow())
+    }
+
+    pub fn reduce_polys(
+        &mut self,
+        polys: impl DoubleEndedIterator<Item = impl Borrow<PolynomialCoeffs<F>>>,
+    ) -> PolynomialCoeffs<F> {
+        polys.rev().fold(PolynomialCoeffs::empty(), |mut acc, x| {
+            self.mul_poly(&mut acc);
+            acc += x.borrow();
+            acc
+        })
+    }
+
+    pub fn shift(&mut self, x: F) -> F {
+        let tmp = self.base.exp(self.count) * x;
+        self.count = 0;
+        tmp
+    }
+
+    pub fn shift_poly(&mut self, p: &mut PolynomialCoeffs<F>) {
+        *p *= self.base.exp(self.count);
+        self.count = 0;
+    }
+
+    pub fn reset(&mut self) {
+        self.count = 0;
+    }
+
+    pub fn repeated_frobenius<const D: usize>(&self, count: usize) -> Self
+    where
+        F: Frobenius<D>,
+    {
+        Self {
+            base: self.base.repeated_frobenius(count),
+            count: self.count,
+        }
+    }
+}