From c3327b8884b56c84e8ae0c96e68a6da28bbdf745 Mon Sep 17 00:00:00 2001 From: Shahar Papini <43779613+spapinistarkware@users.noreply.github.com> Date: Wed, 3 Apr 2024 14:19:46 +0300 Subject: [PATCH] Dumb down FRI (#529) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This change is [Reviewable](https://reviewable.io/reviews/starkware-libs/stwo/529) --- benches/fri.rs | 10 +- src/core/backend/cpu/fri.rs | 34 ++-- src/core/backend/cpu/mod.rs | 3 - src/core/commitment_scheme/prover.rs | 8 - src/core/commitment_scheme/quotients.rs | 4 +- src/core/fields/secure_column.rs | 52 +++++ src/core/fri.rs | 238 ++++++++++++----------- src/core/poly/circle/secure_poly.rs | 13 +- src/core/poly/line.rs | 240 ++++++------------------ 9 files changed, 264 insertions(+), 338 deletions(-) diff --git a/benches/fri.rs b/benches/fri.rs index 06527d2a1..eb633febc 100644 --- a/benches/fri.rs +++ b/benches/fri.rs @@ -1,6 +1,8 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion}; use stwo::core::backend::CPUBackend; use stwo::core::fields::m31::BaseField; +use stwo::core::fields::qm31::SecureField; +use stwo::core::fields::secure_column::SecureColumn; use stwo::core::fri::FriOps; use stwo::core::poly::circle::CanonicCoset; use stwo::core::poly::line::{LineDomain, LineEvaluation}; @@ -10,9 +12,13 @@ fn folding_benchmark(c: &mut Criterion) { let domain = LineDomain::new(CanonicCoset::new(LOG_SIZE + 1).half_coset()); let evals = LineEvaluation::new( domain, - vec![BaseField::from_u32_unchecked(712837213).into(); 1 << LOG_SIZE], + SecureColumn { + columns: std::array::from_fn(|i| { + vec![BaseField::from_u32_unchecked(i as u32); 1 << LOG_SIZE] + }), + }, ); - let alpha = BaseField::from_u32_unchecked(12389).into(); + let alpha = SecureField::from_u32_unchecked(2213980, 2213981, 2213982, 2213983); c.bench_function("fold_line", |b| { b.iter(|| { black_box(CPUBackend::fold_line(black_box(&evals), black_box(alpha))); diff --git a/src/core/backend/cpu/fri.rs b/src/core/backend/cpu/fri.rs index 5f772ab98..9048d3766 100644 --- a/src/core/backend/cpu/fri.rs +++ b/src/core/backend/cpu/fri.rs @@ -1,21 +1,14 @@ -use std::iter::zip; - use super::CPUBackend; use crate::core::fft::ibutterfly; -use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; -use crate::core::fields::{ExtensionOf, Field, FieldExpOps}; +use crate::core::fields::FieldExpOps; use crate::core::fri::{FriOps, CIRCLE_TO_LINE_FOLD_STEP, FOLD_STEP}; -use crate::core::poly::circle::CircleEvaluation; +use crate::core::poly::circle::SecureEvaluation; use crate::core::poly::line::LineEvaluation; -use crate::core::poly::BitReversedOrder; use crate::core::utils::bit_reverse_index; impl FriOps for CPUBackend { - fn fold_line( - eval: &LineEvaluation, - alpha: SecureField, - ) -> LineEvaluation { + fn fold_line(eval: &LineEvaluation, alpha: SecureField) -> LineEvaluation { let n = eval.len(); assert!(n >= 2, "Evaluation too small"); @@ -23,9 +16,10 @@ impl FriOps for CPUBackend { let folded_values = eval .values + .into_iter() .array_chunks() .enumerate() - .map(|(i, &[f_x, f_neg_x])| { + .map(|(i, [f_x, f_neg_x])| { // TODO(andrew): Inefficient. Update when domain twiddles get stored in a buffer. let x = domain.at(bit_reverse_index(i << FOLD_STEP, domain.log_size())); @@ -37,22 +31,20 @@ impl FriOps for CPUBackend { LineEvaluation::new(domain.double(), folded_values) } - fn fold_circle_into_line( - dst: &mut LineEvaluation, - src: &CircleEvaluation, + fn fold_circle_into_line( + dst: &mut LineEvaluation, + src: &SecureEvaluation, alpha: SecureField, - ) where - F: ExtensionOf, - SecureField: ExtensionOf + Field, - { + ) { assert_eq!(src.len() >> CIRCLE_TO_LINE_FOLD_STEP, dst.len()); let domain = src.domain; let alpha_sq = alpha * alpha; - zip(&mut dst.values, src.array_chunks()) + src.into_iter() + .array_chunks() .enumerate() - .for_each(|(i, (dst, &[f_p, f_neg_p]))| { + .for_each(|(i, [f_p, f_neg_p])| { // TODO(andrew): Inefficient. Update when domain twiddles get stored in a buffer. let p = domain.at(bit_reverse_index( i << CIRCLE_TO_LINE_FOLD_STEP, @@ -64,7 +56,7 @@ impl FriOps for CPUBackend { ibutterfly(&mut f0_px, &mut f1_px, p.y.inverse()); let f_prime = alpha * f1_px + f0_px; - *dst = *dst * alpha_sq + f_prime; + dst.values.set(i, dst.values.at(i) * alpha_sq + f_prime); }); } } diff --git a/src/core/backend/cpu/mod.rs b/src/core/backend/cpu/mod.rs index 094e4bb54..d8f5fc59f 100644 --- a/src/core/backend/cpu/mod.rs +++ b/src/core/backend/cpu/mod.rs @@ -8,7 +8,6 @@ use std::fmt::Debug; use super::{Backend, Column, ColumnOps, FieldOps}; use crate::core::fields::Field; use crate::core::poly::circle::{CircleEvaluation, CirclePoly}; -use crate::core::poly::line::LineEvaluation; use crate::core::utils::bit_reverse; #[derive(Copy, Clone, Debug)] @@ -49,8 +48,6 @@ impl Column for Vec { pub type CPUCirclePoly = CirclePoly; pub type CPUCircleEvaluation = CircleEvaluation; -// TODO(spapini): Remove the EvalOrder on LineEvaluation. -pub type CPULineEvaluation = LineEvaluation; #[cfg(test)] mod tests { diff --git a/src/core/commitment_scheme/prover.rs b/src/core/commitment_scheme/prover.rs index adfc29c4d..e2e5385e5 100644 --- a/src/core/commitment_scheme/prover.rs +++ b/src/core/commitment_scheme/prover.rs @@ -22,7 +22,6 @@ use crate::commitment_scheme::blake2_hash::{Blake2sHash, Blake2sHasher}; use crate::commitment_scheme::blake2_merkle::Blake2sMerkleHasher; use crate::commitment_scheme::prover::{MerkleDecommitment, MerkleProver}; use crate::core::channel::Channel; -use crate::core::poly::circle::SecureEvaluation; type MerkleHasher = Blake2sMerkleHasher; type ProofChannel = Blake2sChannel; @@ -89,13 +88,6 @@ impl CommitmentSchemeProver { let columns = self.evaluations().flatten(); let quotients = compute_fri_quotients(&columns, &samples.flatten(), channel.draw_felt()); - // TODO(spapini): Conversion to CircleEvaluation can be removed when FRI supports - // SecureColumn. - let quotients = quotients - .into_iter() - .map(SecureEvaluation::to_cpu) - .collect_vec(); - // Run FRI commitment phase on the oods quotients. let fri_config = FriConfig::new(LOG_LAST_LAYER_DEGREE_BOUND, LOG_BLOWUP_FACTOR, N_QUERIES); let fri_prover = diff --git a/src/core/commitment_scheme/quotients.rs b/src/core/commitment_scheme/quotients.rs index ffb0415f2..7cd42d68e 100644 --- a/src/core/commitment_scheme/quotients.rs +++ b/src/core/commitment_scheme/quotients.rs @@ -96,7 +96,7 @@ pub fn fri_answers( random_coeff: SecureField, query_domain_per_log_size: BTreeMap, queried_values_per_column: &[Vec], -) -> Result>, VerificationError> { +) -> Result, VerificationError> { izip!(column_log_sizes, samples, queried_values_per_column) .sorted_by_key(|(log_size, ..)| Reverse(*log_size)) .group_by(|(log_size, ..)| *log_size) @@ -121,7 +121,7 @@ pub fn fri_answers_for_log_size( random_coeff: SecureField, query_domain: &SparseSubCircleDomain, queried_values_per_column: &[&Vec], -) -> Result, VerificationError> { +) -> Result { let commitment_domain = CanonicCoset::new(log_size).circle_domain(); let sample_batches = ColumnSampleBatch::new_vec(samples); for queried_values in queried_values_per_column { diff --git a/src/core/fields/secure_column.rs b/src/core/fields/secure_column.rs index 78fdc48fe..86d2e3bd6 100644 --- a/src/core/fields/secure_column.rs +++ b/src/core/fields/secure_column.rs @@ -44,4 +44,56 @@ impl> SecureColumn { pub fn is_empty(&self) -> bool { self.columns[0].is_empty() } + + pub fn to_cpu(&self) -> SecureColumn { + SecureColumn { + columns: self.columns.clone().map(|c| c.to_vec()), + } + } +} + +pub struct SecureColumnIter<'a> { + column: &'a SecureColumn, + index: usize, +} +impl Iterator for SecureColumnIter<'_> { + type Item = SecureField; + + fn next(&mut self) -> Option { + if self.index < self.column.len() { + let value = self.column.at(self.index); + self.index += 1; + Some(value) + } else { + None + } + } +} +impl<'a> IntoIterator for &'a SecureColumn { + type Item = SecureField; + type IntoIter = SecureColumnIter<'a>; + + fn into_iter(self) -> Self::IntoIter { + SecureColumnIter { + column: self, + index: 0, + } + } +} +impl FromIterator for SecureColumn { + fn from_iter>(iter: I) -> Self { + let mut columns = std::array::from_fn(|_| vec![]); + for value in iter.into_iter() { + let vals = value.to_m31_array(); + for j in 0..SECURE_EXTENSION_DEGREE { + columns[j].push(vals[j]); + } + } + SecureColumn { columns } + } +} +impl From> for Vec { + fn from(column: SecureColumn) -> Self { + column.into_iter().collect() + } } diff --git a/src/core/fri.rs b/src/core/fri.rs index 608c5523b..1dcaaa4e6 100644 --- a/src/core/fri.rs +++ b/src/core/fri.rs @@ -8,13 +8,11 @@ use itertools::Itertools; use num_traits::Zero; use thiserror::Error; -use super::backend::cpu::CPULineEvaluation; -use super::backend::CPUBackend; +use super::backend::{Backend, CPUBackend}; use super::channel::Channel; use super::fields::m31::BaseField; use super::fields::qm31::SecureField; -use super::fields::{ExtensionOf, Field, FieldOps}; -use super::poly::circle::CircleEvaluation; +use super::poly::circle::{CircleEvaluation, SecureEvaluation}; use super::poly::line::{LineEvaluation, LinePoly}; use super::poly::BitReversedOrder; // TODO(andrew): Create fri/ directory, move queries.rs there and split this file up. @@ -22,7 +20,6 @@ use super::queries::{Queries, SparseSubCircleDomain}; use crate::commitment_scheme::hasher::Hasher; use crate::commitment_scheme::merkle_decommitment::MerkleDecommitment; use crate::commitment_scheme::merkle_tree::MerkleTree; -use crate::core::backend::Column; use crate::core::circle::Coset; use crate::core::poly::line::LineDomain; use crate::core::utils::bit_reverse_index; @@ -70,7 +67,7 @@ impl FriConfig { } } -pub trait FriOps: FieldOps + Sized { +pub trait FriOps: Backend + Sized { /// Folds a degree `d` polynomial into a degree `d/2` polynomial. /// /// Let `eval` be a polynomial evaluated on a [LineDomain] `E`, `alpha` be a random field @@ -81,10 +78,7 @@ pub trait FriOps: FieldOps + Sized { /// # Panics /// /// Panics if there are less than two evaluations. - fn fold_line( - eval: &LineEvaluation, - alpha: SecureField, - ) -> LineEvaluation; + fn fold_line(eval: &LineEvaluation, alpha: SecureField) -> LineEvaluation; /// Folds and accumulates a degree `d` circle polynomial into a degree `d/2` univariate /// polynomial. @@ -102,21 +96,18 @@ pub trait FriOps: FieldOps + Sized { /// [`CircleDomain`]: super::poly::circle::CircleDomain // TODO(andrew): Make folding factor generic. // TODO(andrew): Fold directly into FRI layer to prevent allocation. - fn fold_circle_into_line( - dst: &mut LineEvaluation, - src: &CircleEvaluation, + fn fold_circle_into_line( + dst: &mut LineEvaluation, + src: &SecureEvaluation, alpha: SecureField, - ) where - F: ExtensionOf, - SecureField: ExtensionOf + Field, - Self: FieldOps; + ); } /// A FRI prover that applies the FRI protocol to prove a set of polynomials are of low degree. pub struct FriProver { config: FriConfig, inner_layers: Vec>, - last_layer_poly: LinePoly, + last_layer_poly: LinePoly, /// Unique sizes of committed columns sorted in descending order. column_log_sizes: Vec, } @@ -142,16 +133,11 @@ impl> FriProver { /// /// [`CircleDomain`]: super::poly::circle::CircleDomain // TODO(andrew): Add docs for all evaluations needing to be from canonic domains. - pub fn commit( + pub fn commit( channel: &mut impl Channel, config: FriConfig, - columns: &[CircleEvaluation], - ) -> Self - where - F: ExtensionOf, - SecureField: ExtensionOf, - B: FieldOps, - { + columns: &[SecureEvaluation], + ) -> Self { assert!(!columns.is_empty(), "no columns"); assert!(columns.is_sorted_by_key(|e| Reverse(e.len())), "not sorted"); assert!(columns.iter().all(|e| e.domain.is_canonic()), "not canonic"); @@ -177,21 +163,13 @@ impl> FriProver { /// All `columns` must be provided in descending order by size. /// /// Returns all inner layers and the evaluation of the last layer. - fn commit_inner_layers( + fn commit_inner_layers( channel: &mut impl Channel, config: FriConfig, - columns: &[CircleEvaluation], - ) -> ( - Vec>, - LineEvaluation, - ) - where - F: ExtensionOf, - SecureField: ExtensionOf, - B: FriOps + FieldOps, - { + columns: &[SecureEvaluation], + ) -> (Vec>, LineEvaluation) { // Returns the length of the [LineEvaluation] a [CircleEvaluation] gets folded into. - let folded_len = |e: &CircleEvaluation| e.len() >> CIRCLE_TO_LINE_FOLD_STEP; + let folded_len = |e: &SecureEvaluation| e.len() >> CIRCLE_TO_LINE_FOLD_STEP; let first_layer_size = folded_len(&columns[0]); let first_layer_domain = LineDomain::new(Coset::half_odds(first_layer_size.ilog2())); @@ -238,11 +216,11 @@ impl> FriProver { fn commit_last_layer( channel: &mut impl Channel, config: FriConfig, - evaluation: LineEvaluation, - ) -> LinePoly { + evaluation: LineEvaluation, + ) -> LinePoly { assert_eq!(evaluation.len(), config.last_layer_domain_size()); - let evaluation = evaluation.bit_reverse().to_cpu(); + let evaluation = evaluation.to_cpu(); let mut coeffs = evaluation.interpolate().into_ordered_coefficients(); let last_layer_degree_bound = 1 << config.log_last_layer_degree_bound; @@ -305,7 +283,7 @@ pub struct FriVerifier { column_bounds: Vec, inner_layers: Vec>, last_layer_domain: LineDomain, - last_layer_poly: LinePoly, + last_layer_poly: LinePoly, /// The queries used for decommitment. Initialized when calling /// [`FriVerifier::column_opening_positions`]. queries: Option, @@ -404,27 +382,19 @@ impl> FriVerifier { /// * The queries were sampled on the wrong domain size. /// * There aren't the same number of decommitted values as degree bounds. // TODO(andrew): Finish docs. - pub fn decommit( + pub fn decommit( mut self, - decommitted_values: Vec>, - ) -> Result<(), FriVerificationError> - where - F: ExtensionOf, - SecureField: ExtensionOf, - { + decommitted_values: Vec, + ) -> Result<(), FriVerificationError> { let queries = self.queries.take().expect("queries not sampled"); self.decommit_on_queries(&queries, decommitted_values) } - fn decommit_on_queries( + fn decommit_on_queries( self, queries: &Queries, - decommitted_values: Vec>, - ) -> Result<(), FriVerificationError> - where - F: ExtensionOf, - SecureField: ExtensionOf, - { + decommitted_values: Vec, + ) -> Result<(), FriVerificationError> { assert_eq!(queries.log_domain_size, self.expected_query_log_domain_size); assert_eq!(decommitted_values.len(), self.column_bounds.len()); @@ -437,15 +407,11 @@ impl> FriVerifier { /// Verifies all inner layer decommitments. /// /// Returns the queries and query evaluations needed for verifying the last FRI layer. - fn decommit_inner_layers( + fn decommit_inner_layers( &self, queries: &Queries, - decommitted_values: Vec>, - ) -> Result<(Queries, Vec), FriVerificationError> - where - F: ExtensionOf, - SecureField: ExtensionOf + Field, - { + decommitted_values: Vec, + ) -> Result<(Queries, Vec), FriVerificationError> { let circle_poly_alpha = self.circle_poly_alpha; let circle_poly_alpha_sq = circle_poly_alpha * circle_poly_alpha; @@ -555,7 +521,7 @@ pub trait FriChannel { fn reseed_with_inner_layer(&mut self, commitment: &Self::Digest); /// Reseeds the channel with the FRI last layer polynomial. - fn reseed_with_last_layer(&mut self, last_layer: &LinePoly); + fn reseed_with_last_layer(&mut self, last_layer: &LinePoly); /// Draws a random field element. fn draw(&mut self) -> Self::Field; @@ -627,7 +593,7 @@ impl LinePolyDegreeBound { #[derive(Debug)] pub struct FriProof { pub inner_layers: Vec>, - pub last_layer_poly: LinePoly, + pub last_layer_poly: LinePoly, } /// Number of folds for univariate polynomials. @@ -645,7 +611,7 @@ pub struct FriLayerProof { /// The subset stored corresponds to the set of evaluations the verifier doesn't have but needs /// to fold and verify the merkle decommitment. pub evals_subset: Vec, - pub decommitment: MerkleDecommitment, + pub decommitment: MerkleDecommitment, pub commitment: H::Hash, } @@ -686,8 +652,8 @@ impl> FriLayerVerifier { for leaf in decommitment.values() { // Ensure each leaf is a single value. - if let [eval] = *leaf { - expected_decommitment_evals.push(eval); + if let Ok(evals) = leaf.try_into() { + expected_decommitment_evals.push(SecureField::from_m31_array(evals)); } else { return Err(FriVerificationError::InnerLayerCommitmentInvalid { layer: self.layer_index, @@ -698,9 +664,9 @@ impl> FriLayerVerifier { let actual_decommitment_evals = sparse_evaluation .subline_evals .iter() - .flat_map(|e| &e.values); + .flat_map(|e| e.values.into_iter()); - if !actual_decommitment_evals.eq(&expected_decommitment_evals) { + if !actual_decommitment_evals.eq(expected_decommitment_evals) { return Err(FriVerificationError::InnerLayerCommitmentInvalid { layer: self.layer_index, }); @@ -780,7 +746,10 @@ impl> FriLayerVerifier { let subline_initial = self.domain.coset().index_at(subline_initial_index); let subline_domain = LineDomain::new(Coset::new(subline_initial, FOLD_STEP)); - all_subline_evals.push(LineEvaluation::new(subline_domain, subline_evals)); + all_subline_evals.push(LineEvaluation::new( + subline_domain, + subline_evals.into_iter().collect(), + )); } // Check all proof evals have been consumed. @@ -800,15 +769,15 @@ impl> FriLayerVerifier { /// of size two. Each leaf of the merkle tree commits to a single coset evaluation. // TODO(andrew): Support different step sizes. struct FriLayerProver { - evaluation: LineEvaluation, - merkle_tree: MerkleTree, + evaluation: LineEvaluation, + merkle_tree: MerkleTree, } impl> FriLayerProver { - fn new(evaluation: LineEvaluation) -> Self { + fn new(evaluation: LineEvaluation) -> Self { // TODO(spapini): Commit on slice. // TODO(spapini): Merkle tree in backend. - let merkle_tree = MerkleTree::commit(vec![evaluation.values.to_vec()]); + let merkle_tree = MerkleTree::commit(evaluation.values.to_cpu().columns.into()); #[allow(unreachable_code)] FriLayerProver { evaluation, @@ -856,39 +825,46 @@ impl> FriLayerProver { /// Holds a foldable subset of circle polynomial evaluations. #[derive(Debug, Clone)] -pub struct SparseCircleEvaluation> { - subcircle_evals: Vec>, +pub struct SparseCircleEvaluation { + subcircle_evals: Vec>, } -impl> SparseCircleEvaluation { +impl SparseCircleEvaluation { /// # Panics /// /// Panics if the evaluation domain sizes don't equal the folding factor. - pub fn new(subcircle_evals: Vec>) -> Self { + pub fn new( + subcircle_evals: Vec>, + ) -> Self { let folding_factor = 1 << CIRCLE_TO_LINE_FOLD_STEP; assert!(subcircle_evals.iter().all(|e| e.len() == folding_factor)); Self { subcircle_evals } } - fn fold(self, alpha: SecureField) -> Vec - where - SecureField: ExtensionOf, - { + fn fold(self, alpha: SecureField) -> Vec { self.subcircle_evals .into_iter() .map(|e| { let buffer_domain = LineDomain::new(e.domain.half_coset); - let mut buffer = LineEvaluation::new(buffer_domain, vec![SecureField::zero()]); - CPUBackend::fold_circle_into_line(&mut buffer, &e, alpha); - buffer.values[0] + let mut buffer = LineEvaluation::new_zero(buffer_domain); + CPUBackend::fold_circle_into_line( + &mut buffer, + &SecureEvaluation { + domain: e.domain, + values: e.values.into_iter().collect(), + }, + alpha, + ); + buffer.values.at(0) }) .collect() } } -impl<'a, F: ExtensionOf> IntoIterator for &'a mut SparseCircleEvaluation { - type Item = &'a mut CircleEvaluation; - type IntoIter = std::slice::IterMut<'a, CircleEvaluation>; +impl<'a> IntoIterator for &'a mut SparseCircleEvaluation { + type Item = &'a mut CircleEvaluation; + type IntoIter = + std::slice::IterMut<'a, CircleEvaluation>; fn into_iter(self) -> Self::IntoIter { self.subcircle_evals.iter_mut() @@ -899,14 +875,14 @@ impl<'a, F: ExtensionOf> IntoIterator for &'a mut SparseCircleEvaluat /// Evaluation is held at the CPU backend. #[derive(Debug, Clone)] struct SparseLineEvaluation { - subline_evals: Vec>, + subline_evals: Vec>, } impl SparseLineEvaluation { /// # Panics /// /// Panics if the evaluation domain sizes don't equal the folding factor. - fn new(subline_evals: Vec>) -> Self { + fn new(subline_evals: Vec>) -> Self { let folding_factor = 1 << FOLD_STEP; assert!(subline_evals.iter().all(|e| e.len() == folding_factor)); Self { subline_evals } @@ -915,7 +891,7 @@ impl SparseLineEvaluation { fn fold(self, alpha: SecureField) -> Vec { self.subline_evals .into_iter() - .map(|e| CPUBackend::fold_line(&e, alpha).values[0]) + .map(|e| CPUBackend::fold_line(&e, alpha).values.at(0)) .collect() } } @@ -924,24 +900,27 @@ impl SparseLineEvaluation { mod tests { use std::iter::zip; + use itertools::Itertools; use num_traits::{One, Zero}; use super::{get_opening_positions, FriVerificationError, SparseCircleEvaluation}; use crate::commitment_scheme::blake2_hash::Blake2sHasher; - use crate::core::backend::cpu::{CPUCircleEvaluation, CPUCirclePoly, CPULineEvaluation}; - use crate::core::backend::CPUBackend; + use crate::core::backend::cpu::{CPUCircleEvaluation, CPUCirclePoly}; + use crate::core::backend::{CPUBackend, Col, Column, ColumnOps}; use crate::core::circle::{CirclePointIndex, Coset}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; - use crate::core::fields::{ExtensionOf, Field}; + use crate::core::fields::secure_column::SecureColumn; + use crate::core::fields::Field; use crate::core::fri::{ CirclePolyDegreeBound, FriConfig, FriOps, FriVerifier, CIRCLE_TO_LINE_FOLD_STEP, }; - use crate::core::poly::circle::{CircleDomain, CircleEvaluation}; + use crate::core::poly::circle::{CircleDomain, SecureEvaluation}; use crate::core::poly::line::{LineDomain, LineEvaluation, LinePoly}; - use crate::core::poly::{BitReversedOrder, NaturalOrder}; + use crate::core::poly::NaturalOrder; use crate::core::queries::{Queries, SparseSubCircleDomain}; use crate::core::test_utils::test_channel; + use crate::core::utils::bit_reverse_index; /// Default blowup factor used for tests. const LOG_BLOWUP_FACTOR: u32 = 2; @@ -964,13 +943,19 @@ mod tests { let alpha = BaseField::from_u32_unchecked(19283).into(); let domain = LineDomain::new(Coset::half_odds(DEGREE.ilog2())); let drp_domain = domain.double(); - let evals = poly.evaluate(domain).bit_reverse(); + let mut values = domain + .iter() + .map(|p| poly.eval_at_point(p.into())) + .collect(); + CPUBackend::bit_reverse_column(&mut values); + let evals = LineEvaluation::new(domain, values.into_iter().collect()); let drp_evals = CPUBackend::fold_line(&evals, alpha); + let mut drp_evals = drp_evals.values.into_iter().collect_vec(); + CPUBackend::bit_reverse_column(&mut drp_evals); assert_eq!(drp_evals.len(), DEGREE / 2); - let drp_evals = drp_evals.bit_reverse(); - for (i, (&drp_eval, x)) in zip(&drp_evals.values, drp_domain).enumerate() { + for (i, (&drp_eval, x)) in zip(&drp_evals, drp_domain).enumerate() { let f_e: SecureField = even_poly.eval_at_point(x.into()); let f_o: SecureField = odd_poly.eval_at_point(x.into()); assert_eq!(drp_eval, (f_e + alpha * f_o).double(), "mismatch at {i}"); @@ -981,12 +966,10 @@ mod tests { fn fold_circle_to_line_works() { const LOG_DEGREE: u32 = 4; let circle_evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let num_folded_evals = circle_evaluation.domain.size() >> CIRCLE_TO_LINE_FOLD_STEP; let alpha = SecureField::one(); let folded_domain = LineDomain::new(circle_evaluation.domain.half_coset); - let mut folded_evaluation = - LineEvaluation::new(folded_domain, vec![Zero::zero(); num_folded_evals]); + let mut folded_evaluation = LineEvaluation::new_zero(folded_domain); CPUBackend::fold_circle_into_line(&mut folded_evaluation, &circle_evaluation, alpha); assert_eq!( @@ -1011,7 +994,10 @@ mod tests { fn committing_evaluation_from_invalid_domain_fails() { let invalid_domain = CircleDomain::new(Coset::new(CirclePointIndex::generator(), 3)); assert!(!invalid_domain.is_canonic(), "must be an invalid domain"); - let evaluation = CircleEvaluation::new(invalid_domain, vec![BaseField::one(); 1 << 4]); + let evaluation = SecureEvaluation { + domain: invalid_domain, + values: vec![SecureField::one(); 1 << 4].into_iter().collect(), + }; FriProver::commit(&mut test_channel(), FriConfig::new(2, 2, 3), &[evaluation]); } @@ -1247,49 +1233,61 @@ mod tests { fn polynomial_evaluation( log_degree: u32, log_blowup_factor: u32, - ) -> CPUCircleEvaluation { + ) -> SecureEvaluation { let poly = CPUCirclePoly::new(vec![BaseField::one(); 1 << log_degree]); let coset = Coset::half_odds(log_degree + log_blowup_factor - 1); let domain = CircleDomain::new(coset); - poly.evaluate(domain) + let values = poly.evaluate(domain); + SecureEvaluation { + domain, + values: SecureColumn { + columns: [ + values.values, + Col::::zeros(1 << (log_degree + log_blowup_factor)), + Col::::zeros(1 << (log_degree + log_blowup_factor)), + Col::::zeros(1 << (log_degree + log_blowup_factor)), + ], + }, + } } /// Returns the log degree bound of a polynomial. - fn log_degree_bound>( - polynomial: CPULineEvaluation, - ) -> u32 { - let polynomial = polynomial.bit_reverse(); + fn log_degree_bound(polynomial: LineEvaluation) -> u32 { let coeffs = polynomial.interpolate().into_ordered_coefficients(); let degree = coeffs.into_iter().rposition(|c| !c.is_zero()).unwrap_or(0); (degree + 1).ilog2() } // TODO: Remove after SubcircleDomain integration. - fn query_polynomial>( - polynomial: &CPUCircleEvaluation, + fn query_polynomial( + polynomial: &SecureEvaluation, queries: &Queries, - ) -> SparseCircleEvaluation { + ) -> SparseCircleEvaluation { let polynomial_log_size = polynomial.domain.log_size(); let positions = get_opening_positions(queries, &[queries.log_domain_size, polynomial_log_size]); open_polynomial(polynomial, &positions[&polynomial_log_size]) } - fn open_polynomial>( - polynomial: &CPUCircleEvaluation, + fn open_polynomial( + polynomial: &SecureEvaluation, positions: &SparseSubCircleDomain, - ) -> SparseCircleEvaluation { - let polynomial = polynomial.clone().bit_reverse(); - + ) -> SparseCircleEvaluation { let coset_evals = positions .iter() .map(|position| { let coset_domain = position.to_circle_domain(&polynomial.domain); let evals = coset_domain .iter_indices() - .map(|p| polynomial.get_at(p)) + .map(|p| { + polynomial.at(bit_reverse_index( + polynomial.domain.find(p).unwrap(), + polynomial.domain.log_size(), + )) + }) .collect(); - let coset_eval = CPUCircleEvaluation::::new(coset_domain, evals); + let coset_eval = + CPUCircleEvaluation::::new(coset_domain, evals); coset_eval.bit_reverse() }) .collect(); diff --git a/src/core/poly/circle/secure_poly.rs b/src/core/poly/circle/secure_poly.rs index 715b3220c..8902f49de 100644 --- a/src/core/poly/circle/secure_poly.rs +++ b/src/core/poly/circle/secure_poly.rs @@ -1,6 +1,6 @@ use std::ops::Deref; -use super::CircleDomain; +use super::{CircleDomain, CircleEvaluation}; use crate::core::backend::cpu::{CPUCircleEvaluation, CPUCirclePoly}; use crate::core::backend::CPUBackend; use crate::core::circle::CirclePoint; @@ -70,3 +70,14 @@ impl SecureEvaluation { CPUCircleEvaluation::new(self.domain, self.values.to_vec()) } } + +impl From> + for SecureEvaluation +{ + fn from(evaluation: CircleEvaluation) -> Self { + Self { + domain: evaluation.domain, + values: evaluation.values.into_iter().collect(), + } + } +} diff --git a/src/core/poly/line.rs b/src/core/poly/line.rs index e369401a0..8390438ae 100644 --- a/src/core/poly/line.rs +++ b/src/core/poly/line.rs @@ -1,20 +1,19 @@ use std::cmp::Ordering; use std::fmt::Debug; use std::iter::Map; -use std::marker::PhantomData; use std::ops::{Deref, DerefMut}; +use itertools::Itertools; use num_traits::Zero; use super::utils::fold; -use super::{BitReversedOrder, NaturalOrder}; -use crate::core::backend::cpu::CPULineEvaluation; -use crate::core::backend::{Col, Column}; +use crate::core::backend::{Backend, CPUBackend, ColumnOps}; use crate::core::circle::{CirclePoint, Coset, CosetIterator}; -use crate::core::fft::{butterfly, ibutterfly}; +use crate::core::fft::ibutterfly; use crate::core::fields::m31::BaseField; -use crate::core::fields::{ExtensionOf, Field, FieldExpOps, FieldOps}; -use crate::core::poly::utils::repeat_value; +use crate::core::fields::qm31::SecureField; +use crate::core::fields::secure_column::SecureColumn; +use crate::core::fields::{ExtensionOf, FieldExpOps}; use crate::core::utils::bit_reverse; /// Domain comprising of the x-coordinates of points in a [Coset]. @@ -100,29 +99,29 @@ type LineDomainIterator = /// A univariate polynomial defined on a [LineDomain]. #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] -pub struct LinePoly { +pub struct LinePoly { /// Coefficients of the polynomial in [line_ifft] algorithm's basis. /// /// The coefficients are stored in bit-reversed order. - coeffs: Vec, + coeffs: Vec, /// The number of coefficients stored as `log2(len(coeffs))`. log_size: u32, } -impl> LinePoly { +impl LinePoly { /// Creates a new line polynomial from bit reversed coefficients. /// /// # Panics /// /// Panics if the number of coefficients is not a power of two. - pub fn new(coeffs: Vec) -> Self { + pub fn new(coeffs: Vec) -> Self { assert!(coeffs.len().is_power_of_two()); let log_size = coeffs.len().ilog2(); Self { coeffs, log_size } } /// Evaluates the polynomial at a single point. - pub fn eval_at_point>(&self, mut x: E) -> E { + pub fn eval_at_point(&self, mut x: SecureField) -> SecureField { // TODO(Andrew): Allocation here expensive for small polynomials. let mut doublings = Vec::new(); for _ in 0..self.log_size { @@ -132,21 +131,6 @@ impl> LinePoly { fold(&self.coeffs, &doublings) } - /// Evaluates the polynomial at all points in the domain. - pub fn evaluate(mut self, domain: LineDomain) -> CPULineEvaluation { - assert!(domain.size() >= self.coeffs.len()); - - // The first few FFT layers may just copy coefficients so we do it directly. - // See the docs for `n_skipped_layers` in [line_fft]. - let log_degree_bound = self.log_size; - let n_skipped_layers = (domain.log_size() - log_degree_bound) as usize; - let duplicity = 1 << n_skipped_layers; - self.coeffs = repeat_value(&self.coeffs, duplicity); - - line_fft(&mut self.coeffs, domain, n_skipped_layers); - LineEvaluation::new(domain, self.coeffs) - } - /// Returns the number of coefficients. #[allow(clippy::len_without_is_empty)] pub fn len(&self) -> usize { @@ -157,7 +141,7 @@ impl> LinePoly { } /// Returns the polynomial's coefficients in their natural order. - pub fn into_ordered_coefficients(mut self) -> Vec { + pub fn into_ordered_coefficients(mut self) -> Vec { bit_reverse(&mut self.coeffs); self.coeffs } @@ -167,22 +151,22 @@ impl> LinePoly { /// # Panics /// /// Panics if the number of coefficients is not a power of two. - pub fn from_ordered_coefficients(mut coeffs: Vec) -> Self { + pub fn from_ordered_coefficients(mut coeffs: Vec) -> Self { bit_reverse(&mut coeffs); Self::new(coeffs) } } -impl> Deref for LinePoly { - type Target = [F]; +impl Deref for LinePoly { + type Target = [SecureField]; - fn deref(&self) -> &[F] { + fn deref(&self) -> &[SecureField] { &self.coeffs } } -impl> DerefMut for LinePoly { - fn deref_mut(&mut self) -> &mut [F] { +impl DerefMut for LinePoly { + fn deref_mut(&mut self) -> &mut [SecureField] { &mut self.coeffs } } @@ -192,30 +176,25 @@ impl> DerefMut for LinePoly { // only used by FRI where evaluations are in bit-reversed order. // TODO(spapini): Remove pub. #[derive(Clone, Debug)] -pub struct LineEvaluation, F: Field, EvalOrder = NaturalOrder> { +pub struct LineEvaluation { /// Evaluations of a univariate polynomial on `domain`. - pub values: Col, + pub values: SecureColumn, domain: LineDomain, - _eval_order: PhantomData, } -impl, F: Field, EvalOrder> LineEvaluation { +impl LineEvaluation { /// Creates new [LineEvaluation] from a set of polynomial evaluations over a [LineDomain]. /// /// # Panics /// /// Panics if the number of evaluations does not match the size of the domain. - pub fn new(domain: LineDomain, values: Col) -> Self { + pub fn new(domain: LineDomain, values: SecureColumn) -> Self { assert_eq!(values.len(), domain.size()); - Self { - values, - domain, - _eval_order: PhantomData, - } + Self { values, domain } } pub fn new_zero(domain: LineDomain) -> Self { - Self::new(domain, Col::::zeros(domain.size())) + Self::new(domain, SecureColumn::zeros(domain.size())) } /// Returns the number of evaluations. @@ -229,41 +208,21 @@ impl, F: Field, EvalOrder> LineEvaluation { } /// Clones the values into a new line evaluation in the CPU. - pub fn to_cpu(&self) -> CPULineEvaluation { - CPULineEvaluation::new(self.domain, self.values.to_vec()) + pub fn to_cpu(&self) -> LineEvaluation { + LineEvaluation::new(self.domain, self.values.to_cpu()) } } -impl> CPULineEvaluation { +impl LineEvaluation { /// Interpolates the polynomial as evaluations on `domain`. - pub fn interpolate(mut self) -> LinePoly { - line_ifft(&mut self.values, self.domain); + pub fn interpolate(self) -> LinePoly { + let mut values = self.values.into_iter().collect_vec(); + CPUBackend::bit_reverse_column(&mut values); + line_ifft(&mut values, self.domain); // Normalize the coefficients. - let len_inv = BaseField::from(self.values.len()).inverse(); - self.values.iter_mut().for_each(|v| *v *= len_inv); - LinePoly::new(self.values) - } -} - -impl, F: Field> LineEvaluation { - pub fn bit_reverse(mut self) -> LineEvaluation { - B::bit_reverse_column(&mut self.values); - LineEvaluation { - values: self.values, - domain: self.domain, - _eval_order: PhantomData, - } - } -} - -impl, F: Field> LineEvaluation { - pub fn bit_reverse(mut self) -> LineEvaluation { - B::bit_reverse_column(&mut self.values); - LineEvaluation { - values: self.values, - domain: self.domain, - _eval_order: PhantomData, - } + let len_inv = BaseField::from(values.len()).inverse(); + values.iter_mut().for_each(|v| *v *= len_inv); + LinePoly::new(values) } } @@ -303,55 +262,18 @@ fn line_ifft>(values: &mut [F], mut domain: LineDomain } } -/// Performs a univariate FFT of a polynomial over a [LineDomain]. -/// -/// The transform happens in-place. `values` consist of coefficients in [line_ifft] algorithm's -/// basis need to be stored in bit-reversed order. After the transformation `values` becomes -/// evaluations of the polynomial over `domain` stored in natural order. -/// -/// The `n_skipped_layers` argument allows specifying how many of the initial butterfly layers of -/// the FFT to skip. This is useful when doing more efficient degree aware FFTs as the butterflies -/// in the first layers of the FFT only involve copying coefficients to different locations (because -/// one or more of the coefficients is zero). This new algorithm is `O(n log d)` vs `O(n log n)` -/// where `n` is the domain size and `d` is the number of coefficients. -/// -/// # Panics -/// -/// Panics if the number of values doesn't match the size of the domain. -fn line_fft>( - values: &mut [F], - mut domain: LineDomain, - n_skipped_layers: usize, -) { - assert_eq!(values.len(), domain.size()); - - // Construct the domains we need. - let mut domains = vec![]; - while domain.size() > 1 << n_skipped_layers { - domains.push(domain); - domain = domain.double(); - } - - // Execute the butterfly layers. - for domain in domains.iter().rev() { - for chunk in values.chunks_exact_mut(domain.size()) { - let (l, r) = chunk.split_at_mut(domain.size() / 2); - for (i, x) in domain.iter().take(domain.size() / 2).enumerate() { - butterfly(&mut l[i], &mut r[i], x); - } - } - } -} - #[cfg(test)] mod tests { type B = CPUBackend; + use itertools::Itertools; + use super::LineDomain; - use crate::core::backend::CPUBackend; + use crate::core::backend::{CPUBackend, ColumnOps}; use crate::core::circle::{CirclePoint, Coset}; use crate::core::fields::m31::BaseField; use crate::core::poly::line::{LineEvaluation, LinePoly}; + use crate::core::utils::bit_reverse_index; #[test] #[should_panic] @@ -425,16 +347,16 @@ mod tests { } #[test] - fn line_polynomial_evaluation() { + fn line_evaluation_interpolation() { let poly = LinePoly::new(vec![ - BaseField::from(7), // 7 * 1 - BaseField::from(9), // 9 * pi(x) - BaseField::from(5), // 5 * x - BaseField::from(3), // 3 * pi(x)*x + BaseField::from(7).into(), // 7 * 1 + BaseField::from(9).into(), // 9 * pi(x) + BaseField::from(5).into(), // 5 * x + BaseField::from(3).into(), // 3 * pi(x)*x ]); let coset = Coset::half_odds(poly.len().ilog2()); let domain = LineDomain::new(coset); - let expected_evals = domain + let mut values = domain .iter() .map(|x| { let pi_x = CirclePoint::double_x(x); @@ -443,62 +365,9 @@ mod tests { + poly.coeffs[2] * x + poly.coeffs[3] * pi_x * x }) - .collect::>(); - - let actual_evals = poly.evaluate(domain); - - assert_eq!(actual_evals.values, expected_evals); - } - - #[test] - fn line_polynomial_evaluation_on_larger_domain() { - let poly = LinePoly::new(vec![ - BaseField::from(7), // 7 * 1 - BaseField::from(9), // 9 * pi(x) - BaseField::from(5), // 5 * x - BaseField::from(3), // 3 * pi(x)*x - ]); - let coset = Coset::half_odds(4 + poly.len().ilog2()); - let domain = LineDomain::new(coset); - let expected_evals = domain - .iter() - .map(|x| { - let pi_x = CirclePoint::double_x(x); - poly.coeffs[0] - + poly.coeffs[1] * pi_x - + poly.coeffs[2] * x - + poly.coeffs[3] * pi_x * x - }) - .collect::>(); - - let actual_evals = poly.evaluate(domain); - - assert_eq!(actual_evals.values, expected_evals); - } - - #[test] - fn line_evaluation_interpolation() { - let poly = LinePoly::new(vec![ - BaseField::from(7), // 7 * 1 - BaseField::from(9), // 9 * pi(x) - BaseField::from(5), // 5 * x - BaseField::from(3), // 3 * pi(x)*x - ]); - let coset = Coset::half_odds(poly.len().ilog2()); - let domain = LineDomain::new(coset); - let evals = LineEvaluation::::new( - domain, - domain - .iter() - .map(|x| { - let pi_x = CirclePoint::double_x(x); - poly.coeffs[0] - + poly.coeffs[1] * pi_x - + poly.coeffs[2] * x - + poly.coeffs[3] * pi_x * x - }) - .collect::>(), - ); + .collect_vec(); + CPUBackend::bit_reverse_column(&mut values); + let evals = LineEvaluation::::new(domain, values.into_iter().collect()); let interpolated_poly = evals.interpolate(); @@ -510,12 +379,21 @@ mod tests { const LOG_SIZE: u32 = 2; let coset = Coset::half_odds(LOG_SIZE); let domain = LineDomain::new(coset); - let evals = - LineEvaluation::::new(domain, (0..1 << LOG_SIZE).map(BaseField::from).collect()); + let evals = LineEvaluation::::new( + domain, + (0..1 << LOG_SIZE) + .map(BaseField::from) + .map(|x| x.into()) + .collect(), + ); let poly = evals.clone().interpolate(); for (i, x) in domain.iter().enumerate() { - assert_eq!(poly.eval_at_point(x), evals.values[i], "mismatch at {i}"); + assert_eq!( + poly.eval_at_point(x.into()), + evals.values.at(bit_reverse_index(i, domain.log_size())), + "mismatch at {i}" + ); } } }