Added LogupAtRow functionality to EvalAtRow.

starkware-libs · Nov 10, 2024 · f4aa5d7 · f4aa5d7
1 parent 4aa5a22
commit f4aa5d7
Show file tree

Hide file tree

Showing 19 changed files with 230 additions and 122 deletions.
diff --git a/crates/prover/src/constraint_framework/assert.rs b/crates/prover/src/constraint_framework/assert.rs
@@ -1,10 +1,12 @@
 use num_traits::{One, Zero};
 
+use super::logup::LogupAtRow;
 use super::EvalAtRow;
 use crate::core::backend::{Backend, Column};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::lookups::utils::Fraction;
 use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::{CanonicCoset, CirclePoly};
 use crate::core::utils::circle_domain_order_to_coset_order;
@@ -14,13 +16,15 @@ pub struct AssertEvaluator<'a> {
     pub trace: &'a TreeVec<Vec<Vec<BaseField>>>,
     pub col_index: TreeVec<usize>,
     pub row: usize,
+    pub logup: LogupAtRow<Self>,
 }
 impl<'a> AssertEvaluator<'a> {
     pub fn new(trace: &'a TreeVec<Vec<Vec<BaseField>>>, row: usize) -> Self {
         Self {
             trace,
             col_index: TreeVec::new(vec![0; trace.len()]),
             row,
+            logup: LogupAtRow::dummy(),
         }
     }
 }
@@ -57,6 +61,8 @@ impl<'a> EvalAtRow for AssertEvaluator<'a> {
     fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {
         SecureField::from_m31_array(values)
     }
+
+    super::logup_proxy!();
 }
 
 pub fn assert_constraints<B: Backend>(

diff --git a/crates/prover/src/constraint_framework/cpu_domain.rs b/crates/prover/src/constraint_framework/cpu_domain.rs
@@ -2,11 +2,13 @@ use std::ops::Mul;
 
 use num_traits::Zero;
 
+use super::logup::LogupAtRow;
 use super::EvalAtRow;
 use crate::core::backend::CpuBackend;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::lookups::utils::Fraction;
 use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::CircleEvaluation;
 use crate::core::poly::BitReversedOrder;
@@ -22,6 +24,7 @@ pub struct CpuDomainEvaluator<'a> {
     pub constraint_index: usize,
     pub domain_log_size: u32,
     pub eval_domain_log_size: u32,
+    pub logup: LogupAtRow<Self>,
 }
 
 impl<'a> CpuDomainEvaluator<'a> {
@@ -42,6 +45,7 @@ impl<'a> CpuDomainEvaluator<'a> {
             constraint_index: 0,
             domain_log_size,
             eval_domain_log_size: eval_log_size,
+            logup: LogupAtRow::dummy(),
         }
     }
 }
@@ -88,4 +92,6 @@ impl<'a> EvalAtRow for CpuDomainEvaluator<'a> {
     fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {
         SecureField::from_m31_array(values)
     }
+
+    super::logup_proxy!();
 }
diff --git a/crates/prover/src/constraint_framework/expr.rs b/crates/prover/src/constraint_framework/expr.rs
@@ -2,10 +2,12 @@ use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub};
 
 use num_traits::{One, Zero};
 
+use super::logup::LogupAtRow;
 use super::EvalAtRow;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::FieldExpOps;
+use crate::core::lookups::utils::Fraction;
 
 /// A single base field column at index `idx` of interaction `interaction`, at mask offset `offset`.
 #[derive(Clone, Debug, PartialEq)]
@@ -149,8 +151,9 @@ impl AddAssign<BaseField> for Expr {
 /// An Evaluator that saves all constraint expressions.
 #[derive(Default)]
 struct ExprEvaluator {
-    cur_var_index: usize,
-    constraints: Vec<Expr>,
+    pub cur_var_index: usize,
+    pub constraints: Vec<Expr>,
+    pub logup: LogupAtRow<Self>,
 }
 
 impl EvalAtRow for ExprEvaluator {
@@ -189,6 +192,8 @@ impl EvalAtRow for ExprEvaluator {
             Box::new(values[3].clone()),
         ])
     }
+
+    super::logup_proxy!();
 }
 
 #[cfg(test)]

diff --git a/crates/prover/src/constraint_framework/info.rs b/crates/prover/src/constraint_framework/info.rs
@@ -2,9 +2,11 @@ use std::ops::Mul;
 
 use num_traits::One;
 
+use super::logup::LogupAtRow;
 use super::EvalAtRow;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
+use crate::core::lookups::utils::Fraction;
 use crate::core::pcs::TreeVec;
 
 /// Collects information about the constraints.
