Barycentric Evaluation in Poly Ops

Author: Gali-StarkWare

crates/prover/src/core/backend/cpu/circle.rs

@@ -1,16 +1,22 @@
-use num_traits::Zero;
+use itertools::Itertools;
+use num_traits::{One, Zero};
 
 use super::CpuBackend;
 use crate::core::backend::cpu::bit_reverse;
-use crate::core::circle::{CirclePoint, Coset};
+use crate::core::backend::Col;
+use crate::core::circle::{CirclePoint, CirclePointIndex, Coset};
+use crate::core::constraints::{coset_vanishing, coset_vanishing_derivative, point_vanishing};
 use crate::core::fft::{butterfly, ibutterfly};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::{batch_inverse_in_place, ExtensionOf};
-use crate::core::poly::circle::{CircleDomain, CircleEvaluation, CirclePoly, PolyOps};
+use crate::core::poly::circle::{
+    CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
+};
 use crate::core::poly::twiddles::TwiddleTree;
 use crate::core::poly::utils::{domain_line_twiddles_from_tree, fold};
 use crate::core::poly::BitReversedOrder;
+use crate::core::utils::bit_reverse_index;
 
 impl PolyOps for CpuBackend {
     type Twiddles = Vec<BaseField>;
@@ -86,6 +92,84 @@ impl PolyOps for CpuBackend {
         fold(&poly.coeffs, &mappings)
     }
 
+    fn weights(log_size: u32, sample_point: CirclePoint<SecureField>) -> Col<Self, SecureField> {
+        let domain = CanonicCoset::new(log_size).circle_domain();
+
+        for i in 0..domain.size() {
+            if domain.at(i).into_ef() == sample_point {
+                let mut weights = vec![SecureField::zero(); domain.size()];
+                weights[i] = SecureField::one();
+                return weights;
+            }
+        }
+
+        let p_0 = domain.at(0).into_ef::<SecureField>();
+        let weights_first_half = SecureField::one()
+            / (-(p_0.y + p_0.y)
+                * coset_vanishing_derivative(
+                    Coset::new(CirclePointIndex::generator(), domain.log_size()),
+                    p_0,
+                ));
+        let p_0_inverse = domain.at(domain.half_coset.size()).into_ef::<SecureField>();
+        let weights_second_half = SecureField::one()
+            / (-(p_0_inverse.y + p_0_inverse.y)
+                * coset_vanishing_derivative(
+                    Coset::new(CirclePointIndex::generator(), domain.log_size()),
+                    p_0_inverse,
+                ));
+
+        let domain_points_vanishing_evaluated_at_point = (0..domain.size())
+            .map(|i| {
+                point_vanishing(
+                    domain.at(i).into_ef::<SecureField>(),
+                    sample_point.into_ef::<SecureField>(),
+                )
+            })
+            .collect_vec();
+        let mut inversed_domain_points_vanishing_evaluated_at_point =
+            vec![unsafe { std::mem::zeroed() }; domain.size()];
+
+        batch_inverse_in_place(
+            &domain_points_vanishing_evaluated_at_point,
+            &mut inversed_domain_points_vanishing_evaluated_at_point,
+        );
+
+        let coset_vanishing_evaluated_at_point: SecureField = coset_vanishing(
+            CanonicCoset::new(domain.log_size()).coset,
+            sample_point.into_ef::<SecureField>(),
+        );
+
+        (0..domain.size())
+            .map(|i| {
+                if i < domain.half_coset.size() {
+                    weights_first_half
+                        * inversed_domain_points_vanishing_evaluated_at_point[i]
+                        * coset_vanishing_evaluated_at_point
+                } else {
+                    weights_second_half
+                        * inversed_domain_points_vanishing_evaluated_at_point[i]
+                        * coset_vanishing_evaluated_at_point
+                }
+            })
+            .collect_vec()
+    }
+
+    fn barycentric_eval_at_point(
+        evals: &CircleEvaluation<Self, BaseField, BitReversedOrder>,
+        point: CirclePoint<SecureField>,
+        weights: &Col<Self, SecureField>,
+    ) -> SecureField {
+        for i in 0..evals.domain.size() {
+            if point == evals.domain.at(i).into_ef() {
+                return evals.values[bit_reverse_index(i, evals.domain.log_size())].into();
+            }
+        }
+
+        (0..evals.domain.size()).fold(SecureField::zero(), |acc, i| {
+            acc + (evals.values[bit_reverse_index(i, evals.domain.log_size())] * weights[i])
+        })
+    }
+
     fn extend(poly: &CirclePoly<Self>, log_size: u32) -> CirclePoly<Self> {
         assert!(log_size >= poly.log_size());
         let mut coeffs = Vec::with_capacity(1 << log_size);

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

@@ -3,6 +3,8 @@ use std::mem::transmute;
 use std::simd::Simd;
 
 use bytemuck::Zeroable;
+use itertools::Itertools;
+use num_traits::{One, Zero};
 #[cfg(feature = "parallel")]
 use rayon::prelude::*;
 
@@ -14,10 +16,11 @@ use crate::core::backend::cpu::circle::slow_precompute_twiddles;
 use crate::core::backend::simd::column::BaseColumn;
 use crate::core::backend::simd::m31::PackedM31;
 use crate::core::backend::{Col, Column, CpuBackend};
-use crate::core::circle::{CirclePoint, Coset, M31_CIRCLE_LOG_ORDER};
+use crate::core::circle::{CirclePoint, CirclePointIndex, Coset, M31_CIRCLE_LOG_ORDER};
+use crate::core::constraints::{coset_vanishing, coset_vanishing_derivative, point_vanishing};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
-use crate::core::fields::{Field, FieldExpOps};
+use crate::core::fields::{batch_inverse_in_place, Field, FieldExpOps};
 use crate::core::poly::circle::{
     CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
 };
@@ -221,6 +224,117 @@ impl PolyOps for SimdBackend {
         (sum * twiddle_lows).pointwise_sum()
     }
 
+    fn weights(log_size: u32, sample_point: CirclePoint<SecureField>) -> Col<Self, SecureField> {
+        let domain = CanonicCoset::new(log_size).circle_domain();
+        let weights_vec_len = domain.size().div_ceil(N_LANES);
+
+        for i in 0..domain.size() {
+            if domain.at(i).into_ef() == sample_point {
+                let mut weights = Col::<Self, SecureField>::zeros(domain.size());
+                weights.set(i, SecureField::one());
+                return weights;
+            }
+        }
+
+        let p_0 = domain.at(0).into_ef::<SecureField>();
+        let weights_first_half = SecureField::one()
+            / (-(p_0.y + p_0.y)
+                * coset_vanishing_derivative(
+                    Coset::new(CirclePointIndex::generator(), domain.log_size()),
+                    p_0,
+                ));
+        let p_0_inverse = domain.at(domain.half_coset.size()).into_ef::<SecureField>();
+        let weights_second_half = SecureField::one()
+            / (-(p_0_inverse.y + p_0_inverse.y)
+                * coset_vanishing_derivative(
+                    Coset::new(CirclePointIndex::generator(), domain.log_size()),
+                    p_0_inverse,
+                ));
+
+        let domain_points_vanishing_evaluated_at_point = (0..weights_vec_len)
+            .map(|i| {
+                PackedSecureField::from_array(std::array::from_fn(|j| {
+                    if domain.size() <= i * N_LANES + j {
+                        SecureField::one()
+                    } else {
+                        point_vanishing(
+                            domain.at(i * N_LANES + j).into_ef::<SecureField>(),
+                            sample_point.into_ef::<SecureField>(),
+                        )
+                    }
+                }))
+            })
+            .collect_vec();
+        let mut inversed_domain_points_vanishing_evaluated_at_point =
+            vec![unsafe { std::mem::zeroed() }; weights_vec_len];
+
+        batch_inverse_in_place(
+            &domain_points_vanishing_evaluated_at_point,
+            &mut inversed_domain_points_vanishing_evaluated_at_point,
+        );
+
+        let coset_vanishing_evaluated_at_point: SecureField = coset_vanishing(
+            CanonicCoset::new(domain.log_size()).coset,
+            sample_point.into_ef::<SecureField>(),
+        );
+
+        if weights_vec_len == 1 {
+            return (0..N_LANES)
+                .map(|i| {
+                    let inversed_domain_points_vanishing_evaluated_at_point =
+                        inversed_domain_points_vanishing_evaluated_at_point[0].to_array();
+                    if i < domain.size() / 2 {
+                        inversed_domain_points_vanishing_evaluated_at_point[i]
+                            * (weights_first_half * coset_vanishing_evaluated_at_point)
+                    } else {
+                        inversed_domain_points_vanishing_evaluated_at_point[i]
+                            * (weights_second_half * coset_vanishing_evaluated_at_point)
+                    }
+                })
+                .collect();
+        }
+
+        let weights: Col<Self, SecureField> = (0..weights_vec_len)
+            .map(|i| {
+                if i < weights_vec_len / 2 {
+                    inversed_domain_points_vanishing_evaluated_at_point[i]
+                        * (weights_first_half * coset_vanishing_evaluated_at_point)
+                } else {
+                    inversed_domain_points_vanishing_evaluated_at_point[i]
+                        * (weights_second_half * coset_vanishing_evaluated_at_point)
+                }
+            })
+            .collect();
+
+        weights
+    }
+
+    fn barycentric_eval_at_point(
+        evals: &CircleEvaluation<Self, BaseField, BitReversedOrder>,
+        point: CirclePoint<SecureField>,
+        weights: &Col<Self, SecureField>,
+    ) -> SecureField {
+        for i in 0..evals.domain.size() {
+            if point == evals.domain.at(i).into_ef() {
+                return evals
+                    .values
+                    .at(bit_reverse_index(i, evals.domain.log_size()))
+                    .into();
+            }
+        }
+        let evals = evals.clone().bit_reverse();
+        if evals.domain.size() < N_LANES {
+            return (0..evals.domain.size()).fold(SecureField::zero(), |acc, i| {
+                acc + (weights.at(i) * evals.values.at(i))
+            });
+        }
+        (0..evals.domain.size().div_ceil(N_LANES))
+            .fold(PackedSecureField::zero(), |acc, i| {
+                acc + (weights.data[i] * evals.values.data[i])
+            })
+            .pointwise_sum()
+    }
+
     fn extend(poly: &CirclePoly<Self>, log_size: u32) -> CirclePoly<Self> {
         // TODO(shahars): Get rid of extends.
         poly.evaluate(CanonicCoset::new(log_size).circle_domain())