Skip to content

Commit

Permalink
Speed-up lde_onto_coset
Browse files Browse the repository at this point in the history
  • Loading branch information
Nashtare committed Jul 31, 2023
1 parent 880f39f commit a02dfc2
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 17 deletions.
28 changes: 22 additions & 6 deletions evm/src/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@ use anyhow::{ensure, Result};
use itertools::Itertools;
use once_cell::sync::Lazy;
use plonky2::field::extension::Extendable;
use plonky2::field::fft::ifft;
use plonky2::field::packable::Packable;
use plonky2::field::packed::PackedField;
use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
use plonky2::field::types::Field;
use plonky2::field::types::{Field, SmallPowers};
use plonky2::field::zero_poly_coset::ZeroPolyOnCoset;
use plonky2::fri::oracle::PolynomialBatch;
use plonky2::hash::hash_types::RichField;
Expand Down Expand Up @@ -498,11 +499,26 @@ where
// When opening the `Z`s polys at the "next" point, need to look at the point `next_step` steps away.
let next_step = 1 << quotient_degree_bits;

// Evaluation of the first Lagrange polynomial on the LDE domain.
let lagrange_first = PolynomialValues::selector(degree, 0).lde_onto_coset(quotient_degree_bits);
// Evaluation of the last Lagrange polynomial on the LDE domain.
let lagrange_last =
PolynomialValues::selector(degree, degree - 1).lde_onto_coset(quotient_degree_bits);
// Evaluation of the first Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
let lagrange_first = PolynomialValues::<F>::selector(degree, 0);
let coeffs = ifft(lagrange_first).lde(quotient_degree_bits);
let modified_poly: PolynomialCoeffs<F> =
SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
.zip(&coeffs.coeffs)
.map(|(r, &c)| r * c)
.collect::<Vec<_>>()
.into();
let lagrange_first = modified_poly.fft_with_options(None, None);
// Evaluation of the last Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
let lagrange_last = PolynomialValues::<F>::selector(degree, degree - 1);
let coeffs = ifft(lagrange_last).lde(quotient_degree_bits);
let modified_poly: PolynomialCoeffs<F> =
SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
.zip(&coeffs.coeffs)
.map(|(r, &c)| r * c)
.collect::<Vec<_>>()
.into();
let lagrange_last = modified_poly.fft_with_options(None, None);

let z_h_on_coset = ZeroPolyOnCoset::<F>::new(degree_bits, quotient_degree_bits);

Expand Down
25 changes: 21 additions & 4 deletions plonky2/src/fri/oracle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use alloc::format;
use alloc::vec::Vec;

use itertools::Itertools;
use plonky2_field::types::Field;
use plonky2_field::types::{Field, SmallPowers};
use plonky2_maybe_rayon::*;

use crate::field::extension::Extendable;
Expand Down Expand Up @@ -112,8 +112,17 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
.par_iter()
.map(|p| {
assert_eq!(p.len(), degree, "Polynomial degrees inconsistent");
p.lde(rate_bits)
.coset_fft_with_options(F::coset_shift(), Some(rate_bits), fft_root_table)
let poly = p.lde(rate_bits);

// Custom coset_fft with small shift
let modified_poly: PolynomialCoeffs<F> =
SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
.zip(&poly.coeffs)
.map(|(r, &c)| r * c)
.collect::<Vec<_>>()
.into();
modified_poly
.fft_with_options(Some(rate_bits), fft_root_table)
.values
})
.chain(
Expand Down Expand Up @@ -197,10 +206,18 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
}

let lde_final_poly = final_poly.lde(fri_params.config.rate_bits);

// Custom coset_fft with small shift
let modified_poly: PolynomialCoeffs<F::Extension> =
SmallPowers::<F::Extension>::new(F::coset_shift().to_noncanonical_u64() as u32)
.zip(&lde_final_poly.coeffs)
.map(|(r, &c)| r * c)
.collect::<Vec<_>>()
.into();
let lde_final_values = timed!(
timing,
&format!("perform final FFT {}", lde_final_poly.len()),
lde_final_poly.coset_fft(F::coset_shift().into())
modified_poly.fft_with_options(None, None)
);

let fri_proof = fri_proof::<F, C, D>(
Expand Down
29 changes: 22 additions & 7 deletions starky/src/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@ use core::iter::once;
use anyhow::{ensure, Result};
use itertools::Itertools;
use plonky2::field::extension::Extendable;
use plonky2::field::fft::ifft;
use plonky2::field::packable::Packable;
use plonky2::field::packed::PackedField;
use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
use plonky2::field::types::Field;
use plonky2::field::types::{Field, SmallPowers};
use plonky2::field::zero_poly_coset::ZeroPolyOnCoset;
use plonky2::fri::oracle::PolynomialBatch;
use plonky2::hash::hash_types::RichField;
Expand Down Expand Up @@ -227,12 +228,26 @@ where
// When opening the `Z`s polys at the "next" point, need to look at the point `next_step` steps away.
let next_step = 1 << quotient_degree_bits;

// Evaluation of the first Lagrange polynomial on the LDE domain.
let lagrange_first = PolynomialValues::selector(degree, 0).lde_onto_coset(quotient_degree_bits);
// Evaluation of the last Lagrange polynomial on the LDE domain.
let lagrange_last =
PolynomialValues::selector(degree, degree - 1).lde_onto_coset(quotient_degree_bits);

// Evaluation of the first Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
let lagrange_first = PolynomialValues::<F>::selector(degree, 0);
let coeffs = ifft(lagrange_first).lde(quotient_degree_bits);
let modified_poly: PolynomialCoeffs<F> =
SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
.zip(&coeffs.coeffs)
.map(|(r, &c)| r * c)
.collect::<Vec<_>>()
.into();
let lagrange_first = modified_poly.fft_with_options(None, None);
// Evaluation of the last Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
let lagrange_last = PolynomialValues::<F>::selector(degree, degree - 1);
let coeffs = ifft(lagrange_last).lde(quotient_degree_bits);
let modified_poly: PolynomialCoeffs<F> =
SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
.zip(&coeffs.coeffs)
.map(|(r, &c)| r * c)
.collect::<Vec<_>>()
.into();
let lagrange_last = modified_poly.fft_with_options(None, None);
let z_h_on_coset = ZeroPolyOnCoset::<F>::new(degree_bits, quotient_degree_bits);

// Retrieve the LDE values at index `i`.
Expand Down

0 comments on commit a02dfc2

Please sign in to comment.