Eval framework

crates/prover/benches/poseidon.rs

@@ -15,7 +15,7 @@ pub fn simd_poseidon(c: &mut Criterion) {
     group.bench_function(format!("poseidon2 2^{} instances", LOG_N_ROWS + 3), |b| {
         b.iter(|| {
             let component = PoseidonComponent {
-                log_n_instances: LOG_N_ROWS,
+                log_n_rows: LOG_N_ROWS,
             };
             let trace = gen_trace(component.log_column_size());
             let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));

crates/prover/src/builder/assert.rs

@@ -0,0 +1,51 @@
+use num_traits::{One, Zero};
+
+use super::EvalAtRow;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::pcs::TreeVec;
+
+pub struct AssertEvaluator<'a> {
+    pub trace: &'a TreeVec<Vec<Vec<BaseField>>>,
+    pub col_index: TreeVec<usize>,
+    pub row: usize,
+}
+impl<'a> AssertEvaluator<'a> {
+    pub fn new(trace: &'a TreeVec<Vec<Vec<BaseField>>>) -> Self {
+        Self {
+            trace,
+            col_index: TreeVec::new(vec![0; trace.len()]),
+            row: 0,
+        }
+    }
+}
+impl<'a> EvalAtRow for AssertEvaluator<'a> {
+    type F = BaseField;
+    type EF = SecureField;
+
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N] {
+        let col_index = self.col_index[interaction];
+        self.col_index[interaction] += 1;
+        offsets.map(|off| {
+            self.trace[interaction][col_index][(self.row as isize + off)
+                .rem_euclid(self.trace[interaction][col_index].len() as isize)
+                as usize]
+        })
+    }
+
+    fn add_constraint<G>(&mut self, constraint: G)
+    where
+        Self::EF: std::ops::Mul<G, Output = Self::EF>,
+    {
+        let res = SecureField::one() * constraint;
+        assert_eq!(res, SecureField::zero(), "row: {}", self.row);
+    }
+
+    fn combine_ef(values: [Self::F; 4]) -> Self::EF {
+        SecureField::from_m31_array(values)
+    }
+}

crates/prover/src/builder/domain.rs

@@ -0,0 +1,89 @@
+use std::ops::Mul;
+
+use num_traits::Zero;
+
+use super::EvalAtRow;
+use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES};
+use crate::core::backend::simd::qm31::PackedSecureField;
+use crate::core::backend::simd::SimdBackend;
+use crate::core::backend::Column;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::pcs::TreeVec;
+use crate::core::poly::circle::CircleEvaluation;
+use crate::core::poly::BitReversedOrder;
+use crate::core::utils::offset_bit_reversed_circle_domain_index;
+
+pub struct DomainEvaluator<'a> {
+    pub trace_eval:
+        &'a TreeVec<Vec<&'a CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>>,
+    pub col_index: Vec<usize>,
+    pub vec_row: usize,
+    pub random_coeff_powers: &'a [SecureField],
+    pub row_res: PackedSecureField,
+    pub constraint_index: usize,
+    pub domain_log_size: u32,
+    pub eval_log_size: u32,
+}
+impl<'a> DomainEvaluator<'a> {
+    pub fn new(
+        trace_eval: &'a TreeVec<Vec<&CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>>,
+        vec_row: usize,
+        random_coeff_powers: &'a [SecureField],
+        domain_log_size: u32,
+        eval_log_size: u32,
+    ) -> Self {
+        Self {
+            trace_eval,
+            col_index: vec![0; trace_eval.len()],
+            vec_row,
+            random_coeff_powers,
+            row_res: PackedSecureField::zero(),
+            constraint_index: 0,
+            domain_log_size,
+            eval_log_size,
+        }
+    }
+}
+impl<'a> EvalAtRow for DomainEvaluator<'a> {
+    type F = PackedBaseField;
+    type EF = PackedSecureField;
+
+    // TODO(spapini): Remove all boundary checks.
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N] {
+        let col_index = self.col_index[interaction];
+        self.col_index[interaction] += 1;
+        offsets.map(|off| {
+            // TODO(spapini): Optimize.
+            if off == 0 {
+                return self.trace_eval[interaction][col_index].data[self.vec_row];
+            }
+            PackedBaseField::from_array(std::array::from_fn(|i| {
+                let index = offset_bit_reversed_circle_domain_index(
+                    (self.vec_row << LOG_N_LANES) + i,
+                    self.domain_log_size,
+                    self.eval_log_size,
+                    off,
+                );
+                self.trace_eval[interaction][col_index].at(index)
+            }))
+        })
+    }
+    fn add_constraint<G>(&mut self, constraint: G)
+    where
+        Self::EF: Mul<G, Output = Self::EF>,
+    {
+        self.row_res +=
+            PackedSecureField::broadcast(self.random_coeff_powers[self.constraint_index])
+                * constraint;
+        self.constraint_index += 1;
+    }
+
+    fn combine_ef(values: [Self::F; 4]) -> Self::EF {
+        PackedSecureField::from_packed_m31s(values)
+    }
+}

crates/prover/src/builder/info.rs

