Implement QuotientOps for SIMD backend

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

@@ -2,7 +2,7 @@ use self::bit_reverse::bit_reverse_m31;
 use self::column::{BaseFieldVec, SecureFieldVec};
 use self::m31::PackedBaseField;
 use self::qm31::PackedSecureField;
-use super::ColumnOps;
+use super::{Backend, ColumnOps};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::{FieldExpOps, FieldOps};
@@ -18,11 +18,14 @@ pub mod fft;
 pub mod fri;
 pub mod m31;
 pub mod qm31;
+pub mod quotients;
 mod utils;
 
 #[derive(Copy, Clone, Debug)]
 pub struct SimdBackend;
 
+impl Backend for SimdBackend {}
+
 impl ColumnOps<BaseField> for SimdBackend {
     type Column = BaseFieldVec;
 

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

@@ -0,0 +1,224 @@
+use itertools::{izip, zip_eq, Itertools};
+use num_traits::Zero;
+
+use super::m31::N_LANES;
+use super::qm31::PackedSecureField;
+use super::{SecureFieldVec, SimdBackend};
+use crate::core::backend::cpu::quotients::{
+    batch_random_coeffs, column_line_coeffs, QuotientConstants,
+};
+use crate::core::backend::simd::m31::LOG_N_LANES;
+use crate::core::backend::simd::PackedBaseField;
+use crate::core::backend::{Col, Column};
+use crate::core::circle::CirclePoint;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::fields::secure_column::SecureColumn;
+use crate::core::fields::{ComplexConjugate, FieldOps};
+use crate::core::pcs::quotients::{ColumnSampleBatch, QuotientOps};
+use crate::core::poly::circle::{CircleDomain, CircleEvaluation, SecureEvaluation};
+use crate::core::poly::BitReversedOrder;
+use crate::core::utils::bit_reverse_index;
+
+impl QuotientOps for SimdBackend {
+    fn accumulate_quotients(
+        domain: CircleDomain,
+        columns: &[&CircleEvaluation<Self, BaseField, BitReversedOrder>],
+        random_coeff: SecureField,
+        sample_batches: &[ColumnSampleBatch],
+    ) -> SecureEvaluation<Self> {
+        assert!(domain.log_size() >= LOG_N_LANES);
+        let mut values = SecureColumn::<Self>::zeros(domain.size());
+        let quotient_constants = quotient_constants(sample_batches, random_coeff, domain);
+
+        // TODO(spapini): bit reverse iterator.
+        for vec_row in 0..1 << (domain.log_size() - LOG_N_LANES) {
+            // TODO(spapini): Optimize this, for the small number of columns case.
+            let points = std::array::from_fn(|i| {
+                domain.at(bit_reverse_index(
+                    (vec_row << LOG_N_LANES) + i,
+                    domain.log_size(),
+                ))
+            });
+            let domain_points_x = PackedBaseField::from_array(points.map(|p| p.x));
+            let domain_points_y = PackedBaseField::from_array(points.map(|p| p.y));
+            let row_accumulator = accumulate_row_quotients(
+                sample_batches,
+                columns,
+                &quotient_constants,
+                vec_row,
+                (domain_points_x, domain_points_y),
+            );
+            unsafe { values.set_packed(vec_row, row_accumulator) };
+        }
+        SecureEvaluation { domain, values }
+    }
+}
+
+pub fn accumulate_row_quotients(
+    sample_batches: &[ColumnSampleBatch],
+    columns: &[&CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>],
+    quotient_constants: &QuotientConstants<SimdBackend>,
+    vec_row: usize,
+    domain_point_vec: (PackedBaseField, PackedBaseField),
+) -> PackedSecureField {
+    let mut row_accumulator = PackedSecureField::zero();
+    for (sample_batch, line_coeffs, batch_coeff, denominator_inverses) in izip!(
+        sample_batches,
+        &quotient_constants.line_coeffs,
+        &quotient_constants.batch_random_coeffs,
+        &quotient_constants.denominator_inverses
+    ) {
+        let mut numerator = PackedSecureField::zero();
+        for ((column_index, _), (a, b, c)) in zip_eq(&sample_batch.columns_and_values, line_coeffs)
+        {
+            let column = &columns[*column_index];
+            let value = PackedSecureField::broadcast(*c) * column.data[vec_row];
+            // The numerator is a line equation passing through
+            //   (sample_point.y, sample_value), (conj(sample_point), conj(sample_value))
+            // evaluated at (domain_point.y, value).
+            // When substituting a polynomial in this line equation, we get a polynomial with a root
+            // at sample_point and conj(sample_point) if the original polynomial had the values
+            // sample_value and conj(sample_value) at these points.
+            // TODO(AlonH): Use single point vanishing to save a multiplication.
+            let linear_term = PackedSecureField::broadcast(*a) * domain_point_vec.1
+                + PackedSecureField::broadcast(*b);
+            numerator += value - linear_term;
+        }
+
+        row_accumulator = row_accumulator * PackedSecureField::broadcast(*batch_coeff)
+            + numerator * denominator_inverses.data[vec_row];
+    }
+    row_accumulator
+}
+
+/// Pair vanishing for the packed representation of the points. See
+/// [crate::core::constraints::pair_vanishing] for more details.
+fn packed_pair_vanishing(
+    excluded0: CirclePoint<SecureField>,
+    excluded1: CirclePoint<SecureField>,
+    packed_p: (PackedBaseField, PackedBaseField),
+) -> PackedSecureField {
+    PackedSecureField::broadcast(excluded0.y - excluded1.y) * packed_p.0
+        + PackedSecureField::broadcast(excluded1.x - excluded0.x) * packed_p.1
+        + PackedSecureField::broadcast(excluded0.x * excluded1.y - excluded0.y * excluded1.x)
+}
+
+fn denominator_inverses(
+    sample_batches: &[ColumnSampleBatch],
+    domain: CircleDomain,
+) -> Vec<Col<SimdBackend, SecureField>> {
+    let flat_denominators: SecureFieldVec = sample_batches
+        .iter()
+        .flat_map(|sample_batch| {
+            (0..1 << (domain.log_size() - LOG_N_LANES))
+                .map(|vec_row| {
+                    // TODO(spapini): Optimize this, for the small number of columns case.
+                    let points = std::array::from_fn(|i| {
+                        domain.at(bit_reverse_index(
+                            (vec_row << LOG_N_LANES) + i,
+                            domain.log_size(),
+                        ))
+                    });
+                    let domain_points_x = PackedBaseField::from_array(points.map(|p| p.x));
+                    let domain_points_y = PackedBaseField::from_array(points.map(|p| p.y));
+                    let domain_point_vec = (domain_points_x, domain_points_y);
+                    packed_pair_vanishing(
+                        sample_batch.point,
+                        sample_batch.point.complex_conjugate(),
+                        domain_point_vec,
+                    )
+                })
+                .collect_vec()
+        })
+        .collect();
+
+    let mut flat_denominator_inverses = SecureFieldVec::zeros(flat_denominators.len());
+    <SimdBackend as FieldOps<SecureField>>::batch_inverse(
+        &flat_denominators,
+        &mut flat_denominator_inverses,
+    );
+
+    flat_denominator_inverses
+        .data
+        .chunks(domain.size() / N_LANES)
+        .map(|denominator_inverses| denominator_inverses.iter().copied().collect())
+        .collect()
+}
+
+fn quotient_constants(
+    sample_batches: &[ColumnSampleBatch],
+    random_coeff: SecureField,
+    domain: CircleDomain,
+) -> QuotientConstants<SimdBackend> {
+    let line_coeffs = column_line_coeffs(sample_batches, random_coeff);
+    let batch_random_coeffs = batch_random_coeffs(sample_batches, random_coeff);
+    let denominator_inverses = denominator_inverses(sample_batches, domain);
+    QuotientConstants {
+        line_coeffs,
+        batch_random_coeffs,
+        denominator_inverses,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use itertools::Itertools;
+
+    use crate::core::backend::simd::column::BaseFieldVec;
+    use crate::core::backend::simd::SimdBackend;
+    use crate::core::backend::{CPUBackend, Column};
+    use crate::core::circle::SECURE_FIELD_CIRCLE_GEN;
+    use crate::core::fields::m31::BaseField;
+    use crate::core::pcs::quotients::{ColumnSampleBatch, QuotientOps};
+    use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+    use crate::core::poly::BitReversedOrder;
+    use crate::qm31;
+
+    #[test]
+    fn test_accumulate_quotients() {
+        const LOG_SIZE: u32 = 8;
+        let domain = CanonicCoset::new(LOG_SIZE).circle_domain();
+        let e0: BaseFieldVec = (0..domain.size()).map(BaseField::from).collect();
+        let e1: BaseFieldVec = (0..domain.size()).map(|i| BaseField::from(2 * i)).collect();
+        let columns = vec![
+            CircleEvaluation::<SimdBackend, BaseField, BitReversedOrder>::new(domain, e0),
+            CircleEvaluation::<SimdBackend, BaseField, BitReversedOrder>::new(domain, e1),
+        ];
+        let random_coeff = qm31!(1, 2, 3, 4);
+        let a = qm31!(3, 6, 9, 12);
+        let b = qm31!(4, 8, 12, 16);
+        let samples = vec![ColumnSampleBatch {
+            point: SECURE_FIELD_CIRCLE_GEN,
+            columns_and_values: vec![(0, a), (1, b)],
+        }];
+        let cpu_columns = columns
+            .iter()
+            .map(|c| {
+                CircleEvaluation::<CPUBackend, _, BitReversedOrder>::new(
+                    c.domain,
+                    c.values.to_cpu(),
+                )
+            })
+            .collect::<Vec<_>>();
+        let cpu_result = CPUBackend::accumulate_quotients(
+            domain,
+            &cpu_columns.iter().collect_vec(),
+            random_coeff,
+            &samples,
+        )
+        .values
+        .to_vec();
+
+        let res = SimdBackend::accumulate_quotients(
+            domain,
+            &columns.iter().collect_vec(),
+            random_coeff,
+            &samples,
+        )
+        .values
+        .to_vec();
+
+        assert_eq!(res, cpu_result);
+    }
+}