Add get lde functions (#100)

I have to revert #91 to align the value domains. --------- Co-authored-by: Kapil <[email protected]>
0xmozak · Jul 18, 2024 · d6cc0b9 · d6cc0b9
1 parent 3a545c1
commit d6cc0b9
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 40 deletions.
diff --git a/plonky2/src/fri/batch_oracle.rs b/plonky2/src/fri/batch_oracle.rs
@@ -4,6 +4,7 @@ use alloc::{format, vec::Vec};
 use itertools::Itertools;
 use plonky2_field::extension::Extendable;
 use plonky2_field::fft::FftRootTable;
+use plonky2_field::packed::PackedField;
 use plonky2_field::polynomial::{PolynomialCoeffs, PolynomialValues};
 use plonky2_field::types::Field;
 use plonky2_maybe_rayon::*;
@@ -21,7 +22,7 @@ use crate::plonk::config::GenericConfig;
 use crate::timed;
 use crate::util::reducing::ReducingFactor;
 use crate::util::timing::TimingTree;
-use crate::util::transpose;
+use crate::util::{reverse_bits, transpose};
 
 /// Represents a batch FRI oracle, i.e. a batch of polynomials with different degrees which have
 /// been Merkle-ized in a Field Merkle Tree.
@@ -30,7 +31,7 @@ pub struct BatchFriOracle<F: RichField + Extendable<D>, C: GenericConfig<D, F =
 {
     pub polynomials: Vec<PolynomialCoeffs<F>>,
     pub field_merkle_tree: FieldMerkleTree<F, C::Hasher>,
-    // The degree bits of each polynomial.
+    // The degree bits of each polynomial group.
     pub degree_bits: Vec<usize>,
     pub rate_bits: usize,
     pub blinding: bool,
@@ -73,17 +74,15 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
         timing: &mut TimingTree,
         fft_root_table: &[Option<&FftRootTable<F>>],
     ) -> Self {
-        let degree_bits = polynomials
+        let mut degree_bits = polynomials
             .iter()
             .map(|p| log2_strict(p.len()))
             .collect_vec();
         assert!(degree_bits.windows(2).all(|pair| { pair[0] >= pair[1] }));
-        let max_degree_bits = degree_bits[0];
 
         let num_polynomials = polynomials.len();
         let mut group_start = 0;
         let mut leaves = Vec::new();
-        let shift = F::coset_shift();
 
         for (i, d) in degree_bits.iter().enumerate() {
             if i == num_polynomials - 1 || *d > degree_bits[i + 1] {
@@ -93,7 +92,6 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
                     PolynomialBatch::<F, C, D>::lde_values(
                         &polynomials[group_start..i + 1],
                         rate_bits,
-                        shift.exp_power_of_2(max_degree_bits - d),
                         blinding,
                         fft_root_table[i]
                     )
@@ -113,6 +111,10 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             FieldMerkleTree::new(leaves, cap_height)
         );
 
+        degree_bits.sort_unstable();
+        degree_bits.dedup();
+        degree_bits.reverse();
+        assert_eq!(field_merkle_tree.leaves.len(), degree_bits.len());
         Self {
             polynomials,
             field_merkle_tree,
@@ -170,11 +172,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             let lde_final_values = timed!(
                 timing,
                 &format!("perform final FFT {}", lde_final_poly.len()),
-                lde_final_poly.coset_fft(
-                    F::coset_shift()
-                        .exp_u64(1 << (degree_bits[0] - degree_bits[i]))
-                        .into()
-                )
+                lde_final_poly.coset_fft(F::coset_shift().into())
             );
             final_lde_polynomial_coeff.push(lde_final_poly);
             final_lde_polynomial_values.push(lde_final_values);
@@ -185,13 +183,69 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
                 .iter()
                 .map(|o| &o.field_merkle_tree)
                 .collect::<Vec<_>>(),
-            &final_lde_polynomial_coeff,
+            final_lde_polynomial_coeff[0].clone(),
             &final_lde_polynomial_values,
             challenger,
             fri_params,
             timing,
         )
     }
+
+    /// Fetches LDE values at the `index * step`th point.
+    pub fn get_lde_values(
+        &self,
+        degree_bits_index: usize,
+        index: usize,
+        step: usize,
+        slice_start: usize,
+        slice_len: usize,
+    ) -> &[F] {
+        let index = index * step;
+        let index = reverse_bits(index, self.degree_bits[degree_bits_index] + self.rate_bits);
+        let slice = &self.field_merkle_tree.leaves[degree_bits_index][index];
+        &slice[slice_start..slice_start + slice_len]
+    }
+
+    /// Like `get_lde_values`, but fetches LDE values from a batch of `P::WIDTH` points, and returns
+    /// packed values.
+    pub fn get_lde_values_packed<P>(
+        &self,
+        degree_bits_index: usize,
+        index_start: usize,
+        step: usize,
+        slice_start: usize,
+        slice_len: usize,
+    ) -> Vec<P>
+    where
+        P: PackedField<Scalar = F>,
+    {
+        let row_wise = (0..P::WIDTH)
+            .map(|i| {
+                self.get_lde_values(
+                    degree_bits_index,
+                    index_start + i,
+                    step,
+                    slice_start,
+                    slice_len,
+                )
+            })
+            .collect_vec();
+
+        // This is essentially a transpose, but we will not use the generic transpose method as we
+        // want inner lists to be of type P, not Vecs which would involve allocation.
+        let leaf_size = row_wise[0].len();
+        (0..leaf_size)
+            .map(|j| {
+                let mut packed = P::ZEROS;
+                packed
+                    .as_slice_mut()
+                    .iter_mut()
+                    .zip(&row_wise)
+                    .for_each(|(packed_i, row_i)| *packed_i = row_i[j]);
+                packed
+            })
+            .collect_vec()
+    }
 }
 
 #[cfg(test)]

diff --git a/plonky2/src/fri/batch_prover.rs b/plonky2/src/fri/batch_prover.rs
@@ -1,8 +1,9 @@
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 
-use itertools::zip_eq;
 use plonky2_field::extension::flatten;
+#[allow(unused_imports)]
+use plonky2_field::types::Field;
 use plonky2_maybe_rayon::*;
 use plonky2_util::reverse_index_bits_in_place;
 
@@ -23,13 +24,13 @@ use crate::util::timing::TimingTree;
 /// Builds a batch FRI proof.
 pub fn batch_fri_proof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
     initial_merkle_trees: &[&FieldMerkleTree<F, C::Hasher>],
-    lde_polynomial_coeffs: &[PolynomialCoeffs<F::Extension>],
+    lde_polynomial_coeffs: PolynomialCoeffs<F::Extension>,
     lde_polynomial_values: &[PolynomialValues<F::Extension>],
     challenger: &mut Challenger<F, C::Hasher>,
     fri_params: &FriParams,
     timing: &mut TimingTree,
 ) -> FriProof<F, C::Hasher, D> {
-    let n = lde_polynomial_coeffs[0].len();
+    let n = lde_polynomial_coeffs.len();
     assert_eq!(lde_polynomial_values[0].len(), n);
     // The polynomial vectors should be sorted by degree, from largest to smallest, with no duplicate degrees.
     assert!(lde_polynomial_values
@@ -77,7 +78,7 @@ pub(crate) fn batch_fri_committed_trees<
     C: GenericConfig<D, F = F>,
     const D: usize,
 >(
-    coeffs: &[PolynomialCoeffs<F::Extension>],
+    mut final_coeffs: PolynomialCoeffs<F::Extension>,
     values: &[PolynomialValues<F::Extension>],
     challenger: &mut Challenger<F, C::Hasher>,
     fri_params: &FriParams,
@@ -86,7 +87,6 @@ pub(crate) fn batch_fri_committed_trees<
     let mut shift = F::MULTIPLICATIVE_GROUP_GENERATOR;
     let mut polynomial_index = 1;
     let mut final_values = values[0].clone();
-    let mut final_coeffs = coeffs[0].clone();
     for arity_bits in &fri_params.reduction_arity_bits {
         let arity = 1 << arity_bits;
 
@@ -112,17 +112,16 @@ pub(crate) fn batch_fri_committed_trees<
         if polynomial_index != values.len() && final_values.len() == values[polynomial_index].len()
         {
             final_values = PolynomialValues::new(
-                zip_eq(final_values.values.iter(), &values[polynomial_index].values)
-                    .map(|(&f, &v)| f * beta + v)
-                    .collect::<Vec<_>>(),
-            );
-            final_coeffs = PolynomialCoeffs::new(
-                zip_eq(final_coeffs.coeffs.iter(), &coeffs[polynomial_index].coeffs)
+                final_values
+                    .values
+                    .iter()
+                    .zip(&values[polynomial_index].values)
                     .map(|(&f, &v)| f * beta + v)
                     .collect::<Vec<_>>(),
             );
             polynomial_index += 1;
         }
+        final_coeffs = final_values.clone().coset_ifft(shift.into());
     }
     assert_eq!(polynomial_index, values.len());
 
@@ -216,7 +215,6 @@ mod tests {
     use plonky2_field::types::{Field64, Sample};
 
     use super::*;
-    use crate::field::types::Field;
     use crate::fri::batch_oracle::BatchFriOracle;
     use crate::fri::batch_verifier::verify_batch_fri_proof;
     use crate::fri::reduction_strategies::FriReductionStrategy;
@@ -295,7 +293,7 @@ mod tests {
 
         let proof = batch_fri_proof::<F, C, D>(
             &[&polynomial_batch.field_merkle_tree],
-            &[lde_final_poly],
+            lde_final_poly,
             &[lde_final_values],
             &mut challenger,
             &fri_params,
@@ -387,19 +385,17 @@ mod tests {
         let mut quotient = composition_poly.divide_by_linear(zeta);
         quotient.coeffs.push(<F as Extendable<D>>::Extension::ZERO);
         let lde_final_poly_1 = quotient.lde(fri_params.config.rate_bits);
-        let lde_final_values_1 =
-            lde_final_poly_1.coset_fft(F::coset_shift().exp_u64(1 << (k0 - k1)).into());
+        let lde_final_values_1 = lde_final_poly_1.coset_fft(F::coset_shift().into());
 
         let composition_poly = poly2.mul_extension::<D>(<F as Extendable<D>>::Extension::ONE);
         let mut quotient = composition_poly.divide_by_linear(zeta);
         quotient.coeffs.push(<F as Extendable<D>>::Extension::ZERO);
         let lde_final_poly_2 = quotient.lde(fri_params.config.rate_bits);
-        let lde_final_values_2 =
-            lde_final_poly_2.coset_fft(F::coset_shift().exp_u64(1 << (k0 - k2)).into());
+        let lde_final_values_2 = lde_final_poly_2.coset_fft(F::coset_shift().into());
 
         let proof = batch_fri_proof::<F, C, D>(
             &[&trace_oracle.field_merkle_tree],
-            &[lde_final_poly_0, lde_final_poly_1, lde_final_poly_2],
+            lde_final_poly_0,
             &[lde_final_values_0, lde_final_values_1, lde_final_values_2],
             &mut challenger,
             &fri_params,

diff --git a/plonky2/src/fri/batch_verifier.rs b/plonky2/src/fri/batch_verifier.rs
@@ -221,12 +221,14 @@ fn batch_fri_verifier_query_round<
         if batch_index < degree_bits.len()
             && n == degree_bits[batch_index] + params.config.rate_bits
         {
+            let subgroup_x_init = F::MULTIPLICATIVE_GROUP_GENERATOR
+                * F::primitive_root_of_unity(n).exp_u64(reverse_bits(x_index, n) as u64);
             let eval = batch_fri_combine_initial::<F, C, D>(
                 instances,
                 batch_index,
                 &round_proof.initial_trees_proof,
                 challenges.fri_alpha,
-                subgroup_x,
+                subgroup_x_init,
                 &precomputed_reduced_evals[batch_index],
                 params,
             );

diff --git a/plonky2/src/fri/oracle.rs b/plonky2/src/fri/oracle.rs
@@ -91,13 +91,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
         let lde_values = timed!(
             timing,
             "FFT + blinding",
-            Self::lde_values(
-                &polynomials,
-                rate_bits,
-                F::coset_shift(),
-                blinding,
-                fft_root_table
-            )
+            Self::lde_values(&polynomials, rate_bits, blinding, fft_root_table)
         );
 
         let mut leaves = timed!(timing, "transpose LDEs", transpose(&lde_values));
@@ -120,7 +114,6 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
     pub(crate) fn lde_values(
         polynomials: &[PolynomialCoeffs<F>],
         rate_bits: usize,
-        shift: F,
         blinding: bool,
         fft_root_table: Option<&FftRootTable<F>>,
     ) -> Vec<Vec<F>> {
@@ -134,7 +127,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             .map(|p| {
                 assert_eq!(p.len(), degree, "Polynomial degrees inconsistent");
                 p.lde(rate_bits)
-                    .coset_fft_with_options(shift, Some(rate_bits), fft_root_table)
+                    .coset_fft_with_options(F::coset_shift(), Some(rate_bits), fft_root_table)
                     .values
             })
             .chain(