Include Lagrange kernel trace polynomial term into DEEP composition p…

…olynomial (#274)

air/src/air/coefficients.rs

@@ -72,17 +72,40 @@ pub struct LagrangeConstraintsCompositionCoefficients<E: FieldElement> {
 /// https://eprint.iacr.org/2022/1216 and it relies on two points:
 ///
 /// 1. The evaluation proofs for each trace polynomial at $z$ and $g \cdot z$ can be batched using
-/// the non-normalized Lagrange kernel over the set $\{z, g \cdot z\}$. This, however, requires
-/// that the FRI protocol is run with rate $\rho^{+} := \frac{\kappa + 2}{\nu}$ where $\kappa$ and
-/// $\nu$ are the length of the execution trace and the LDE domain size, respectively.
+///    the non-normalized Lagrange kernel over the set $\{z, g \cdot z\}$. This, however, requires
+///    that the FRI protocol is run with a larger agreement parameter
+///    $\alpha^{+} = (1 + 1/2m)\cdot\sqrt{\rho^{+}}$ where $\rho^{+} := \frac{\kappa + 2}{\nu}$,
+///    $\kappa$ and $\nu$ are the length of the execution trace and the LDE domain size,
+///    respectively.
 /// 2. The resulting $Y(x)$ do not need to be degree adjusted but the soundness error of the
-/// protocol needs to be updated. For most combinations of batching parameters, this leads to a
-/// negligible increase in soundness error. The formula for the updated error can be found in
-/// Theorem 8 of https://eprint.iacr.org/2022/1216.
+///    protocol needs to be updated. For most combinations of batching parameters, this leads to a
+///    negligible increase in soundness error. The formula for the updated error can be found in
+///    Theorem 8 of https://eprint.iacr.org/2022/1216.
+///
+/// In the case when the trace polynomials contain a trace polynomial corresponding to a Lagrange
+/// kernel column, the above expression of $Y(x)$ includes the additional term given by
+///
+/// $$
+/// \gamma \cdot \frac{T_l(x) - p_S(x)}{Z_S(x)}
+/// $$
+///
+/// where:
+///
+/// 1. $\gamma$ is the composition coefficient for the Lagrange kernel trace polynomial.
+/// 2. $T_l(x) is the evaluation of the Lagrange trace polynomial at $x$.
+/// 3. $S$ is the set of opening points for the Lagrange kernel i.e.,
+///    $S := {z, z.g, z.g^2, ..., z.g^{2^{log_2(\nu) - 1}}}$.
+/// 4. $p_S(X)$ is the polynomial of minimal degree interpolating the set ${(a, T_l(a)): a \in S}$.
+/// 5. $Z_S(X)$ is the polynomial of minimal degree vanishing over the set $S$.
+///
+/// Note that, if a Lagrange kernel trace polynomial is present, then $\rho^{+}$ from above should
+/// be updated to be $\rho^{+} := \frac{\kappa + log_2(\nu) + 1}{\nu}$.
 #[derive(Debug, Clone)]
 pub struct DeepCompositionCoefficients<E: FieldElement> {
     /// Trace polynomial composition coefficients $\alpha_i$.
     pub trace: Vec<E>,
     /// Constraint column polynomial composition coefficients $\beta_j$.
     pub constraints: Vec<E>,
+    /// Lagrange kernel trace polynomial composition coefficient $\gamma$.
+    pub lagrange: Option<E>,
 }

air/src/air/context.rs

@@ -124,7 +124,11 @@ impl<B: StarkField> AirContext<B> {
 
         // validate Lagrange kernel aux column, if any
         if let Some(lagrange_kernel_aux_column_idx) = lagrange_kernel_aux_column_idx {
-            assert!(lagrange_kernel_aux_column_idx < trace_info.get_aux_segment_width(), "Lagrange kernel column index out of bounds: index={}, but only {} columns in segment", lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width());
+            assert!(
+            lagrange_kernel_aux_column_idx == trace_info.get_aux_segment_width() - 1,
+            "Lagrange kernel column should be the last column of the auxiliary trace: index={}, but aux trace width is {}",
+            lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width()
+            );
         }
 
         // determine minimum blowup factor needed to evaluate transition constraints by taking

air/src/air/mod.rs

@@ -586,9 +586,16 @@ pub trait Air: Send + Sync {
             c_coefficients.push(public_coin.draw()?);
         }
 
+        let lagrange_cc = if self.context().has_lagrange_kernel_aux_column() {
+            Some(public_coin.draw()?)
+        } else {
+            None
+        };
+
         Ok(DeepCompositionCoefficients {
             trace: t_coefficients,
             constraints: c_coefficients,
+            lagrange: lagrange_cc,
         })
     }
 }

air/src/proof/ood_frame.rs

@@ -124,7 +124,21 @@ impl OodFrame {
         assert!(main_trace_width > 0, "trace width cannot be zero");
         assert!(num_evaluations > 0, "number of evaluations cannot be zero");
 
-        // Parse main and auxiliary trace evaluation frames. This does the reverse operation done in
+        // parse Lagrange kernel column trace, if any
+        let mut reader = SliceReader::new(&self.lagrange_kernel_trace_states);
+        let lagrange_kernel_frame_size = reader.read_u8()? as usize;
+        let lagrange_kernel_frame = if lagrange_kernel_frame_size > 0 {
+            let lagrange_kernel_trace = reader.read_many(lagrange_kernel_frame_size)?;
+
+            Some(LagrangeKernelEvaluationFrame::new(lagrange_kernel_trace))
+        } else {
+            None
+        };
+
+        // if there is a Lagrange kernel, we treat its associated entries separately above
+        let aux_trace_width = aux_trace_width - (lagrange_kernel_frame.is_some() as usize);
+
+        // parse main and auxiliary trace evaluation frames. This does the reverse operation done in
         // `set_trace_states()`.
         let (current_row, next_row) = {
             let mut reader = SliceReader::new(&self.trace_states);
@@ -146,17 +160,6 @@ impl OodFrame {
             (current_row, next_row)
         };
 
-        // parse Lagrange kernel column trace
-        let mut reader = SliceReader::new(&self.lagrange_kernel_trace_states);
-        let lagrange_kernel_frame_size = reader.read_u8()? as usize;
-        let lagrange_kernel_frame = if lagrange_kernel_frame_size > 0 {
-            let lagrange_kernel_trace = reader.read_many(lagrange_kernel_frame_size)?;
-
-            Some(LagrangeKernelEvaluationFrame::new(lagrange_kernel_trace))
-        } else {
-            None
-        };
-
         // parse the constraint evaluations
         let mut reader = SliceReader::new(&self.evaluations);
         let evaluations = reader.read_many(num_evaluations)?;

math/Cargo.toml

@@ -27,6 +27,10 @@ harness = false
 name = "polynom"
 harness = false
 
+[[bench]]
+name = "poly_div_lagrange"
+harness = false
+
 [features]
 concurrent = ["utils/concurrent", "std"]
 default = ["std"]

math/benches/poly_div_lagrange.rs

@@ -0,0 +1,101 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use rand_utils::{rand_value, rand_vector};
+use std::time::Duration;
+use winter_math::{
+    fields::{f64::BaseElement, QuadExtension},
+    polynom::{self, eval_many, syn_div_roots_in_place},
+    ExtensionOf, StarkField,
+};
+
+const TRACE_LENS: [usize; 4] = [2_usize.pow(16), 2_usize.pow(18), 2_usize.pow(20), 2_usize.pow(22)];
+
+fn polynomial_division_naive(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Naive polynomial division");
+    group.sample_size(10);
+    group.measurement_time(Duration::from_secs(20));
+
+    for &trace_len in TRACE_LENS.iter() {
+        group.bench_function(BenchmarkId::new("prover", trace_len), |b| {
+            let poly: Vec<QuadExtension<BaseElement>> = rand_vector(trace_len);
+            let z: QuadExtension<BaseElement> = rand_value();
+            let log_trace_len = trace_len.ilog2();
+            let g: BaseElement = BaseElement::get_root_of_unity(log_trace_len);
+            let mut xs = Vec::with_capacity(log_trace_len as usize + 1);
+
+            // push z
+            xs.push(z);
+
+            // compute the values (z * g), (z * g^2), (z * g^4), ..., (z * g^(2^(v-1)))
+            let mut g_exp = g;
+            for _ in 0..log_trace_len {
+                let x = z.mul_base(g_exp);
+                xs.push(x);
+                g_exp *= g_exp;
+            }
+            let ood_evaluations = eval_many(&poly, &xs);
+
+            let p_s = polynom::interpolate(&xs, &ood_evaluations, true);
+            let numerator = polynom::sub(&poly, &p_s);
+            let z_s = polynom::poly_from_roots(&xs);
+
+            b.iter_batched(
+                || {
+                    let numerator = numerator.clone();
+                    let z_s = z_s.clone();
+                    (numerator, z_s)
+                },
+                |(numerator, z_s)| polynom::div(&numerator, &z_s),
+                BatchSize::SmallInput,
+            )
+        });
+    }
+}
+
+fn polynomial_division_synthetic(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Synthetic polynomial division");
+    group.sample_size(10);
+    group.measurement_time(Duration::from_secs(20));
+
+    for &trace_len in TRACE_LENS.iter() {
+        group.bench_function(BenchmarkId::new("prover", trace_len), |b| {
+            let poly: Vec<QuadExtension<BaseElement>> = rand_vector(trace_len);
+            let z: QuadExtension<BaseElement> = rand_value();
+            let log_trace_len = trace_len.ilog2();
+            let g: BaseElement = BaseElement::get_root_of_unity(log_trace_len);
+            let mut xs = Vec::with_capacity(log_trace_len as usize + 1);
+
+            // push z
+            xs.push(z);
+
+            // compute the values (z * g), (z * g^2), (z * g^4), ..., (z * g^(2^(v-1)))
+            let mut g_exp = g;
+            for _ in 0..log_trace_len {
+                let x = z.mul_base(g_exp);
+                xs.push(x);
+                g_exp *= g_exp;
+            }
+            let ood_evaluations = eval_many(&poly, &xs);
+
+            let p_s = polynom::interpolate(&xs, &ood_evaluations, true);
+            let numerator = polynom::sub(&poly, &p_s);
+
+            b.iter_batched(
+                || {
+                    let numerator = numerator.clone();
+                    let xs = xs.clone();
+                    (numerator, xs)
+                },
+                |(mut numerator, xs)| syn_div_roots_in_place(&mut numerator, &xs),
+                BatchSize::SmallInput,
+            )
+        });
+    }
+}
+
+criterion_group!(poly_division, polynomial_division_naive, polynomial_division_synthetic);
+criterion_main!(poly_division);

math/src/polynom/mod.rs

@@ -114,7 +114,7 @@ where
 {
     debug_assert!(xs.len() == ys.len(), "number of X and Y coordinates must be the same");
 
-    let roots = get_zero_roots(xs);
+    let roots = poly_from_roots(xs);
     let numerators: Vec<Vec<E>> = xs.iter().map(|&x| syn_div(&roots, 1, x)).collect();
 
     let denominators: Vec<E> = numerators.iter().zip(xs).map(|(e, &x)| eval(e, x)).collect();
@@ -549,6 +549,64 @@ where
     }
 }
 
+/// Divides a polynomial by a polynomial given its roots and saves the result into the original
+/// polynomial.
+///
+/// Specifically, divides polynomial `p` by polynomial \prod_{i = 1}^m (x - `x_i`) using
+/// [synthetic division](https://en.wikipedia.org/wiki/Synthetic_division) method and saves the
+/// result into `p`. If the polynomials don't divide evenly, the remainder is ignored. Polynomial
+/// `p` is expected to be in the coefficient form, and the result will be in coefficient form as
+/// well.
+///
+/// This function is significantly faster than the generic `polynom::div()` function, using
+/// the coefficients of the divisor.
+///
+/// # Panics
+/// Panics if:
+/// * `roots.len()` is zero;
+/// * `p.len()` is smaller than or equal to `roots.len()`.
+///
+/// # Examples
+/// ```
+/// # use winter_math::polynom::*;
+/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
+/// // p(x) = x^3 - 7 * x + 6
+/// let mut p = [
+///     BaseElement::new(6),
+///     -BaseElement::new(7),
+///     BaseElement::new(0),
+///     BaseElement::new(1),
+/// ];
+///
+/// // divide by (x - 1) * (x - 2)
+/// let zeros = vec![BaseElement::new(1), BaseElement::new(2)];
+/// syn_div_roots_in_place(&mut p, &zeros);
+///
+/// // expected result = x + 3
+/// let expected = [
+///     BaseElement::new(3),
+///     BaseElement::new(1),
+///     BaseElement::ZERO,
+///     BaseElement::ZERO,
+/// ];
+///
+/// assert_eq!(expected, p);
+pub fn syn_div_roots_in_place<E>(p: &mut [E], roots: &[E])
+where
+    E: FieldElement,
+{
+    assert!(!roots.is_empty(), "divisor should contain at least one linear factor");
+    assert!(p.len() > roots.len(), "divisor degree cannot be greater than dividend size");
+
+    for root in roots {
+        let mut c = E::ZERO;
+        for coeff in p.iter_mut().rev() {
+            *coeff += *root * c;
+            mem::swap(coeff, &mut c);
+        }
+    }
+}
+
 // DEGREE INFERENCE
 // ================================================================================================
 
@@ -621,14 +679,35 @@ where
     vec![]
 }
 
-// HELPER FUNCTIONS
-// ================================================================================================
-fn get_zero_roots<E: FieldElement>(xs: &[E]) -> Vec<E> {
+/// Returns the coefficients of polynomial given its roots.
+///
+/// # Examples
+/// ```
+/// # use winter_math::polynom::*;
+/// # use winter_math::{fields::{f128::BaseElement}, FieldElement};
+/// let xs = vec![1u128, 2]
+///     .into_iter()
+///     .map(BaseElement::new)
+///     .collect::<Vec<_>>();
+///
+/// let mut expected_poly = vec![2u128, 3, 1]
+///     .into_iter()
+///     .map(BaseElement::new)
+///     .collect::<Vec<_>>();
+/// expected_poly[1] *= -BaseElement::ONE;
+///
+/// let poly = poly_from_roots(&xs);
+/// assert_eq!(expected_poly, poly);
+/// ```
+pub fn poly_from_roots<E: FieldElement>(xs: &[E]) -> Vec<E> {
     let mut result = unsafe { utils::uninit_vector(xs.len() + 1) };
     fill_zero_roots(xs, &mut result);
     result
 }
 
+// HELPER FUNCTIONS
+// ================================================================================================
+
 fn fill_zero_roots<E: FieldElement>(xs: &[E], result: &mut [E]) {
     let mut n = result.len();
     n -= 1;