Logup trace generation

crates/prover/benches/poseidon.rs

@@ -1,27 +1,12 @@
 use criterion::{criterion_group, criterion_main, Criterion, Throughput};
-use stwo_prover::core::backend::simd::SimdBackend;
-use stwo_prover::core::channel::{Blake2sChannel, Channel};
-use stwo_prover::core::fields::m31::BaseField;
-use stwo_prover::core::fields::IntoSlice;
-use stwo_prover::core::vcs::blake2_hash::Blake2sHasher;
-use stwo_prover::core::vcs::hasher::Hasher;
-use stwo_prover::examples::poseidon::{gen_trace, PoseidonAir, PoseidonComponent};
-use stwo_prover::trace_generation::commit_and_prove;
+use stwo_prover::examples::poseidon::prove_poseidon;
 
 pub fn simd_poseidon(c: &mut Criterion) {
-    const LOG_N_ROWS: u32 = 15;
+    const LOG_N_INSTANCES: u32 = 18;
     let mut group = c.benchmark_group("poseidon2");
-    group.throughput(Throughput::Elements(1u64 << (LOG_N_ROWS + 3)));
-    group.bench_function(format!("poseidon2 2^{} instances", LOG_N_ROWS + 3), |b| {
-        b.iter(|| {
-            let component = PoseidonComponent {
-                log_n_rows: LOG_N_ROWS,
-            };
-            let trace = gen_trace(component.log_column_size());
-            let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));
-            let air = PoseidonAir { component };
-            commit_and_prove::<SimdBackend>(&air, channel, trace).unwrap()
-        });
+    group.throughput(Throughput::Elements(1u64 << LOG_N_INSTANCES));
+    group.bench_function(format!("poseidon2 2^{} instances", LOG_N_INSTANCES), |b| {
+        b.iter(|| prove_poseidon(LOG_N_INSTANCES));
     });
 }
 

crates/prover/src/constraint_framework/logup.rs

@@ -0,0 +1,150 @@
+use std::ops::{Add, Mul, Sub};
+
+use itertools::Itertools;
+use num_traits::Zero;
+
+use crate::core::backend::simd::column::SecureColumn;
+use crate::core::backend::simd::m31::LOG_N_LANES;
+use crate::core::backend::simd::prefix_sum::inclusive_prefix_sum;
+use crate::core::backend::simd::qm31::PackedSecureField;
+use crate::core::backend::simd::SimdBackend;
+use crate::core::backend::Column;
+use crate::core::channel::{Blake2sChannel, Channel};
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::fields::secure_column::SecureColumnByCoords;
+use crate::core::fields::FieldExpOps;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
+use crate::core::utils::shifted_secure_combination;
+use crate::core::ColumnVec;
+
+/// Interaction elements for the logup protocol.
+#[derive(Copy, Clone, Debug)]
+pub struct LookupElements {
+    pub z: SecureField,
+    pub alpha: SecureField,
+}
+impl LookupElements {
+    pub fn draw(channel: &mut Blake2sChannel) -> Self {
+        let [z, alpha] = channel.draw_felts(2).try_into().unwrap();
+        Self { z, alpha }
+    }
+    pub fn combine<F: Copy, EF>(&self, values: &[F]) -> EF
+    where
+        EF: Copy
+            + Zero
+            + Mul<EF, Output = EF>
+            + Add<F, Output = EF>
+            + Sub<EF, Output = EF>
+            + From<SecureField>,
+    {
+        shifted_secure_combination(values, EF::from(self.alpha), EF::from(self.z))
+    }
+}
+
+// SIMD backend generator for logup interaction trace.
+pub struct LogupTraceGenerator {
+    log_size: u32,
+    /// Current allocated interaction columns.
+    trace: Vec<SecureColumnByCoords<SimdBackend>>,
+    /// Denominator expressions (z + sum_i alpha^i * x_i) being generated for the current lookup.
+    denom: SecureColumn,
+    /// Preallocated buffer for the Inverses of the denominators.
+    denom_inv: SecureColumn,
+}
+impl LogupTraceGenerator {
+    pub fn new(log_size: u32) -> Self {
+        let trace = vec![];
+        let denom = SecureColumn::zeros(1 << log_size);
+        let denom_inv = SecureColumn::zeros(1 << log_size);
+        Self {
+            log_size,
+            trace,
+            denom,
+            denom_inv,
+        }
+    }
+
+    /// Allocate a new lookup column.
+    pub fn new_col(&mut self) -> LogupColGenerator<'_> {
+        let log_size = self.log_size;
+        LogupColGenerator {
+            gen: self,
+            numerator: SecureColumnByCoords::<SimdBackend>::zeros(1 << log_size),
+        }
+    }
+
+    /// Finalize the trace. Returns the trace and the claimed sum of the last column.
+    pub fn finalize(
+        mut self,
+    ) -> (
+        ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>,
+        SecureField,
+    ) {
+        // Prefix sum the last column.
+        let last_col_coords = self.trace.pop().unwrap().columns;
+        let coord_prefix_sum = last_col_coords.map(inclusive_prefix_sum);
+        self.trace.push(SecureColumnByCoords {
+            columns: coord_prefix_sum,
+        });
+        let claimed_sum = self.trace.last().unwrap().at(1);
+
+        let trace = self
+            .trace
+            .into_iter()
+            .flat_map(|eval| {
+                eval.columns.map(|c| {
+                    CircleEvaluation::<SimdBackend, _, BitReversedOrder>::new(
+                        CanonicCoset::new(self.log_size).circle_domain(),
+                        c,
+                    )
+                })
+            })
+            .collect_vec();
+        (trace, claimed_sum)
+    }
+}
+
+/// Trace generator for a single lookup column.
+pub struct LogupColGenerator<'a> {
+    gen: &'a mut LogupTraceGenerator,
+    /// Numerator expressions (i.e. multiplicities) being generated for the current lookup.
+    numerator: SecureColumnByCoords<SimdBackend>,
+}
+impl<'a> LogupColGenerator<'a> {
+    /// Write a fraction to the column at a row.
+    pub fn write_frac(
+        &mut self,
+        vec_row: usize,
+        numerator: PackedSecureField,
+        denom: PackedSecureField,
+    ) {
+        unsafe {
+            self.numerator.set_packed(vec_row, numerator);
+            *self.gen.denom.data.get_unchecked_mut(vec_row) = denom;
+        }
+    }
+
+    /// Finalizes generating the column.
+    pub fn finalize_col(mut self) {
+        FieldExpOps::batch_inverse(&self.gen.denom.data, &mut self.gen.denom_inv.data);
+
+        #[allow(clippy::needless_range_loop)]
+        for vec_row in 0..(1 << (self.gen.log_size - LOG_N_LANES)) {
+            unsafe {
+                let value = self.numerator.packed_at(vec_row)
+                    * *self.gen.denom_inv.data.get_unchecked(vec_row);
+                let prev_value = self
+                    .gen
+                    .trace
+                    .last()
+                    .map(|col| col.packed_at(vec_row))
+                    .unwrap_or_else(PackedSecureField::zero);
+                self.numerator.set_packed(vec_row, value + prev_value)
+            };
+        }
+
+        self.gen.trace.push(self.numerator)
+    }
+}