@@ -0,0 +1,44 @@
+use std::ops::Mul;
+
+use num_traits::One;
+
+use super::EvalAtRow;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::pcs::TreeVec;
+
+#[derive(Default)]
+pub struct InfoEvaluator {
+    pub mask_offsets: TreeVec<Vec<Vec<isize>>>,
+    pub n_constraints: usize,
+}
+impl InfoEvaluator {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+impl EvalAtRow for InfoEvaluator {
+    type F = BaseField;
+    type EF = SecureField;
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N] {
+        if self.mask_offsets.len() <= interaction {
+            self.mask_offsets.resize(interaction + 1, vec![]);
+        }
+        self.mask_offsets[interaction].push(offsets.into_iter().collect());
+        [BaseField::one(); N]
+    }
+    fn add_constraint<G>(&mut self, _constraint: G)
+    where
+        Self::EF: Mul<G, Output = Self::EF>,
+    {
+        self.n_constraints += 1;
+    }
+
+    fn combine_ef(_values: [Self::F; 4]) -> Self::EF {
+        SecureField::one()
+    }
+}

crates/prover/src/builder/mod.rs

@@ -0,0 +1,55 @@
+mod assert;
+mod domain;
+mod info;
+mod point;
+
+use std::fmt::Debug;
+use std::ops::{Add, AddAssign, Mul, Sub};
+
+pub use assert::AssertEvaluator;
+pub use domain::DomainEvaluator;
+pub use info::InfoEvaluator;
+use num_traits::{One, Zero};
+pub use point::PointEvaluator;
+
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::fields::FieldExpOps;
+
+pub trait EvalAtRow {
+    type F: FieldExpOps
+        + Copy
+        + Debug
+        + AddAssign<Self::F>
+        + AddAssign<BaseField>
+        + Add<Self::F, Output = Self::F>
+        + Sub<Self::F, Output = Self::F>
+        + Mul<BaseField, Output = Self::F>
+        + Add<SecureField, Output = Self::EF>
+        + Mul<SecureField, Output = Self::EF>
+        + From<BaseField>;
+    type EF: One
+        + Copy
+        + Debug
+        + Zero
+        + Add<SecureField, Output = Self::EF>
+        + Sub<SecureField, Output = Self::EF>
+        + Mul<SecureField, Output = Self::EF>
+        + Add<Self::F, Output = Self::EF>
+        + Mul<Self::F, Output = Self::EF>
+        + Sub<Self::EF, Output = Self::EF>
+        + Mul<Self::EF, Output = Self::EF>;
+
+    fn next_mask(&mut self) -> Self::F {
+        self.next_interaction_mask(0, [0])[0]
+    }
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N];
+    fn add_constraint<G>(&mut self, constraint: G)
+    where
+        Self::EF: Mul<G, Output = Self::EF>;
+    fn combine_ef(values: [Self::F; 4]) -> Self::EF;
+}

crates/prover/src/builder/point.rs

@@ -0,0 +1,57 @@
+use std::ops::Mul;
+
+use super::EvalAtRow;
+use crate::core::air::accumulation::PointEvaluationAccumulator;
+use crate::core::backend::CpuBackend;
+use crate::core::fields::qm31::SecureField;
+use crate::core::pcs::TreeVec;
+use crate::core::poly::circle::SecureCirclePoly;
+use crate::core::ColumnVec;
+
+pub struct PointEvaluator<'a> {
+    pub mask: TreeVec<&'a ColumnVec<Vec<SecureField>>>,
+    pub evaluation_accumulator: &'a mut PointEvaluationAccumulator,
+    pub col_index: Vec<usize>,
+    pub denom_inverse: SecureField,
+}
+impl<'a> PointEvaluator<'a> {
+    pub fn new(
+        mask: TreeVec<&'a ColumnVec<Vec<SecureField>>>,
+        evaluation_accumulator: &'a mut PointEvaluationAccumulator,
+        denom_inverse: SecureField,
+    ) -> Self {
+        let col_index = vec![0; mask.len()];
+        Self {
+            mask,
+            evaluation_accumulator,
+            col_index,
+            denom_inverse,
+        }
+    }
+}
+impl<'a> EvalAtRow for PointEvaluator<'a> {
+    type F = SecureField;
+    type EF = SecureField;
+
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        _offsets: [isize; N],
+    ) -> [Self::F; N] {
+        let col_index = self.col_index[interaction];
+        self.col_index[interaction] += 1;
+        let mask = self.mask[interaction][col_index].clone();
+        assert_eq!(mask.len(), N);
+        mask.try_into().unwrap()
+    }
+    fn add_constraint<G>(&mut self, constraint: G)
+    where
+        Self::EF: Mul<G, Output = Self::EF>,
+    {
+        self.evaluation_accumulator
+            .accumulate(self.denom_inverse * constraint);
+    }
+    fn combine_ef(values: [Self::F; 4]) -> Self::EF {
+        SecureCirclePoly::<CpuBackend>::eval_from_partial_evals(values)
+    }
+}

crates/prover/src/core/backend/simd/m31.rs

@@ -253,6 +253,12 @@ impl From<[BaseField; N_LANES]> for PackedM31 {
     }
 }
 
+impl From<BaseField> for PackedM31 {
+    fn from(v: BaseField) -> Self {
+        Self::broadcast(v)
+    }
+}
+
 impl Distribution<PackedM31> for Standard {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> PackedM31 {
         PackedM31::from_array(rng.gen())