@@ -13,6 +15,7 @@ use crate::core::pcs::TreeVec;
 pub struct InfoEvaluator {
     pub mask_offsets: TreeVec<Vec<Vec<isize>>>,
     pub n_constraints: usize,
+    pub logup: LogupAtRow<Self>,
 }
 impl InfoEvaluator {
     pub fn new() -> Self {
@@ -45,4 +48,6 @@ impl EvalAtRow for InfoEvaluator {
     fn combine_ef(_values: [Self::F; 4]) -> Self::EF {
         SecureField::one()
     }
+
+    super::logup_proxy!();
 }
diff --git a/crates/prover/src/constraint_framework/interaction_gen.rs b/crates/prover/src/constraint_framework/interaction_gen.rs
@@ -0,0 +1,7 @@
+use crate::core::backend::simd::SimdBackend;
+use crate::core::fields::m31::BaseField;
+use crate::core::pcs::TreeVec;
+use crate::core::poly::circle::CircleEvaluation;
+use crate::core::poly::BitReversedOrder;
+
+pub struct InteractionTraceGenerator<'a> {}
diff --git a/crates/prover/src/constraint_framework/logup.rs b/crates/prover/src/constraint_framework/logup.rs
@@ -44,6 +44,12 @@ pub struct LogupAtRow<E: EvalAtRow> {
     /// See [`super::preprocessed_columns::gen_is_first()`].
     pub is_first: E::F,
 }
+
+impl<E: EvalAtRow> Default for LogupAtRow<E> {
+    fn default() -> Self {
+        Self::dummy()
+    }
+}
 impl<E: EvalAtRow> LogupAtRow<E> {
     pub fn new(
         interaction: usize,
@@ -62,6 +68,19 @@ impl<E: EvalAtRow> LogupAtRow<E> {
         }
     }
 
+    // TODO(alont): Remove this once unnecessary LogupAtRows are gone.
+    pub fn dummy() -> Self {
+        Self {
+            interaction: 0,
+            total_sum: SecureField::one(),
+            claimed_sum: None,
+            prev_col_cumsum: E::EF::zero(),
+            cur_frac: None,
+            is_finalized: true,
+            is_first: E::F::zero(),
+        }
+    }
+
     pub fn write_frac(&mut self, eval: &mut E, fraction: Fraction<E::EF, E::EF>) {
         // Add a constraint that num / denom = diff.
         if let Some(cur_frac) = self.cur_frac.clone() {
@@ -73,7 +92,7 @@ impl<E: EvalAtRow> LogupAtRow<E> {
         self.cur_frac = Some(fraction);
     }
 
-    pub fn finalize(mut self, eval: &mut E) {
+    pub fn finalize(&mut self, eval: &mut E) {
         assert!(!self.is_finalized, "LogupAtRow was already finalized");
 
         let frac = self.cur_frac.clone().unwrap();

diff --git a/crates/prover/src/constraint_framework/mod.rs b/crates/prover/src/constraint_framework/mod.rs
@@ -25,6 +25,7 @@ use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
 use crate::core::fields::FieldExpOps;
+use crate::core::lookups::utils::Fraction;
 
 pub const ORIGINAL_TRACE_IDX: usize = 0;
 pub const INTERACTION_TRACE_IDX: usize = 1;
@@ -109,4 +110,93 @@ pub trait EvalAtRow {
 
     /// Combines 4 base field values into a single extension field value.
     fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF;
+
+    /// Adds `elems` to `relation` with `multiplicity`.
+    fn add_to_relation<Relation: RelationType<Self::F, Self::EF>>(
+        &mut self,
+        relation: Relation,
+        multiplicity: usize,
+        elems: &[Self::F],
+    ) {
+        let denom = relation.combine(elems);
+        self.write_frac(Fraction::new(
+            Self::EF::from(SecureField::from(multiplicity)),
+            denom,
+        ));
+    }
+
+    // TODO(alont): Remove these once LogupAtRow is no longer used.
+    fn init_logup(
+        &mut self,
+        _total_sum: SecureField,
+        _claimed_sum: Option<crate::constraint_framework::logup::ClaimedPrefixSum>,
+        _log_size: u32,
+    ) {
+        unimplemented!()
+    }
+    fn write_frac(&mut self, _fraction: Fraction<Self::EF, Self::EF>) {
+        unimplemented!()
+    }
+    fn finalize_logup(&mut self) {
+        unimplemented!()
+    }
+}
+
+/// Default implementation for evaluators that have an element called "logup" that works like a
+/// LogupAtRow, where the logup functionality can be proxied.
+/// TODO(alont): Remove once LogupAtRow is no longer used.
+macro_rules! logup_proxy {
+    () => {
+        fn init_logup(
+            &mut self,
+            total_sum: SecureField,
+            claimed_sum: Option<crate::constraint_framework::logup::ClaimedPrefixSum>,
+            log_size: u32,
+        ) {
+            let is_first = self.get_preprocessed_column(
+                crate::constraint_framework::preprocessed_columns::PreprocessedColumn::IsFirst(
+                    log_size,
+                ),
+            );
+            self.logup = crate::constraint_framework::logup::LogupAtRow::new(
+                crate::constraint_framework::INTERACTION_TRACE_IDX,
+                total_sum,
+                claimed_sum,
+                is_first,
+            );
+        }
+
+        fn write_frac(&mut self, fraction: Fraction<Self::EF, Self::EF>) {
+            let mut logup = std::mem::take(&mut self.logup);
+            logup.write_frac(self, fraction);
+            self.logup = logup;
+        }
+
+        fn finalize_logup(&mut self) {
+            let mut logup = std::mem::take(&mut self.logup);
+            logup.finalize(self);
+            self.logup = logup;
+        }
+    };
+}
+pub(crate) use logup_proxy;
+
+pub trait RelationType<F, EF>
+where
+    F: Clone,
+    EF: Clone + Zero + From<F> + From<SecureField> + Mul<F, Output = EF> + Sub<EF, Output = EF>,
+{
+    fn combine(&self, values: &[F]) -> EF {
+        values
+            .iter()
+            .zip(self.get_alpha_powers())
+            .fold(EF::zero(), |acc, (value, power)| {
+                acc + power.clone() * value.clone()
+            })
+            - self.get_z()
+    }
+
+    fn get_z(&self) -> EF;
+    fn get_alpha_powers(&self) -> &[EF];
+    fn name(&self) -> &str;
 }
diff --git a/crates/prover/src/constraint_framework/point.rs b/crates/prover/src/constraint_framework/point.rs
@@ -1,9 +1,11 @@
 use std::ops::Mul;
 
+use super::logup::LogupAtRow;
 use super::EvalAtRow;
 use crate::core::air::accumulation::PointEvaluationAccumulator;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::lookups::utils::Fraction;
 use crate::core::pcs::TreeVec;
 use crate::core::ColumnVec;
 
@@ -13,6 +15,7 @@ pub struct PointEvaluator<'a> {
     pub evaluation_accumulator: &'a mut PointEvaluationAccumulator,
     pub col_index: Vec<usize>,
     pub denom_inverse: SecureField,
+    pub logup: LogupAtRow<Self>,
 }
 impl<'a> PointEvaluator<'a> {
     pub fn new(
@@ -26,6 +29,7 @@ impl<'a> PointEvaluator<'a> {
             evaluation_accumulator,
             col_index,
             denom_inverse,
+            logup: LogupAtRow::dummy(),
         }
     }
 }
@@ -54,4 +58,6 @@ impl<'a> EvalAtRow for PointEvaluator<'a> {
     fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {
         SecureField::from_partial_evals(values)
     }
+
+    super::logup_proxy!();
 }
diff --git a/crates/prover/src/constraint_framework/simd_domain.rs b/crates/prover/src/constraint_framework/simd_domain.rs
@@ -2,6 +2,7 @@ use std::ops::Mul;
 
 use num_traits::Zero;
 
+use super::logup::LogupAtRow;
 use super::EvalAtRow;
 use crate::core::backend::simd::column::VeryPackedBaseColumn;
 use crate::core::backend::simd::m31::LOG_N_LANES;
@@ -13,6 +14,7 @@ use crate::core::backend::Column;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::lookups::utils::Fraction;
 use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::CircleEvaluation;
 use crate::core::poly::BitReversedOrder;
@@ -30,6 +32,7 @@ pub struct SimdDomainEvaluator<'a> {
     pub constraint_index: usize,
     pub domain_log_size: u32,
     pub eval_domain_log_size: u32,
+    pub logup: LogupAtRow<Self>,
 }
 impl<'a> SimdDomainEvaluator<'a> {
     pub fn new(
@@ -48,6 +51,7 @@ impl<'a> SimdDomainEvaluator<'a> {
             constraint_index: 0,
             domain_log_size,
             eval_domain_log_size: eval_log_size,
+            logup: LogupAtRow::dummy(),
         }
     }
 }
@@ -103,4 +107,6 @@ impl<'a> EvalAtRow for SimdDomainEvaluator<'a> {
     fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {
         VeryPackedSecureField::from_very_packed_m31s(values)
     }
+
+    super::logup_proxy!();
 }