diff --git a/crates/prover/src/constraint_framework/constant_columns.rs b/crates/prover/src/constraint_framework/constant_columns.rs index f9d2ebc0c..e57df28ab 100644 --- a/crates/prover/src/constraint_framework/constant_columns.rs +++ b/crates/prover/src/constraint_framework/constant_columns.rs @@ -4,6 +4,7 @@ use crate::core::backend::{Backend, Col, Column}; use crate::core::fields::m31::BaseField; use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; use crate::core::poly::BitReversedOrder; +use crate::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index}; /// Generates a column with a single one at the first position, and zeros elsewhere. pub fn gen_is_first(log_size: u32) -> CircleEvaluation { @@ -11,3 +12,26 @@ pub fn gen_is_first(log_size: u32) -> CircleEvaluation( + log_size: u32, + log_step: u32, + offset: usize, +) -> CircleEvaluation { + let mut col = Col::::zeros(1 << log_size); + + let size = 1 << log_size; + let step = 1 << log_step; + let step_offset = offset % step; + + for i in (step_offset..size).step_by(step) { + let circle_domain_index = coset_index_to_circle_domain_index(i, log_size); + let circle_domain_index_bit_rev = bit_reverse_index(circle_domain_index, log_size); + col.set(circle_domain_index_bit_rev, BaseField::one()); + } + + CircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), col) +} diff --git a/crates/prover/src/core/backend/simd/column.rs b/crates/prover/src/core/backend/simd/column.rs index 664e8f31a..db98d2669 100644 --- a/crates/prover/src/core/backend/simd/column.rs +++ b/crates/prover/src/core/backend/simd/column.rs @@ -1,4 +1,5 @@ -use std::mem; +use std::iter::zip; +use std::{array, mem}; use bytemuck::allocation::cast_vec; use bytemuck::{cast_slice, cast_slice_mut, Zeroable}; @@ -215,6 +216,29 @@ pub struct SecureColumn { pub length: usize, } +impl SecureColumn { + // Separates a single column of `PackedSecureField` elements into `SECURE_EXTENSION_DEGREE` many + // `PackedBaseField` coordinate columns. + pub fn into_secure_column_by_coords(self) -> SecureColumnByCoords { + if self.len() < N_LANES { + return self.to_cpu().into_iter().collect(); + } + + let length = self.length; + let packed_length = self.data.len(); + let mut columns = array::from_fn(|_| Vec::with_capacity(packed_length)); + + for v in self.data { + let packed_coords = v.into_packed_m31s(); + zip(&mut columns, packed_coords).for_each(|(col, packed_coord)| col.push(packed_coord)); + } + + SecureColumnByCoords { + columns: columns.map(|col| BaseColumn { data: col, length }), + } + } +} + impl Column for SecureColumn { fn zeros(length: usize) -> Self { Self { @@ -276,9 +300,8 @@ impl FromIterator for SecureColumn { impl FromIterator for SecureColumn { fn from_iter>(iter: I) -> Self { - let data = (&mut iter.into_iter()).collect_vec(); + let data = iter.into_iter().collect_vec(); let length = data.len() * N_LANES; - Self { data, length } } } diff --git a/crates/prover/src/core/fields/secure_column.rs b/crates/prover/src/core/fields/secure_column.rs index 3663c204d..073d21bc2 100644 --- a/crates/prover/src/core/fields/secure_column.rs +++ b/crates/prover/src/core/fields/secure_column.rs @@ -1,3 +1,6 @@ +use std::array; +use std::iter::zip; + use super::m31::BaseField; use super::qm31::SecureField; use super::{ExtensionOf, FieldOps}; @@ -89,13 +92,15 @@ impl<'a> IntoIterator for &'a SecureColumnByCoords { } impl FromIterator for SecureColumnByCoords { 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]); - } + let values = iter.into_iter(); + let (lower_bound, _) = values.size_hint(); + let mut columns = array::from_fn(|_| Vec::with_capacity(lower_bound)); + + for value in values { + let coords = value.to_m31_array(); + zip(&mut columns, coords).for_each(|(col, coord)| col.push(coord)); } + SecureColumnByCoords { columns } } } diff --git a/crates/prover/src/core/utils.rs b/crates/prover/src/core/utils.rs index be08c6572..334edb71d 100644 --- a/crates/prover/src/core/utils.rs +++ b/crates/prover/src/core/utils.rs @@ -121,11 +121,15 @@ pub(crate) fn coset_order_to_circle_domain_order(values: &[F]) -> Vec< circle_domain_order } -pub fn coset_order_to_circle_domain_order_index(index: usize, log_size: u32) -> usize { - if index & 1 == 0 { - index / 2 +/// Converts an index within a [`Coset`] to the corresponding index in a [`CircleDomain`]. +/// +/// [`CircleDomain`]: crate::core::poly::circle::CircleDomain +/// [`Coset`]: crate::core::circle::Coset +pub fn coset_index_to_circle_domain_index(coset_index: usize, log_domain_size: u32) -> usize { + if coset_index % 2 == 0 { + coset_index / 2 } else { - (1 << log_size) - (index + 1) / 2 + ((2 << log_domain_size) - coset_index) / 2 } } diff --git a/crates/prover/src/examples/xor/eq_eval_constraints.rs b/crates/prover/src/examples/xor/eq_eval_constraints.rs new file mode 100644 index 000000000..8a4566833 --- /dev/null +++ b/crates/prover/src/examples/xor/eq_eval_constraints.rs @@ -0,0 +1,219 @@ +use std::array; + +use num_traits::{One, Zero}; + +use crate::constraint_framework::EvalAtRow; +use crate::core::fields::qm31::SecureField; +use crate::core::lookups::utils::eq; + +/// Evaluates EqEvals constraints on a column. +/// +/// Returns the evaluation at offset 0 on the column. +/// +/// Given a column `c(P)` defined on a circle domain D, and an MLE evaluation point `(r0, r1, ...)` +/// evaluates constraints that guarantee: `c(D[b0,b1,...]) = eq((b0,b1,...), (r0,r1,...))` +/// +/// See (Section 5.1). +pub fn eval_eq_constraints< + E: EvalAtRow, + const N_VARIABLES: usize, + const EQ_EVALS_TRACE: usize, + const SELECTOR_TRACE: usize, +>( + eval: &mut E, + mle_eval_point: MleEvalPoint, +) -> E::EF { + let [curr, next_next] = eval.next_extension_interaction_mask(EQ_EVALS_TRACE, [0, 2]); + let [is_first, is_second] = eval.next_interaction_mask(SELECTOR_TRACE, [0, -1]); + + // Check the initial value on half_coset0 and final value on half_coset1. + // Combining these constraints is safe because `is_first` and `is_second` are never + // non-zero at the same time on the trace. + let half_coset0_initial_check = (curr - mle_eval_point.eq_0_p) * is_first; + let half_coset1_final_check = (curr - mle_eval_point.eq_1_p) * is_second; + eval.add_constraint(half_coset0_initial_check + half_coset1_final_check); + + // Check all variables except the last (last variable is handled by the constraint above). + #[allow(clippy::needless_range_loop)] + for variable_i in 0..N_VARIABLES.saturating_sub(1) { + let half_coset0_next = next_next; + let half_coset1_prev = next_next; + let [half_coset0_step, half_coset1_step] = + eval.next_interaction_mask(SELECTOR_TRACE, [0, -1]); + let carry_quotient = mle_eval_point.eq_carry_quotients[variable_i]; + // Safe to combine these constraints as `is_step.half_coset0` and `is_step.half_coset1` + // are never non-zero at the same time on the trace. + let half_coset0_check = (curr - half_coset0_next * carry_quotient) * half_coset0_step; + let half_coset1_check = (curr * carry_quotient - half_coset1_prev) * half_coset1_step; + eval.add_constraint(half_coset0_check + half_coset1_check); + } + + curr +} + +#[derive(Debug, Clone, Copy)] +pub struct MleEvalPoint { + // Equals `eq({0}^|p|, p)`. + eq_0_p: SecureField, + // Equals `eq({1}^|p|, p)`. + eq_1_p: SecureField, + // Index `i` stores `eq(({1}^|i|, 0), p[0..i+1]) / eq(({0}^|i|, 1), p[0..i+1])`. + eq_carry_quotients: [SecureField; N_VARIABLES], + // Point `p`. + _p: [SecureField; N_VARIABLES], +} + +impl MleEvalPoint { + /// Creates new metadata from point `p`. + pub fn new(p: [SecureField; N_VARIABLES]) -> Self { + let zero = SecureField::zero(); + let one = SecureField::one(); + + Self { + eq_0_p: eq(&[zero; N_VARIABLES], &p), + eq_1_p: eq(&[one; N_VARIABLES], &p), + eq_carry_quotients: array::from_fn(|i| { + let mut numer_assignment = vec![one; i + 1]; + numer_assignment[i] = zero; + let mut denom_assignment = vec![zero; i + 1]; + denom_assignment[i] = one; + eq(&numer_assignment, &p[..i + 1]) / eq(&denom_assignment, &p[..i + 1]) + }), + _p: p, + } + } +} + +#[cfg(test)] +pub mod tests { + use std::array; + + use num_traits::One; + use rand::rngs::SmallRng; + use rand::{Rng, SeedableRng}; + use test_log::test; + + use super::MleEvalPoint; + use crate::constraint_framework::assert_constraints; + use crate::constraint_framework::constant_columns::{gen_is_first, gen_is_step_with_offset}; + use crate::core::backend::simd::SimdBackend; + use crate::core::fields::m31::BaseField; + use crate::core::fields::qm31::SecureField; + use crate::core::lookups::gkr_prover::GkrOps; + use crate::core::pcs::TreeVec; + use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; + use crate::core::poly::BitReversedOrder; + use crate::core::ColumnVec; + use crate::examples::xor::eq_eval_constraints::eval_eq_constraints; + + const EVALS_TRACE: usize = 0; + const CONST_TRACE: usize = 1; + + #[test] + #[ignore = "SimdBackend `MIN_FFT_LOG_SIZE` is 5"] + fn eq_constraints_with_4_variables() { + const N_VARIABLES: usize = 4; + let mut rng = SmallRng::seed_from_u64(0); + let eval_point: [SecureField; N_VARIABLES] = array::from_fn(|_| rng.gen()); + let base_trace = gen_evals_trace(&eval_point); + let constants_trace = gen_constants_trace(N_VARIABLES); + let traces = TreeVec::new(vec![base_trace, constants_trace]); + let trace_polys = traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect()); + let trace_domain = CanonicCoset::new(eval_point.len() as u32); + let mle_eval_point = MleEvalPoint::new(eval_point); + + assert_constraints(&trace_polys, trace_domain, |mut eval| { + eval_eq_constraints::<_, N_VARIABLES, EVALS_TRACE, CONST_TRACE>( + &mut eval, + mle_eval_point, + ); + }); + } + + #[test] + fn eq_constraints_with_5_variables() { + const N_VARIABLES: usize = 5; + let mut rng = SmallRng::seed_from_u64(0); + let eval_point: [SecureField; N_VARIABLES] = array::from_fn(|_| rng.gen()); + let base_trace = gen_evals_trace(&eval_point); + let constants_trace = gen_constants_trace(N_VARIABLES); + let traces = TreeVec::new(vec![base_trace, constants_trace]); + let trace_polys = traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect()); + let trace_domain = CanonicCoset::new(eval_point.len() as u32); + let mle_eval_point = MleEvalPoint::new(eval_point); + + assert_constraints(&trace_polys, trace_domain, |mut eval| { + eval_eq_constraints::<_, N_VARIABLES, EVALS_TRACE, CONST_TRACE>( + &mut eval, + mle_eval_point, + ); + }); + } + + #[test] + fn eq_constraints_with_8_variables() { + const N_VARIABLES: usize = 8; + let mut rng = SmallRng::seed_from_u64(0); + let eval_point: [SecureField; N_VARIABLES] = array::from_fn(|_| rng.gen()); + let base_trace = gen_evals_trace(&eval_point); + let constants_trace = gen_constants_trace(N_VARIABLES); + let traces = TreeVec::new(vec![base_trace, constants_trace]); + let trace_polys = traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect()); + let trace_domain = CanonicCoset::new(eval_point.len() as u32); + let mle_eval_point = MleEvalPoint::new(eval_point); + + assert_constraints(&trace_polys, trace_domain, |mut eval| { + eval_eq_constraints::<_, N_VARIABLES, EVALS_TRACE, CONST_TRACE>( + &mut eval, + mle_eval_point, + ); + }); + } + + /// Generates a trace. + /// + /// Trace structure: + /// + /// ```text + /// ----------------------------- + /// | eq evals | + /// ----------------------------- + /// | c0 | c1 | c2 | c3 | + /// ----------------------------- + /// ``` + pub fn gen_evals_trace( + eval_point: &[SecureField], + ) -> ColumnVec> { + // TODO(andrew): Consider storing eq evals as a SecureColumn. + let eq_evals = SimdBackend::gen_eq_evals(eval_point, SecureField::one()).into_evals(); + let eq_evals_coordinate_columns = eq_evals.into_secure_column_by_coords().columns; + + let n_variables = eval_point.len(); + let domain = CanonicCoset::new(n_variables as u32).circle_domain(); + eq_evals_coordinate_columns + .map(|col| CircleEvaluation::new(domain, col)) + .into() + } + + pub fn gen_constants_trace( + n_variables: usize, + ) -> ColumnVec> { + let log_size = n_variables as u32; + let mut constants_trace = Vec::new(); + constants_trace.push(gen_is_first(log_size)); + + // TODO(andrew): Note the last selector column is not needed. The column for `is_first` + // with an offset for each half coset midpoint can be used instead. + for variable_i in 1..n_variables as u32 { + let half_coset_log_step = variable_i; + let half_coset_offset = (1 << (half_coset_log_step - 1)) - 1; + + let log_step = half_coset_log_step + 1; + let offset = half_coset_offset * 2; + + constants_trace.push(gen_is_step_with_offset(log_size, log_step, offset)) + } + + constants_trace + } +} diff --git a/crates/prover/src/examples/xor/mod.rs b/crates/prover/src/examples/xor/mod.rs index 8cd1686f1..80f94597d 100644 --- a/crates/prover/src/examples/xor/mod.rs +++ b/crates/prover/src/examples/xor/mod.rs @@ -1 +1,2 @@ +pub mod eq_eval_constraints; pub mod prefix_sum_constraints; diff --git a/crates/prover/src/examples/xor/prefix_sum_constraints.rs b/crates/prover/src/examples/xor/prefix_sum_constraints.rs index 306bf06c4..5982d2433 100644 --- a/crates/prover/src/examples/xor/prefix_sum_constraints.rs +++ b/crates/prover/src/examples/xor/prefix_sum_constraints.rs @@ -29,16 +29,15 @@ impl PrefixSumMask { #[cfg(test)] mod tests { use itertools::Itertools; - use num_traits::One; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; use test_log::test; use super::inclusive_prefix_sum_check; + use crate::constraint_framework::constant_columns::gen_is_first; use crate::constraint_framework::{assert_constraints, EvalAtRow}; use crate::core::backend::simd::prefix_sum::inclusive_prefix_sum; use crate::core::backend::simd::SimdBackend; - use crate::core::backend::{Col, Column}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::secure_column::SecureColumnByCoords; @@ -117,10 +116,6 @@ mod tests { fn gen_constants_trace( log_size: u32, ) -> Vec> { - let trace_domain = CanonicCoset::new(log_size).circle_domain(); - // Column is `1` at the first trace point and `0` on all other trace points. - let mut is_first = Col::::zeros(1 << log_size); - is_first.as_mut_slice()[0] = BaseField::one(); - vec![CircleEvaluation::new(trace_domain, is_first)] + vec![gen_is_first(log_size)] } } diff --git a/crates/prover/src/lib.rs b/crates/prover/src/lib.rs index d820075ce..a06b0c2d1 100644 --- a/crates/prover/src/lib.rs +++ b/crates/prover/src/lib.rs @@ -1,18 +1,20 @@ +#![allow(incomplete_features)] #![feature( array_methods, array_chunks, - iter_array_chunks, + assert_matches, exact_size_is_empty, - is_sorted, - new_uninit, - slice_group_by, - stdsimd, + generic_const_exprs, get_many_mut, int_roundings, + is_sorted, + iter_array_chunks, + new_uninit, + portable_simd, slice_first_last_chunk, slice_flatten, - assert_matches, - portable_simd + slice_group_by, + stdsimd )] pub mod constraint_framework; pub mod core;