crates/prover/src/constraint_framework/mod.rs

@@ -2,6 +2,7 @@
 mod assert;
 pub mod constant_columns;
 mod info;
+pub mod logup;
 mod point;
 mod simd_domain;
 

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

@@ -279,6 +279,23 @@ impl Distribution<PackedQM31> for Standard {
     }
 }
 
+impl From<PackedM31> for PackedQM31 {
+    fn from(value: PackedM31) -> Self {
+        PackedQM31::from_packed_m31s([
+            value,
+            PackedM31::zero(),
+            PackedM31::zero(),
+            PackedM31::zero(),
+        ])
+    }
+}
+
+impl From<QM31> for PackedQM31 {
+    fn from(value: QM31) -> Self {
+        PackedQM31::broadcast(value)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::array;

crates/prover/src/core/fields/secure_column.rs

@@ -2,7 +2,6 @@ use super::m31::BaseField;
 use super::qm31::SecureField;
 use super::{ExtensionOf, FieldOps};
 use crate::core::backend::{Col, Column, CpuBackend};
-use crate::core::utils::IteratorMutExt;
 
 pub const SECURE_EXTENSION_DEGREE: usize =
     <SecureField as ExtensionOf<BaseField>>::EXTENSION_DEGREE;
@@ -14,13 +13,6 @@ pub struct SecureColumnByCoords<B: FieldOps<BaseField>> {
     pub columns: [Col<B, BaseField>; SECURE_EXTENSION_DEGREE],
 }
 impl SecureColumnByCoords<CpuBackend> {
-    pub fn set(&mut self, index: usize, value: SecureField) {
-        self.columns
-            .iter_mut()
-            .map(|c| &mut c[index])
-            .assign(value.to_m31_array());
-    }
-
     // TODO(spapini): Remove when we no longer use CircleEvaluation<SecureField>.
     pub fn to_vec(&self) -> Vec<SecureField> {
         (0..self.len()).map(|i| self.at(i)).collect()
@@ -50,6 +42,12 @@ impl<B: FieldOps<BaseField>> SecureColumnByCoords<B> {
             columns: self.columns.clone().map(|c| c.to_cpu()),
         }
     }
+
+    pub fn set(&mut self, index: usize, value: SecureField) {
+        for i in 0..SECURE_EXTENSION_DEGREE {
+            self.columns[i].set(index, value.to_m31_array()[i]);
+        }
+    }
 }
 
 pub struct SecureColumnByCoordsIter<'a> {