From 9d91b8a1cafe990d2b5cd75895bcca1731095276 Mon Sep 17 00:00:00 2001 From: Shahar Papini <43779613+spapinistarkware@users.noreply.github.com> Date: Thu, 4 Apr 2024 12:11:11 +0300 Subject: [PATCH] Fri AVX ops (#563) 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/563) --- benches/fri.rs | 9 +- src/core/backend/avx512/fft/mod.rs | 2 +- src/core/backend/avx512/fri.rs | 160 +++++++++++++++++ src/core/backend/avx512/mod.rs | 12 +- src/core/backend/avx512/qm31.rs | 17 +- src/core/backend/cpu/fri.rs | 59 ++----- src/core/commitment_scheme/prover.rs | 7 +- src/core/fri.rs | 246 +++++++++++++++++++++------ src/core/poly/circle/secure_poly.rs | 1 + src/core/poly/line.rs | 9 + src/core/poly/utils.rs | 14 +- src/core/prover/mod.rs | 8 +- 12 files changed, 427 insertions(+), 117 deletions(-) create mode 100644 src/core/backend/avx512/fri.rs diff --git a/benches/fri.rs b/benches/fri.rs index eb633febc..2e645bfbf 100644 --- a/benches/fri.rs +++ b/benches/fri.rs @@ -4,7 +4,7 @@ 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::circle::{CanonicCoset, PolyOps}; use stwo::core::poly::line::{LineDomain, LineEvaluation}; fn folding_benchmark(c: &mut Criterion) { @@ -19,9 +19,14 @@ fn folding_benchmark(c: &mut Criterion) { }, ); let alpha = SecureField::from_u32_unchecked(2213980, 2213981, 2213982, 2213983); + let twiddles = CPUBackend::precompute_twiddles(domain.coset()); c.bench_function("fold_line", |b| { b.iter(|| { - black_box(CPUBackend::fold_line(black_box(&evals), black_box(alpha))); + black_box(CPUBackend::fold_line( + black_box(&evals), + black_box(alpha), + &twiddles, + )); }) }); } diff --git a/src/core/backend/avx512/fft/mod.rs b/src/core/backend/avx512/fft/mod.rs index 9e4eab25a..37dbbd686 100644 --- a/src/core/backend/avx512/fft/mod.rs +++ b/src/core/backend/avx512/fft/mod.rs @@ -63,7 +63,7 @@ pub unsafe fn transpose_vecs(values: *mut i32, log_n_vecs: usize) { /// Computes the twiddles for the first fft layer from the second, and loads both to AVX registers. /// Returns the twiddles for the first layer and the twiddles for the second layer. /// # Safety -unsafe fn compute_first_twiddles(twiddle1_dbl: [i32; 8]) -> (__m512i, __m512i) { +pub unsafe fn compute_first_twiddles(twiddle1_dbl: [i32; 8]) -> (__m512i, __m512i) { // Start by loading the twiddles for the second layer (layer 1): // The twiddles for layer 1 are replicated in the following pattern: // 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 diff --git a/src/core/backend/avx512/fri.rs b/src/core/backend/avx512/fri.rs new file mode 100644 index 000000000..1bc311c2c --- /dev/null +++ b/src/core/backend/avx512/fri.rs @@ -0,0 +1,160 @@ +use super::AVX512Backend; +use crate::core::backend::avx512::fft::compute_first_twiddles; +use crate::core::backend::avx512::fft::ifft::avx_ibutterfly; +use crate::core::backend::avx512::qm31::PackedSecureField; +use crate::core::backend::avx512::VECS_LOG_SIZE; +use crate::core::fields::qm31::SecureField; +use crate::core::fields::secure_column::SecureColumn; +use crate::core::fri::{self, FriOps}; +use crate::core::poly::circle::SecureEvaluation; +use crate::core::poly::line::LineEvaluation; +use crate::core::poly::twiddles::TwiddleTree; +use crate::core::poly::utils::domain_line_twiddles_from_tree; + +impl FriOps for AVX512Backend { + fn fold_line( + eval: &LineEvaluation, + alpha: SecureField, + twiddles: &TwiddleTree, + ) -> LineEvaluation { + let log_size = eval.len().ilog2(); + if log_size <= VECS_LOG_SIZE as u32 { + let eval = fri::fold_line(&eval.to_cpu(), alpha); + return LineEvaluation::new(eval.domain(), eval.values.into_iter().collect()); + } + + let domain = eval.domain(); + let itwiddles = domain_line_twiddles_from_tree(domain, &twiddles.itwiddles)[0]; + + let mut folded_values = SecureColumn::zeros(1 << (log_size - 1)); + + for vec_index in 0..(1 << (log_size - 1 - VECS_LOG_SIZE as u32)) { + let value = unsafe { + let twiddle_dbl: [i32; 16] = + std::array::from_fn(|i| *itwiddles.get_unchecked(vec_index * 16 + i)); + let val0 = eval.values.packed_at(vec_index * 2).to_packed_m31s(); + let val1 = eval.values.packed_at(vec_index * 2 + 1).to_packed_m31s(); + let pairs: [_; 4] = std::array::from_fn(|i| { + let (a, b) = val0[i].deinterleave_with(val1[i]); + avx_ibutterfly(a, b, std::mem::transmute(twiddle_dbl)) + }); + let val0 = PackedSecureField::from_packed_m31s(std::array::from_fn(|i| pairs[i].0)); + let val1 = PackedSecureField::from_packed_m31s(std::array::from_fn(|i| pairs[i].1)); + val0 + PackedSecureField::broadcast(alpha) * val1 + }; + folded_values.set_packed(vec_index, value); + } + + LineEvaluation::new(domain.double(), folded_values) + } + + fn fold_circle_into_line( + dst: &mut LineEvaluation, + src: &SecureEvaluation, + alpha: SecureField, + twiddles: &TwiddleTree, + ) { + let log_size = src.len().ilog2(); + assert!(log_size > VECS_LOG_SIZE as u32, "Evaluation too small"); + + let domain = src.domain; + let alpha_sq = alpha * alpha; + let itwiddles = domain_line_twiddles_from_tree(domain, &twiddles.itwiddles)[0]; + + for vec_index in 0..(1 << (log_size - 1 - VECS_LOG_SIZE as u32)) { + let value = unsafe { + // The 16 twiddles of the circle domain can be derived from the 8 twiddles of the + // next line domain. See `compute_first_twiddles()`. + let twiddle_dbl: [i32; 8] = + std::array::from_fn(|i| *itwiddles.get_unchecked(vec_index * 8 + i)); + let (t0, _) = compute_first_twiddles(twiddle_dbl); + let val0 = src.values.packed_at(vec_index * 2).to_packed_m31s(); + let val1 = src.values.packed_at(vec_index * 2 + 1).to_packed_m31s(); + let pairs: [_; 4] = std::array::from_fn(|i| { + let (a, b) = val0[i].deinterleave_with(val1[i]); + avx_ibutterfly(a, b, t0) + }); + let val0 = PackedSecureField::from_packed_m31s(std::array::from_fn(|i| pairs[i].0)); + let val1 = PackedSecureField::from_packed_m31s(std::array::from_fn(|i| pairs[i].1)); + val0 + PackedSecureField::broadcast(alpha) * val1 + }; + dst.values.set_packed( + vec_index, + dst.values.packed_at(vec_index) * PackedSecureField::broadcast(alpha_sq) + value, + ); + } + } +} + +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] +#[cfg(test)] +mod tests { + use crate::core::backend::avx512::AVX512Backend; + use crate::core::backend::CPUBackend; + use crate::core::fields::qm31::SecureField; + use crate::core::fields::secure_column::SecureColumn; + use crate::core::fri::FriOps; + use crate::core::poly::circle::{CanonicCoset, PolyOps, SecureEvaluation}; + use crate::core::poly::line::{LineDomain, LineEvaluation}; + use crate::qm31; + + #[test] + fn test_fold_line() { + const LOG_SIZE: u32 = 7; + let values: Vec = (0..(1 << LOG_SIZE)) + .map(|i| qm31!(4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3)) + .collect(); + let alpha = qm31!(1, 3, 5, 7); + let domain = LineDomain::new(CanonicCoset::new(LOG_SIZE + 1).half_coset()); + let cpu_fold = CPUBackend::fold_line( + &LineEvaluation::new(domain, values.iter().copied().collect()), + alpha, + &CPUBackend::precompute_twiddles(domain.coset()), + ); + + let avx_fold = AVX512Backend::fold_line( + &LineEvaluation::new(domain, values.iter().copied().collect()), + alpha, + &AVX512Backend::precompute_twiddles(domain.coset()), + ); + + assert_eq!(cpu_fold.values.to_vec(), avx_fold.values.to_vec()); + } + + #[test] + fn test_fold_circle_into_line() { + const LOG_SIZE: u32 = 7; + let values: Vec = (0..(1 << LOG_SIZE)) + .map(|i| qm31!(4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3)) + .collect(); + let alpha = qm31!(1, 3, 5, 7); + let circle_domain = CanonicCoset::new(LOG_SIZE).circle_domain(); + let line_domain = LineDomain::new(circle_domain.half_coset); + + let mut cpu_fold = + LineEvaluation::new(line_domain, SecureColumn::zeros(1 << (LOG_SIZE - 1))); + CPUBackend::fold_circle_into_line( + &mut cpu_fold, + &SecureEvaluation { + domain: circle_domain, + values: values.iter().copied().collect(), + }, + alpha, + &CPUBackend::precompute_twiddles(line_domain.coset()), + ); + + let mut avx_fold = + LineEvaluation::new(line_domain, SecureColumn::zeros(1 << (LOG_SIZE - 1))); + AVX512Backend::fold_circle_into_line( + &mut avx_fold, + &SecureEvaluation { + domain: circle_domain, + values: values.iter().copied().collect(), + }, + alpha, + &AVX512Backend::precompute_twiddles(line_domain.coset()), + ); + + assert_eq!(cpu_fold.values.to_vec(), avx_fold.values.to_vec()); + } +} diff --git a/src/core/backend/avx512/mod.rs b/src/core/backend/avx512/mod.rs index 1b62962b7..dd53ea49e 100644 --- a/src/core/backend/avx512/mod.rs +++ b/src/core/backend/avx512/mod.rs @@ -4,6 +4,7 @@ pub mod blake2s_avx; pub mod circle; pub mod cm31; pub mod fft; +mod fri; pub mod m31; pub mod qm31; pub mod quotients; @@ -17,7 +18,7 @@ use self::bit_reverse::bit_reverse_m31; use self::cm31::PackedCM31; pub use self::m31::{PackedBaseField, K_BLOCK_SIZE}; use self::qm31::PackedSecureField; -use super::{Backend, Column, ColumnOps}; +use super::{Backend, CPUBackend, Column, ColumnOps}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::secure_column::SecureColumn; @@ -252,6 +253,15 @@ impl SecureColumn { } } +impl FromIterator for SecureColumn { + fn from_iter>(iter: I) -> Self { + let cpu_col = SecureColumn::::from_iter(iter); + SecureColumn { + columns: cpu_col.columns.map(|col| col.into_iter().collect()), + } + } +} + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] #[cfg(test)] mod tests { diff --git a/src/core/backend/avx512/qm31.rs b/src/core/backend/avx512/qm31.rs index efb9fbbb1..e0d06ffd4 100644 --- a/src/core/backend/avx512/qm31.rs +++ b/src/core/backend/avx512/qm31.rs @@ -46,17 +46,24 @@ impl PackedSecureField { // Multiply packed QM31 by packed M31. pub fn mul_packed_m31(&self, rhs: PackedBaseField) -> PackedSecureField { - let a = self.0[0].0[0] * rhs; - let b = self.0[0].0[1] * rhs; - let c = self.0[1].0[0] * rhs; - let d = self.0[1].0[1] * rhs; - PackedSecureField([PackedCM31([a, b]), PackedCM31([c, d])]) + Self::from_packed_m31s(self.to_packed_m31s().map(|m31| m31 * rhs)) } /// Sums all the elements in the packed M31 element. pub fn pointwise_sum(self) -> QM31 { self.to_array().into_iter().sum() } + + pub fn to_packed_m31s(&self) -> [PackedBaseField; 4] { + [self.a().a(), self.a().b(), self.b().a(), self.b().b()] + } + + pub fn from_packed_m31s(array: [PackedBaseField; 4]) -> Self { + Self([ + PackedCM31([array[0], array[1]]), + PackedCM31([array[2], array[3]]), + ]) + } } impl Add for PackedSecureField { type Output = Self; diff --git a/src/core/backend/cpu/fri.rs b/src/core/backend/cpu/fri.rs index 9048d3766..a30bce6dd 100644 --- a/src/core/backend/cpu/fri.rs +++ b/src/core/backend/cpu/fri.rs @@ -1,62 +1,25 @@ use super::CPUBackend; -use crate::core::fft::ibutterfly; use crate::core::fields::qm31::SecureField; -use crate::core::fields::FieldExpOps; -use crate::core::fri::{FriOps, CIRCLE_TO_LINE_FOLD_STEP, FOLD_STEP}; +use crate::core::fri::{fold_circle_into_line, fold_line, FriOps}; use crate::core::poly::circle::SecureEvaluation; use crate::core::poly::line::LineEvaluation; -use crate::core::utils::bit_reverse_index; +use crate::core::poly::twiddles::TwiddleTree; +// TODO(spapini): Optimized these functions as well. impl FriOps for CPUBackend { - fn fold_line(eval: &LineEvaluation, alpha: SecureField) -> LineEvaluation { - let n = eval.len(); - assert!(n >= 2, "Evaluation too small"); - - let domain = eval.domain(); - - let folded_values = eval - .values - .into_iter() - .array_chunks() - .enumerate() - .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())); - - let (mut f0, mut f1) = (f_x, f_neg_x); - ibutterfly(&mut f0, &mut f1, x.inverse()); - f0 + alpha * f1 - }) - .collect(); - - LineEvaluation::new(domain.double(), folded_values) + fn fold_line( + eval: &LineEvaluation, + alpha: SecureField, + _twiddles: &TwiddleTree, + ) -> LineEvaluation { + fold_line(eval, alpha) } fn fold_circle_into_line( dst: &mut LineEvaluation, src: &SecureEvaluation, alpha: SecureField, + _twiddles: &TwiddleTree, ) { - assert_eq!(src.len() >> CIRCLE_TO_LINE_FOLD_STEP, dst.len()); - - let domain = src.domain; - let alpha_sq = alpha * alpha; - - src.into_iter() - .array_chunks() - .enumerate() - .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, - domain.log_size(), - )); - - // Calculate `f0(px)` and `f1(px)` such that `2f(p) = f0(px) + py * f1(px)`. - let (mut f0_px, mut f1_px) = (f_p, f_neg_p); - ibutterfly(&mut f0_px, &mut f1_px, p.y.inverse()); - let f_prime = alpha * f1_px + f0_px; - - dst.values.set(i, dst.values.at(i) * alpha_sq + f_prime); - }); + fold_circle_into_line(dst, src, alpha) } } diff --git a/src/core/commitment_scheme/prover.rs b/src/core/commitment_scheme/prover.rs index ca8ace0e4..c7c9cb143 100644 --- a/src/core/commitment_scheme/prover.rs +++ b/src/core/commitment_scheme/prover.rs @@ -22,6 +22,7 @@ use crate::commitment_scheme::blake2_hash::Blake2sHash; use crate::commitment_scheme::blake2_merkle::Blake2sMerkleHasher; use crate::commitment_scheme::prover::{MerkleDecommitment, MerkleProver}; use crate::core::channel::Channel; +use crate::core::poly::twiddles::TwiddleTree; type MerkleHasher = Blake2sMerkleHasher; type ProofChannel = Blake2sChannel; @@ -65,6 +66,7 @@ impl CommitmentSchemeProver { &self, sampled_points: TreeVec>>>, channel: &mut ProofChannel, + twiddles: &TwiddleTree, ) -> CommitmentSchemeProof { // Evaluate polynomials on samples points. let samples = self @@ -90,8 +92,9 @@ impl CommitmentSchemeProver { // 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 = - FriProver::::commit(channel, fri_config, "ients); + let fri_prover = FriProver::::commit( + channel, fri_config, "ients, twiddles, + ); // Proof of work. let proof_of_work = ProofOfWork::new(PROOF_OF_WORK_BITS).prove(channel); diff --git a/src/core/fri.rs b/src/core/fri.rs index 0166c97e7..61c5e9376 100644 --- a/src/core/fri.rs +++ b/src/core/fri.rs @@ -14,6 +14,7 @@ use super::fields::qm31::SecureField; use super::fields::secure_column::{SecureColumn, SECURE_EXTENSION_DEGREE}; use super::poly::circle::{CircleEvaluation, SecureEvaluation}; use super::poly::line::{LineEvaluation, LinePoly}; +use super::poly::twiddles::TwiddleTree; use super::poly::BitReversedOrder; // TODO(andrew): Create fri/ directory, move queries.rs there and split this file up. use super::queries::{Queries, SparseSubCircleDomain}; @@ -21,6 +22,8 @@ use crate::commitment_scheme::ops::{MerkleHasher, MerkleOps}; use crate::commitment_scheme::prover::{MerkleDecommitment, MerkleProver}; use crate::commitment_scheme::verifier::{MerkleVerificationError, MerkleVerifier}; use crate::core::circle::Coset; +use crate::core::fft::ibutterfly; +use crate::core::fields::FieldExpOps; use crate::core::poly::line::LineDomain; use crate::core::utils::bit_reverse_index; @@ -78,7 +81,11 @@ pub trait FriOps: Backend + 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, + twiddles: &TwiddleTree, + ) -> LineEvaluation; /// Folds and accumulates a degree `d` circle polynomial into a degree `d/2` univariate /// polynomial. @@ -100,6 +107,7 @@ pub trait FriOps: Backend + Sized { dst: &mut LineEvaluation, src: &SecureEvaluation, alpha: SecureField, + twiddles: &TwiddleTree, ); } @@ -137,12 +145,13 @@ impl, H: MerkleHasher> FriProver { channel: &mut impl Channel, config: FriConfig, columns: &[SecureEvaluation], + twiddles: &TwiddleTree, ) -> 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"); let (inner_layers, last_layer_evaluation) = - Self::commit_inner_layers(channel, config, columns); + Self::commit_inner_layers(channel, config, columns, twiddles); let last_layer_poly = Self::commit_last_layer(channel, config, last_layer_evaluation); let column_log_sizes = columns @@ -167,6 +176,7 @@ impl, H: MerkleHasher> FriProver { channel: &mut impl Channel, config: FriConfig, columns: &[SecureEvaluation], + twiddles: &TwiddleTree, ) -> (Vec>, LineEvaluation) { // Returns the length of the [LineEvaluation] a [CircleEvaluation] gets folded into. let folded_len = |e: &SecureEvaluation| e.len() >> CIRCLE_TO_LINE_FOLD_STEP; @@ -185,13 +195,18 @@ impl, H: MerkleHasher> FriProver { while layer_evaluation.len() > config.last_layer_domain_size() { // Check for any columns (circle poly evaluations) that should be combined. while let Some(column) = columns.next_if(|c| folded_len(c) == layer_evaluation.len()) { - B::fold_circle_into_line(&mut layer_evaluation, column, circle_poly_alpha); + B::fold_circle_into_line( + &mut layer_evaluation, + column, + circle_poly_alpha, + twiddles, + ); } let layer = FriLayerProver::new(layer_evaluation); channel.mix_digest(layer.merkle_tree.root()); let folding_alpha = channel.draw_felt(); - let folded_layer_evaluation = B::fold_line(&layer.evaluation, folding_alpha); + let folded_layer_evaluation = B::fold_line(&layer.evaluation, folding_alpha, twiddles); layer_evaluation = folded_layer_evaluation; layers.push(layer); @@ -848,7 +863,7 @@ impl SparseCircleEvaluation { .map(|e| { let buffer_domain = LineDomain::new(e.domain.half_coset); let mut buffer = LineEvaluation::new_zero(buffer_domain); - CPUBackend::fold_circle_into_line( + fold_circle_into_line( &mut buffer, &SecureEvaluation { domain: e.domain, @@ -892,11 +907,72 @@ impl SparseLineEvaluation { fn fold(self, alpha: SecureField) -> Vec { self.subline_evals .into_iter() - .map(|e| CPUBackend::fold_line(&e, alpha).values.at(0)) + .map(|e| fold_line(&e, alpha).values.at(0)) .collect() } } +/// Folds a degree `d` polynomial into a degree `d/2` polynomial. +/// See [`FriOps::fold_line`]. +pub fn fold_line( + eval: &LineEvaluation, + alpha: SecureField, +) -> LineEvaluation { + let n = eval.len(); + assert!(n >= 2, "Evaluation too small"); + + let domain = eval.domain(); + + let folded_values = eval + .values + .into_iter() + .array_chunks() + .enumerate() + .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())); + + let (mut f0, mut f1) = (f_x, f_neg_x); + ibutterfly(&mut f0, &mut f1, x.inverse()); + f0 + alpha * f1 + }) + .collect(); + + LineEvaluation::new(domain.double(), folded_values) +} + +/// Folds and accumulates a degree `d` circle polynomial into a degree `d/2` univariate +/// polynomial. +/// See [`FriOps::fold_circle_into_line`]. +pub fn fold_circle_into_line( + dst: &mut LineEvaluation, + src: &SecureEvaluation, + alpha: SecureField, +) { + assert_eq!(src.len() >> CIRCLE_TO_LINE_FOLD_STEP, dst.len()); + + let domain = src.domain; + let alpha_sq = alpha * alpha; + + src.into_iter() + .array_chunks() + .enumerate() + .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, + domain.log_size(), + )); + + // Calculate `f0(px)` and `f1(px)` such that `2f(p) = f0(px) + py * f1(px)`. + let (mut f0_px, mut f1_px) = (f_p, f_neg_p); + ibutterfly(&mut f0_px, &mut f1_px, p.y.inverse()); + let f_prime = alpha * f1_px + f0_px; + + dst.values.set(i, dst.values.at(i) * alpha_sq + f_prime); + }); +} + #[cfg(test)] mod tests { use std::iter::zip; @@ -914,9 +990,10 @@ mod tests { 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, + fold_circle_into_line, fold_line, CirclePolyDegreeBound, FriConfig, FriVerifier, + CIRCLE_TO_LINE_FOLD_STEP, }; - use crate::core::poly::circle::{CircleDomain, SecureEvaluation}; + use crate::core::poly::circle::{CircleDomain, PolyOps, SecureEvaluation}; use crate::core::poly::line::{LineDomain, LineEvaluation, LinePoly}; use crate::core::poly::NaturalOrder; use crate::core::queries::{Queries, SparseSubCircleDomain}; @@ -951,7 +1028,7 @@ mod tests { CPUBackend::bit_reverse_column(&mut values); let evals = LineEvaluation::new(domain, values.into_iter().collect()); - let drp_evals = CPUBackend::fold_line(&evals, alpha); + let drp_evals = fold_line(&evals, alpha); let mut drp_evals = drp_evals.values.into_iter().collect_vec(); CPUBackend::bit_reverse_column(&mut drp_evals); @@ -971,7 +1048,7 @@ mod tests { let folded_domain = LineDomain::new(circle_evaluation.domain.half_coset); let mut folded_evaluation = LineEvaluation::new_zero(folded_domain); - CPUBackend::fold_circle_into_line(&mut folded_evaluation, &circle_evaluation, alpha); + fold_circle_into_line(&mut folded_evaluation, &circle_evaluation, alpha); assert_eq!( log_degree_bound(folded_evaluation), @@ -987,7 +1064,12 @@ mod tests { let config = FriConfig::new(2, LOG_EXPECTED_BLOWUP_FACTOR, 3); let evaluation = polynomial_evaluation(6, LOG_INVALID_BLOWUP_FACTOR); - FriProver::commit(&mut test_channel(), config, &[evaluation]); + FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); } #[test] @@ -1000,18 +1082,28 @@ mod tests { values: vec![SecureField::one(); 1 << 4].into_iter().collect(), }; - FriProver::commit(&mut test_channel(), FriConfig::new(2, 2, 3), &[evaluation]); + FriProver::commit( + &mut test_channel(), + FriConfig::new(2, 2, 3), + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); } #[test] fn valid_proof_passes_verification() -> Result<(), FriVerificationError> { const LOG_DEGREE: u32 = 3; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![5], log_domain_size); let config = FriConfig::new(1, LOG_BLOWUP_FACTOR, queries.len()); - let decommitment_value = query_polynomial(&polynomial, &queries); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let decommitment_value = query_polynomial(&evaluation, &queries); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let proof = prover.decommit_on_queries(&queries); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let verifier = FriVerifier::commit(&mut test_channel(), config, proof, bound).unwrap(); @@ -1024,12 +1116,17 @@ mod tests { { const LOG_DEGREE: u32 = 3; const LAST_LAYER_LOG_BOUND: u32 = 0; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![5], log_domain_size); let config = FriConfig::new(LAST_LAYER_LOG_BOUND, LOG_BLOWUP_FACTOR, queries.len()); - let decommitment_value = query_polynomial(&polynomial, &queries); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let decommitment_value = query_polynomial(&evaluation, &queries); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let proof = prover.decommit_on_queries(&queries); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let verifier = FriVerifier::commit(&mut test_channel(), config, proof, bound).unwrap(); @@ -1040,12 +1137,17 @@ mod tests { #[test] fn valid_mixed_degree_proof_passes_verification() -> Result<(), FriVerificationError> { const LOG_DEGREES: [u32; 3] = [6, 5, 4]; - let polynomials = LOG_DEGREES.map(|log_d| polynomial_evaluation(log_d, LOG_BLOWUP_FACTOR)); - let log_domain_size = polynomials[0].domain.log_size(); + let evaluations = LOG_DEGREES.map(|log_d| polynomial_evaluation(log_d, LOG_BLOWUP_FACTOR)); + let log_domain_size = evaluations[0].domain.log_size(); let queries = Queries::from_positions(vec![7, 70], log_domain_size); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, queries.len()); - let prover = FriProver::commit(&mut test_channel(), config, &polynomials); - let decommitment_values = polynomials.map(|p| query_polynomial(&p, &queries)).to_vec(); + let prover = FriProver::commit( + &mut test_channel(), + config, + &evaluations, + &CPUBackend::precompute_twiddles(evaluations[0].domain.half_coset), + ); + let decommitment_values = evaluations.map(|p| query_polynomial(&p, &queries)).to_vec(); let proof = prover.decommit_on_queries(&queries); let bounds = LOG_DEGREES.map(CirclePolyDegreeBound::new).to_vec(); let verifier = FriVerifier::commit(&mut test_channel(), config, proof, bounds).unwrap(); @@ -1057,11 +1159,16 @@ mod tests { fn valid_mixed_degree_end_to_end_proof_passes_verification() -> Result<(), FriVerificationError> { const LOG_DEGREES: [u32; 3] = [6, 5, 4]; - let polynomials = LOG_DEGREES.map(|log_d| polynomial_evaluation(log_d, LOG_BLOWUP_FACTOR)); + let evaluations = LOG_DEGREES.map(|log_d| polynomial_evaluation(log_d, LOG_BLOWUP_FACTOR)); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, 3); - let prover = FriProver::commit(&mut test_channel(), config, &polynomials); + let prover = FriProver::commit( + &mut test_channel(), + config, + &evaluations, + &CPUBackend::precompute_twiddles(evaluations[0].domain.half_coset), + ); let (proof, prover_opening_positions) = prover.decommit(&mut test_channel()); - let decommitment_values = zip(&polynomials, prover_opening_positions.values().rev()) + let decommitment_values = zip(&evaluations, prover_opening_positions.values().rev()) .map(|(poly, positions)| open_polynomial(poly, positions)) .collect(); let bounds = LOG_DEGREES.map(CirclePolyDegreeBound::new).to_vec(); @@ -1076,11 +1183,16 @@ mod tests { #[test] fn proof_with_removed_layer_fails_verification() { const LOG_DEGREE: u32 = 6; - let polynomial = polynomial_evaluation(6, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(6, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![1], log_domain_size); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, queries.len()); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let proof = prover.decommit_on_queries(&queries); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; // Set verifier's config to expect one extra layer than prover config. @@ -1098,11 +1210,16 @@ mod tests { #[test] fn proof_with_added_layer_fails_verification() { const LOG_DEGREE: u32 = 6; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![1], log_domain_size); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, queries.len()); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let proof = prover.decommit_on_queries(&queries); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; // Set verifier's config to expect one less layer than prover config. @@ -1120,12 +1237,17 @@ mod tests { #[test] fn proof_with_invalid_inner_layer_evaluation_fails_verification() { const LOG_DEGREE: u32 = 6; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![5], log_domain_size); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, queries.len()); - let decommitment_value = query_polynomial(&polynomial, &queries); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let decommitment_value = query_polynomial(&evaluation, &queries); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let mut proof = prover.decommit_on_queries(&queries); // Remove an evaluation from the second layer's proof. @@ -1143,12 +1265,17 @@ mod tests { #[test] fn proof_with_invalid_inner_layer_decommitment_fails_verification() { const LOG_DEGREE: u32 = 6; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![5], log_domain_size); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, queries.len()); - let decommitment_value = query_polynomial(&polynomial, &queries); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let decommitment_value = query_polynomial(&evaluation, &queries); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let mut proof = prover.decommit_on_queries(&queries); // Modify the committed values in the second layer. @@ -1167,11 +1294,16 @@ mod tests { fn proof_with_invalid_last_layer_degree_fails_verification() { const LOG_DEGREE: u32 = 6; const LOG_MAX_LAST_LAYER_DEGREE: u32 = 2; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![1, 7, 8], log_domain_size); let config = FriConfig::new(LOG_MAX_LAST_LAYER_DEGREE, LOG_BLOWUP_FACTOR, queries.len()); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let mut proof = prover.decommit_on_queries(&queries); let bad_last_layer_coeffs = vec![One::one(); 1 << (LOG_MAX_LAST_LAYER_DEGREE + 1)]; @@ -1188,12 +1320,17 @@ mod tests { #[test] fn proof_with_invalid_last_layer_fails_verification() { const LOG_DEGREE: u32 = 6; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![1, 7, 8], log_domain_size); let config = FriConfig::new(2, LOG_BLOWUP_FACTOR, queries.len()); - let decommitment_value = query_polynomial(&polynomial, &queries); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let decommitment_value = query_polynomial(&evaluation, &queries); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let mut proof = prover.decommit_on_queries(&queries); // Compromise the last layer polynomial's first coefficient. @@ -1212,12 +1349,17 @@ mod tests { #[should_panic] fn decommit_queries_on_invalid_domain_fails_verification() { const LOG_DEGREE: u32 = 3; - let polynomial = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); - let log_domain_size = polynomial.domain.log_size(); + let evaluation = polynomial_evaluation(LOG_DEGREE, LOG_BLOWUP_FACTOR); + let log_domain_size = evaluation.domain.log_size(); let queries = Queries::from_positions(vec![5], log_domain_size); let config = FriConfig::new(1, LOG_BLOWUP_FACTOR, queries.len()); - let decommitment_value = query_polynomial(&polynomial, &queries); - let prover = FriProver::commit(&mut test_channel(), config, &[polynomial]); + let decommitment_value = query_polynomial(&evaluation, &queries); + let prover = FriProver::commit( + &mut test_channel(), + config, + &[evaluation.clone()], + &CPUBackend::precompute_twiddles(evaluation.domain.half_coset), + ); let proof = prover.decommit_on_queries(&queries); let bound = vec![CirclePolyDegreeBound::new(LOG_DEGREE)]; let verifier = FriVerifier::commit(&mut test_channel(), config, proof, bound).unwrap(); diff --git a/src/core/poly/circle/secure_poly.rs b/src/core/poly/circle/secure_poly.rs index 8902f49de..73c30f153 100644 --- a/src/core/poly/circle/secure_poly.rs +++ b/src/core/poly/circle/secure_poly.rs @@ -52,6 +52,7 @@ impl Deref for SecureCirclePoly { } } +#[derive(Clone)] pub struct SecureEvaluation> { pub domain: CircleDomain, pub values: SecureColumn, diff --git a/src/core/poly/line.rs b/src/core/poly/line.rs index 8390438ae..296684a47 100644 --- a/src/core/poly/line.rs +++ b/src/core/poly/line.rs @@ -6,6 +6,7 @@ use std::ops::{Deref, DerefMut}; use itertools::Itertools; use num_traits::Zero; +use super::circle::CircleDomain; use super::utils::fold; use crate::core::backend::{Backend, CPUBackend, ColumnOps}; use crate::core::circle::{CirclePoint, Coset, CosetIterator}; @@ -93,6 +94,14 @@ impl IntoIterator for LineDomain { } } +impl From for LineDomain { + fn from(domain: CircleDomain) -> Self { + Self { + coset: domain.half_coset, + } + } +} + /// An iterator over the x-coordinates of points in a coset. type LineDomainIterator = Map>, fn(CirclePoint) -> BaseField>; diff --git a/src/core/poly/utils.rs b/src/core/poly/utils.rs index 54f044cf2..60eb2dcb6 100644 --- a/src/core/poly/utils.rs +++ b/src/core/poly/utils.rs @@ -1,4 +1,4 @@ -use super::circle::CircleDomain; +use super::line::LineDomain; use crate::core::fields::{ExtensionOf, Field}; /// Folds values recursively in `O(n)` by a hierarchical application of folding factors. @@ -54,9 +54,15 @@ pub fn repeat_value(values: &[T], duplicity: usize) -> Vec { res } -/// Computes the line twiddles for a [CircleDomain] from the precomputed twiddles tree. -pub fn domain_line_twiddles_from_tree(domain: CircleDomain, twiddle_buffer: &[T]) -> Vec<&[T]> { - (0..domain.half_coset.log_size()) +/// Computes the line twiddles for a [`CircleDomain`] or a [`LineDomain`] from the precomputed +/// twiddles tree. +/// +/// [`CircleDomain`]: super::circle::CircleDomain +pub fn domain_line_twiddles_from_tree( + domain: impl Into, + twiddle_buffer: &[T], +) -> Vec<&[T]> { + (0..domain.into().coset().log_size()) .map(|i| { let len = 1 << i; &twiddle_buffer[twiddle_buffer.len() - len * 2..twiddle_buffer.len() - len] diff --git a/src/core/prover/mod.rs b/src/core/prover/mod.rs index bd30b8c6c..ac0e3bf60 100644 --- a/src/core/prover/mod.rs +++ b/src/core/prover/mod.rs @@ -3,7 +3,7 @@ use thiserror::Error; use super::commitment_scheme::{CommitmentSchemeProof, TreeVec}; use super::fri::FriVerificationError; -use super::poly::circle::{SecureCirclePoly, MAX_CIRCLE_DOMAIN_LOG_SIZE}; +use super::poly::circle::{CanonicCoset, PolyOps, SecureCirclePoly, MAX_CIRCLE_DOMAIN_LOG_SIZE}; use super::proof_of_work::ProofOfWorkVerificationError; use super::ColumnVec; use crate::commitment_scheme::blake2_hash::Blake2sHasher; @@ -92,7 +92,11 @@ pub fn prove( // Second tree - composition polynomial. sample_points.push(vec![vec![oods_point]; 4]); - let commitment_scheme_proof = commitment_scheme.prove_values(sample_points, channel); + // TODO(spapini): Precompute twiddles outside. + let twiddles = CPUBackend::precompute_twiddles( + CanonicCoset::new(composition_polynomial_log_degree_bound + 1).half_coset(), + ); + let commitment_scheme_proof = commitment_scheme.prove_values(sample_points, channel, &twiddles); // Evaluate composition polynomial at OODS point and check that it matches the trace OODS // values. This is a sanity check.