added coset_diff (#583)

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/583)
<!-- Reviewable:end -->

crates/prover/src/core/backend/avx512/fri.rs

@@ -1,8 +1,12 @@
-use super::AVX512Backend;
+use num_traits::Zero;
+
+use super::{AVX512Backend, PackedBaseField, K_BLOCK_SIZE};
 use crate::core::backend::avx512::fft::compute_first_twiddles;
 use crate::core::backend::avx512::fft::ifft::avx_ibutterfly;
 use crate::core::backend::avx512::qm31::PackedSecureField;
 use crate::core::backend::avx512::VECS_LOG_SIZE;
+use crate::core::backend::Column;
+use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::secure_column::SecureColumn;
 use crate::core::fri::{self, FriOps};
@@ -89,17 +93,53 @@ impl FriOps for AVX512Backend {
     }
 }
 
+impl AVX512Backend {
+    /// See [`CPUBackend::decomposition_coefficient`].
+    ///
+    /// [`CPUBackend::decomposition_coefficient`]:
+    /// crate::core::backend::cpu::CPUBackend::decomposition_coefficient
+    // TODO(Ohad): remove pub.
+    pub fn decomposition_coefficient(eval: &SecureEvaluation<Self>) -> SecureField {
+        let cols = &eval.values.columns;
+        let [mut x_sum, mut y_sum, mut z_sum, mut w_sum] = [PackedBaseField::zero(); 4];
+
+        let range = cols[0].len() / K_BLOCK_SIZE;
+        let (half_a, half_b) = (range / 2, range);
+
+        for i in 0..half_a {
+            x_sum += cols[0].data[i];
+            y_sum += cols[1].data[i];
+            z_sum += cols[2].data[i];
+            w_sum += cols[3].data[i];
+        }
+        for i in half_a..half_b {
+            x_sum -= cols[0].data[i];
+            y_sum -= cols[1].data[i];
+            z_sum -= cols[2].data[i];
+            w_sum -= cols[3].data[i];
+        }
+
+        let x = x_sum.pointwise_sum();
+        let y = y_sum.pointwise_sum();
+        let z = z_sum.pointwise_sum();
+        let w = w_sum.pointwise_sum();
+
+        SecureField::from_m31(x, y, z, w)
+            / BaseField::from_u32_unchecked(1 << eval.domain.log_size())
+    }
+}
+
 #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
 #[cfg(test)]
 mod tests {
-    use crate::core::backend::avx512::AVX512Backend;
-    use crate::core::backend::CPUBackend;
+    use crate::core::backend::avx512::{AVX512Backend, BaseFieldVec};
+    use crate::core::backend::{CPUBackend, Column};
     use crate::core::fields::qm31::SecureField;
     use crate::core::fields::secure_column::SecureColumn;
     use crate::core::fri::FriOps;
-    use crate::core::poly::circle::{CanonicCoset, PolyOps, SecureEvaluation};
+    use crate::core::poly::circle::{CanonicCoset, CirclePoly, PolyOps, SecureEvaluation};
     use crate::core::poly::line::{LineDomain, LineEvaluation};
-    use crate::qm31;
+    use crate::{m31, qm31};
 
     #[test]
     fn test_fold_line() {
@@ -160,4 +200,41 @@ mod tests {
 
         assert_eq!(cpu_fold.values.to_vec(), avx_fold.values.to_vec());
     }
+
+    #[test]
+    fn deocompose_coeff_out_fft_space_test() {
+        const DOMAIN_LOG_SIZE: u32 = 5;
+        const DOMAIN_LOG_HALF_SIZE: u32 = DOMAIN_LOG_SIZE - 1;
+        let s = CanonicCoset::new(DOMAIN_LOG_SIZE);
+        let domain = s.circle_domain();
+
+        let mut coeffs = BaseFieldVec::zeros(1 << DOMAIN_LOG_SIZE);
+
+        // Polynomial is out of FFT space.
+        coeffs.as_mut_slice()[1 << DOMAIN_LOG_HALF_SIZE] = m31!(1);
+        let poly = CirclePoly::<AVX512Backend>::new(coeffs);
+        let values = poly.evaluate(domain);
+
+        let avx_column = SecureColumn::<AVX512Backend> {
+            columns: [
+                values.values.clone(),
+                values.values.clone(),
+                values.values.clone(),
+                values.values.clone(),
+            ],
+        };
+        let avx_eval = SecureEvaluation {
+            domain,
+            values: avx_column.clone(),
+        };
+        let cpu_eval = SecureEvaluation::<CPUBackend> {
+            domain,
+            values: avx_eval.to_cpu(),
+        };
+        let cpu_lambda = CPUBackend::decomposition_coefficient(&cpu_eval);
+
+        let lambda = AVX512Backend::decomposition_coefficient(&avx_eval);
+
+        assert_eq!(lambda, cpu_lambda);
+    }
 }

crates/prover/src/core/backend/avx512/m31.rs

@@ -7,6 +7,7 @@ use std::arch::x86_64::{
     _mm512_store_epi32,
 };
 use std::fmt::Display;
+use std::iter::Sum;
 use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
 
 use num_traits::{One, Zero};
@@ -259,6 +260,12 @@ impl FieldExpOps for PackedBaseField {
     }
 }
 
+impl Sum for PackedBaseField {
+    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
+        iter.fold(Self::zero(), Add::add)
+    }
+}
+
 #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
 #[cfg(test)]
 mod tests {

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

@@ -1,3 +1,4 @@
+use std::iter::Sum;
 use std::ops::{Add, AddAssign, Mul, MulAssign, Sub};
 
 use bytemuck::{Pod, Zeroable};
@@ -154,6 +155,12 @@ impl Mul<PackedBaseField> for PackedSecureField {
     }
 }
 
