SIMD Barycentric Evaluation and Test

Author: Gali-StarkWare

crates/prover/src/core/poly/circle/evaluation.rs

@@ -233,13 +233,126 @@ fn barycentric_eval_at_point(
     })
 }
 
+fn simd_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 simd_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()
+}
+
 #[cfg(test)]
 mod tests {
     use crate::core::backend::cpu::{CpuCircleEvaluation, CpuCirclePoly};
+    use crate::core::backend::simd::column::BaseColumn;
+    use crate::core::backend::simd::SimdBackend;
     use crate::core::backend::CpuBackend;
     use crate::core::circle::{CirclePoint, Coset};
     use crate::core::fields::m31::BaseField;
-    use crate::core::poly::circle::CanonicCoset;
+    use crate::core::poly::circle::{CanonicCoset, CirclePoly};
     use crate::core::poly::{BitReversedOrder, NaturalOrder};
     use crate::m31;
 
@@ -285,7 +398,7 @@ mod tests {
     }
 
     #[test]
-    fn test_barycentric_evaluation() {
+    fn test_cpu_barycentric_evaluation() {
         let poly = CpuCirclePoly::new(
             [691, 805673, 5, 435684, 4832, 23876431, 197, 897346068]
                 .map(BaseField::from)
@@ -310,9 +423,56 @@ mod tests {
         let sampled_barycentric_values = sampled_points
             .iter()
             .map(|point| {
-                eval.barycentric_eval_at_point(
+                super::barycentric_eval_at_point(
+                    eval,
+                    *point,
+                    &super::weights(
+                        eval.domain.log_size(),
+                        *point,
+                    ),
+                )
+            })
+            .collect::<Vec<_>>();
+
+        assert_eq!(
+            sampled_barycentric_values, sampled_values,
+            "Barycentric evaluation should be equal to the polynomial evaluation"
+        );
+    }
+
+    #[test]
+    fn test_simd_barycentric_evaluation() {
+        let poly = CirclePoly::<SimdBackend>::new(BaseColumn::from_cpu(
+            [691, 805673, 5, 435684, 4832, 23876431, 197, 897346068]
+                .map(BaseField::from)
+                .to_vec(),
+        ));
+        let s = CanonicCoset::new(3);
+        let domain = s.circle_domain();
+        let eval = poly.evaluate(domain);
+        let sampled_points = [
+            CirclePoint::get_point(348),
+            CirclePoint::get_point(9736524),
+            CirclePoint::get_point(13),
+            CirclePoint::get_point(346752),
+            domain.at(0).into_ef(),
+            domain.at(3).into_ef(),
+        ];
+        let sampled_values = sampled_points
+            .iter()
+            .map(|point| poly.eval_at_point(*point))
+            .collect::<Vec<_>>();
+
+        let sampled_barycentric_values = sampled_points
+            .iter()
+            .map(|point| {
+                super::simd_barycentric_eval_at_point(
+                    eval,
                     *point,
-                    &super::weights(eval.domain.log_size(), *point),
+                    &super::simd_weights(
+                        eval.domain.log_size(),
+                        *point,
+                    ),
                 )
             })
             .collect::<Vec<_>>();