+impl Sum for PackedSecureField {
+    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
+        iter.fold(Self::zero(), Add::add)
+    }
+}
+
 unsafe impl Pod for PackedSecureField {}
 unsafe impl Zeroable for PackedSecureField {
     fn zeroed() -> Self {

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

@@ -1,4 +1,5 @@
 use super::CPUBackend;
+use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fri::{fold_circle_into_line, fold_line, FriOps};
 use crate::core::poly::circle::SecureEvaluation;
@@ -23,3 +24,104 @@ impl FriOps for CPUBackend {
         fold_circle_into_line(dst, src, alpha)
     }
 }
+
+impl CPUBackend {
+    /// Used to decompose a general polynomial to a polynomial inside the fft-space, and
+    /// the remainder terms.
+    /// A coset-diff on a [`CirclePoly`] that is in the FFT space will return zero.
+    ///
+    /// Let N be the domain size, Let h be a coset size N/2. Using lemma #7 from the CircleStark
+    /// paper, <f,V_h> = lambda<V_h,V_h> = lambda\*N => lambda = f(0)\*V_h(0) + f(1)*V_h(1) + .. +
+    /// f(N-1)\*V_h(N-1). The Vanishing polynomial of a cannonic coset sized half the circle
+    /// domain,evaluated on the circle domain, is [(1, -1, -1, 1)] repeating. This becomes
+    /// alternating [+-1] in our NaturalOrder, and [(+, +, +, ... , -, -)] in bit reverse.
+    /// Explicitly, lambda\*N = sum(+f(0..N/2)) + sum(-f(N/2..)).
+    ///
+    /// # Warning
+    /// This function assumes the blowupfactor is 2
+    ///
+    /// [`CirclePoly`]: crate::core::poly::circle::CirclePoly
+    // TODO(Ohad): remove pub.
+    pub fn decomposition_coefficient(eval: &SecureEvaluation<Self>) -> SecureField {
+        let domain_size = 1 << eval.domain.log_size();
+        let half_domain_size = domain_size / 2;
+
+        // eval is in bit-reverse, hence all the positive factors are in the first half, opposite to
+        // the latter.
+        let a_sum = (0..half_domain_size)
+            .map(|i| eval.values.at(i))
+            .sum::<SecureField>();
+        let b_sum = (half_domain_size..domain_size)
+            .map(|i| eval.values.at(i))
+            .sum::<SecureField>();
+
+        // lambda = sum(+-f(p)) / 2N.
+        (a_sum - b_sum) / BaseField::from_u32_unchecked(domain_size as u32)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use num_traits::Zero;
+
+    use crate::core::backend::cpu::{CPUCircleEvaluation, CPUCirclePoly};
+    use crate::core::backend::CPUBackend;
+    use crate::core::fields::m31::BaseField;
+    use crate::core::fields::secure_column::SecureColumn;
+    use crate::core::poly::circle::{CanonicCoset, SecureEvaluation};
+    use crate::m31;
+
+    #[test]
+    fn decompose_coeff_out_fft_space_test() {
+        for domain_log_size in 5..12 {
+            let domain_log_half_size = domain_log_size - 1;
+            let s = CanonicCoset::new(domain_log_size);
+            let domain = s.circle_domain();
+
+            let mut coeffs = vec![BaseField::zero(); 1 << domain_log_size];
+
+            // Polynomial is out of FFT space.
+            coeffs[1 << domain_log_half_size] = m31!(1);
+            assert!(!CPUCirclePoly::new(coeffs.clone()).is_in_fft_space(domain_log_half_size));
+
+            let poly = CPUCirclePoly::new(coeffs);
+            let values = poly.evaluate(domain);
+            let secure_column = SecureColumn {
+                columns: [
+                    values.values.clone(),
+                    values.values.clone(),
+                    values.values.clone(),
+                    values.values.clone(),
+                ],
+            };
+            let secure_eval = SecureEvaluation::<CPUBackend> {
+                domain,
+                values: secure_column.clone(),
+            };
+
+            let lambda = CPUBackend::decomposition_coefficient(&secure_eval.clone());
+
+            // isolate the polynomial that is inside of the fft-space.
+            // TODO(Ohad): this should be a function in itself.
+            let q_x_values: SecureColumn<CPUBackend> = secure_column
+                .into_iter()
+                .enumerate()
+                .map(|(i, x)| {
+                    if i < (1 << domain_log_half_size) {
+                        x - lambda
+                    } else {
+                        x + lambda
+                    }
+                })
+                .collect();
+
+            // Assert the new polynomial is in the FFT space.
+            for i in 0..4 {
+                let basefield_column = q_x_values.columns[i].clone();
+                let eval = CPUCircleEvaluation::new(domain, basefield_column);
+                let coeffs = eval.interpolate().coeffs;
+                assert!(CPUCirclePoly::new(coeffs).is_in_fft_space(domain_log_half_size));
+            }
+        }
+    }
+}

crates/prover/src/core/fri.rs

@@ -73,7 +73,7 @@ impl FriConfig {
     }
 }
 
-pub trait FriOps: FieldOps<BaseField> + PolyOps + Sized {
+pub trait FriOps: FieldOps<BaseField> + PolyOps + Sized + FieldOps<SecureField> {
     /// Folds a degree `d` polynomial into a degree `d/2` polynomial.
     ///
     /// Let `eval` be a polynomial evaluated on a [LineDomain] `E`, `alpha` be a random field