Use precomputed twiddles in avx

benches/fft.rs

@@ -21,13 +21,13 @@ pub fn avx512_ifft(c: &mut criterion::Criterion) {
 
     // Compute.
     let mut values = BaseFieldVec::from_iter(values);
-    let twiddle_dbls = get_itwiddle_dbls(domain);
+    let twiddle_dbls = get_itwiddle_dbls(domain.half_coset);
 
     c.bench_function("avx ifft", |b| {
         b.iter(|| unsafe {
             ifft::ifft(
                 std::mem::transmute(values.data.as_mut_ptr()),
-                &twiddle_dbls[1..]
+                &twiddle_dbls
                     .iter()
                     .map(|x| x.as_slice())
                     .collect::<Vec<_>>(),

src/core/backend/avx512/circle.rs

@@ -13,13 +13,16 @@ use crate::core::poly::circle::{
     CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
 };
 use crate::core::poly::twiddles::TwiddleTree;
-use crate::core::poly::utils::fold;
+use crate::core::poly::utils::{domain_line_twiddles_from_tree, fold};
 use crate::core::poly::BitReversedOrder;
 
 // TODO(spapini): Everything is returned in redundant representation, where values can also be P.
 // Decide if and when it's ok and what to do if it's not.
 impl PolyOps for AVX512Backend {
-    type Twiddles = ();
+    // The twiddles type is i32, and not BaseField. This is because the fast AVX mul implementation
+    //  requries one of the numbers to be shifted left by 1 bit. This is not a reduced
+    //  representation of the field.
+    type Twiddles = Vec<i32>;
 
     fn new_canonical_ordered(
         coset: CanonicCoset,
@@ -32,23 +35,18 @@ impl PolyOps for AVX512Backend {
 
     fn interpolate(
         eval: CircleEvaluation<Self, BaseField, BitReversedOrder>,
-        _itwiddles: &TwiddleTree<Self>,
+        twiddles: &TwiddleTree<Self>,
     ) -> CirclePoly<Self> {
         let mut values = eval.values;
         let log_size = values.length.ilog2();
 
-        // TODO(spapini): Precompute twiddles.
-        let twiddle_dbls = ifft::get_itwiddle_dbls(eval.domain);
-        // TODO(spapini): Handle small cases.
+        let twiddles = domain_line_twiddles_from_tree(eval.domain, &twiddles.itwiddles);
 
         // Safe because [PackedBaseField] is aligned on 64 bytes.
         unsafe {
             ifft::ifft(
                 std::mem::transmute(values.data.as_mut_ptr()),
-                &twiddle_dbls[1..]
-                    .iter()
-                    .map(|x| x.as_slice())
-                    .collect::<Vec<_>>(),
+                &twiddles,
                 log_size as usize,
             );
         }
@@ -98,7 +96,7 @@ impl PolyOps for AVX512Backend {
     fn evaluate(
         poly: &CirclePoly<Self>,
         domain: CircleDomain,
-        _twiddles: &TwiddleTree<Self>,
+        twiddles: &TwiddleTree<Self>,
     ) -> CircleEvaluation<Self, BaseField, BitReversedOrder> {
         // TODO(spapini): Precompute twiddles.
         // TODO(spapini): Handle small cases.
@@ -109,17 +107,17 @@ impl PolyOps for AVX512Backend {
             "Can only evaluate on larger domains"
         );
 
-        let twiddles = rfft::get_twiddle_dbls(domain);
+        let twiddles = domain_line_twiddles_from_tree(domain, &twiddles.twiddles);
 
         // Evaluate on a big domains by evaluating on several subdomains.
         let log_subdomains = log_size - fft_log_size;
         let mut values = Vec::with_capacity(domain.size() >> VECS_LOG_SIZE);
         for i in 0..(1 << log_subdomains) {
             // The subdomain twiddles are a slice of the large domain twiddles.
-            let subdomain_twiddles = (1..fft_log_size)
+            let subdomain_twiddles = (0..(fft_log_size - 1))
                 .map(|layer_i| {
                     &twiddles[layer_i]
-                        [i << (fft_log_size - 1 - layer_i)..(i + 1) << (fft_log_size - 1 - layer_i)]
+                        [i << (fft_log_size - 2 - layer_i)..(i + 1) << (fft_log_size - 2 - layer_i)]
                 })
                 .collect::<Vec<_>>();
 
@@ -150,10 +148,27 @@ impl PolyOps for AVX512Backend {
     }
 
     fn precompute_twiddles(coset: Coset) -> TwiddleTree<Self> {
+        let mut twiddles = Vec::with_capacity(coset.size());
+        let mut itwiddles = Vec::with_capacity(coset.size());
+
+        // TODO(spapini): Optimize.
+        for layer in &rfft::get_twiddle_dbls(coset) {
+            twiddles.extend(layer);
+        }
+        // Pad by any value, to make the size a power of 2.
+        twiddles.push(1);
+        assert_eq!(twiddles.len(), coset.size());
+        for layer in &ifft::get_itwiddle_dbls(coset) {
+            itwiddles.extend(layer);
+        }
+        // Pad by any value, to make the size a power of 2.
+        itwiddles.push(1);
+        assert_eq!(itwiddles.len(), coset.size());
+
         TwiddleTree {
             root_coset: coset,
-            twiddles: (),
-            itwiddles: (),
+            twiddles,
+            itwiddles,
         }
     }
 }

src/core/backend/avx512/fft/ifft.rs

@@ -8,8 +8,8 @@ use std::arch::x86_64::{
 use super::{compute_first_twiddles, EVENS_INTERLEAVE_EVENS, ODDS_INTERLEAVE_ODDS};
 use crate::core::backend::avx512::fft::{transpose_vecs, CACHED_FFT_LOG_SIZE, MIN_FFT_LOG_SIZE};
 use crate::core::backend::avx512::{PackedBaseField, VECS_LOG_SIZE};
+use crate::core::circle::Coset;
 use crate::core::fields::FieldExpOps;
-use crate::core::poly::circle::CircleDomain;
 use crate::core::utils::bit_reverse;
 
 /// Performs an Inverse Circle Fast Fourier Transform (ICFFT) on the given values.
@@ -374,17 +374,9 @@ pub unsafe fn vecwise_ibutterflies(
     val0.deinterleave_with(val1)
 }
 
-pub fn get_itwiddle_dbls(domain: CircleDomain) -> Vec<Vec<i32>> {
-    let mut coset = domain.half_coset;
-
+/// Returns the line twiddles (x points) for an ifft on a coset.
+pub fn get_itwiddle_dbls(mut coset: Coset) -> Vec<Vec<i32>> {
     let mut res = vec![];
-    res.push(
-        coset
-            .iter()
-            .map(|p| (p.y.inverse().0 * 2) as i32)
-            .collect::<Vec<_>>(),
-    );
-    bit_reverse(res.last_mut().unwrap());
     for _ in 0..coset.log_size() {
         res.push(
             coset
@@ -643,19 +635,19 @@ mod tests {
     #[test]
     fn test_vecwise_ibutterflies() {
         let domain = CanonicCoset::new(5).circle_domain();
-        let twiddle_dbls = get_itwiddle_dbls(domain);
-        assert_eq!(twiddle_dbls.len(), 5);
+        let twiddle_dbls = get_itwiddle_dbls(domain.half_coset);
+        assert_eq!(twiddle_dbls.len(), 4);
         let values0: [i32; 16] = std::array::from_fn(|i| i as i32);
         let values1: [i32; 16] = std::array::from_fn(|i| (i + 16) as i32);
         let result: [BaseField; 32] = unsafe {
             let (val0, val1) = vecwise_ibutterflies(
                 std::mem::transmute(values0),
                 std::mem::transmute(values1),
+                twiddle_dbls[0].clone().try_into().unwrap(),
                 twiddle_dbls[1].clone().try_into().unwrap(),
                 twiddle_dbls[2].clone().try_into().unwrap(),
-                twiddle_dbls[3].clone().try_into().unwrap(),
             );
-            let (val0, val1) = avx_ibutterfly(val0, val1, _mm512_set1_epi32(twiddle_dbls[4][0]));
+            let (val0, val1) = avx_ibutterfly(val0, val1, _mm512_set1_epi32(twiddle_dbls[3][0]));
             std::mem::transmute([val0, val1])
         };
 
@@ -681,12 +673,12 @@ mod tests {
 
             // Compute.
             let mut values = BaseFieldVec::from_iter(values);
-            let twiddle_dbls = get_itwiddle_dbls(domain);
+            let twiddle_dbls = get_itwiddle_dbls(domain.half_coset);
 
             unsafe {
                 ifft_lower_with_vecwise(
                     std::mem::transmute(values.data.as_mut_ptr()),
-                    &twiddle_dbls[1..]
+                    &twiddle_dbls
                         .iter()
                         .map(|x| x.as_slice())
                         .collect::<Vec<_>>(),
@@ -709,12 +701,12 @@ mod tests {
 
         // Compute.
         let mut values = BaseFieldVec::from_iter(values);
-        let twiddle_dbls = get_itwiddle_dbls(domain);
+        let twiddle_dbls = get_itwiddle_dbls(domain.half_coset);
 
         unsafe {
             ifft(
                 std::mem::transmute(values.data.as_mut_ptr()),
-                &twiddle_dbls[1..]
+                &twiddle_dbls
                     .iter()
                     .map(|x| x.as_slice())
                     .collect::<Vec<_>>(),

src/core/backend/avx512/fft/mod.rs

@@ -98,34 +98,3 @@ unsafe fn compute_first_twiddles(twiddle1_dbl: [i32; 8]) -> (__m512i, __m512i) {
     let t0 = _mm512_xor_epi32(_mm512_permutexvar_epi32(INDICES_FROM_T1, t1), NEGATION_MASK);
     (t0, t1)
 }
-
-#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
-#[cfg(test)]
-mod tests {
-    use self::ifft::get_itwiddle_dbls;
-    use super::*;
-    use crate::core::fields::m31::BaseField;
-    use crate::core::poly::circle::CanonicCoset;
-
-    #[test]
-    fn test_twiddle_relation() {
-        let ts = get_itwiddle_dbls(CanonicCoset::new(5).circle_domain());
-        let t0 = ts[0]
-            .iter()
-            .copied()
-            .map(|x| BaseField::from_u32_unchecked((x as u32) / 2))
-            .collect::<Vec<_>>();
-        let t1 = ts[1]
-            .iter()
-            .copied()
-            .map(|x| BaseField::from_u32_unchecked((x as u32) / 2))
-            .collect::<Vec<_>>();
-
-        for i in 0..t0.len() / 4 {
-            assert_eq!(t0[i * 4], t1[i * 2 + 1]);
-            assert_eq!(t0[i * 4 + 1], -t1[i * 2 + 1]);
-            assert_eq!(t0[i * 4 + 2], -t1[i * 2]);
-            assert_eq!(t0[i * 4 + 3], t1[i * 2]);
-        }
-    }
-}

src/core/backend/avx512/fft/rfft.rs

@@ -8,7 +8,7 @@ use std::arch::x86_64::{
 use super::{compute_first_twiddles, EVENS_INTERLEAVE_EVENS, ODDS_INTERLEAVE_ODDS};
 use crate::core::backend::avx512::fft::{transpose_vecs, CACHED_FFT_LOG_SIZE, MIN_FFT_LOG_SIZE};
 use crate::core::backend::avx512::{PackedBaseField, VECS_LOG_SIZE};
-use crate::core::poly::circle::CircleDomain;
+use crate::core::circle::Coset;
 use crate::core::utils::bit_reverse;
 
 /// Performs a Circle Fast Fourier Transform (ICFFT) on the given values.
@@ -347,12 +347,9 @@ pub unsafe fn vecwise_butterflies(
     val0.interleave_with(val1)
 }
 
-pub fn get_twiddle_dbls(domain: CircleDomain) -> Vec<Vec<i32>> {
-    let mut coset = domain.half_coset;
-
+/// Returns the line twiddles (x points) for an fft on a coset.
+pub fn get_twiddle_dbls(mut coset: Coset) -> Vec<Vec<i32>> {
     let mut res = vec![];
-    res.push(coset.iter().map(|p| (p.y.0 * 2) as i32).collect::<Vec<_>>());
-    bit_reverse(res.last_mut().unwrap());
     for _ in 0..coset.log_size() {
         res.push(
             coset
@@ -606,22 +603,22 @@ mod tests {
     #[test]
     fn test_vecwise_butterflies() {
         let domain = CanonicCoset::new(5).circle_domain();
-        let twiddle_dbls = get_twiddle_dbls(domain);
-        assert_eq!(twiddle_dbls.len(), 5);
+        let twiddle_dbls = get_twiddle_dbls(domain.half_coset);
+        assert_eq!(twiddle_dbls.len(), 4);
         let values0: [i32; 16] = std::array::from_fn(|i| i as i32);
         let values1: [i32; 16] = std::array::from_fn(|i| (i + 16) as i32);
         let result: [BaseField; 32] = unsafe {
             let (val0, val1) = avx_butterfly(
                 std::mem::transmute(values0),
                 std::mem::transmute(values1),
-                _mm512_set1_epi32(twiddle_dbls[4][0]),
+                _mm512_set1_epi32(twiddle_dbls[3][0]),
             );
             let (val0, val1) = vecwise_butterflies(
                 val0,
                 val1,
+                twiddle_dbls[0].clone().try_into().unwrap(),
                 twiddle_dbls[1].clone().try_into().unwrap(),
                 twiddle_dbls[2].clone().try_into().unwrap(),
-                twiddle_dbls[3].clone().try_into().unwrap(),
             );
             std::mem::transmute([val0, val1])
         };
@@ -648,12 +645,12 @@ mod tests {
 
             // Compute.
             let mut values = BaseFieldVec::from_iter(values);
-            let twiddle_dbls = get_twiddle_dbls(domain);
+            let twiddle_dbls = get_twiddle_dbls(domain.half_coset);
 
             unsafe {
                 fft_lower_with_vecwise(
                     std::mem::transmute(values.data.as_mut_ptr()),
-                    &twiddle_dbls[1..]
+                    &twiddle_dbls
                         .iter()
                         .map(|x| x.as_slice())
                         .collect::<Vec<_>>(),
@@ -676,7 +673,7 @@ mod tests {
 
         // Compute.
         let mut values = BaseFieldVec::from_iter(values);
-        let twiddle_dbls = get_twiddle_dbls(domain);
+        let twiddle_dbls = get_twiddle_dbls(domain.half_coset);
 
         unsafe {
             transpose_vecs(
@@ -685,7 +682,7 @@ mod tests {
             );
             fft(
                 std::mem::transmute(values.data.as_mut_ptr()),
-                &twiddle_dbls[1..]
+                &twiddle_dbls
                     .iter()
                     .map(|x| x.as_slice())
                     .collect::<Vec<_>>(),

src/core/backend/cpu/circle.rs

@@ -9,7 +9,7 @@ use crate::core::poly::circle::{
     CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly, PolyOps,
 };
 use crate::core::poly::twiddles::TwiddleTree;
-use crate::core::poly::utils::fold;
+use crate::core::poly::utils::{domain_line_twiddles_from_tree, fold};
 use crate::core::poly::BitReversedOrder;
 use crate::core::utils::bit_reverse;
 
@@ -156,20 +156,6 @@ impl PolyOps for CPUBackend {
     }
 }
 
-/// Computes the line twiddles for a [CircleDomain] from the precomputed twiddles tree.
-fn domain_line_twiddles_from_tree(
-    domain: CircleDomain,
-    twiddle_buffer: &[BaseField],
-) -> Vec<&[BaseField]> {
-    (0..domain.half_coset.log_size())
-        .map(|i| {
-            let len = 1 << i;
-            &twiddle_buffer[twiddle_buffer.len() - len * 2..twiddle_buffer.len() - len]
-        })
-        .rev()
-        .collect()
-}
-
 fn fft_layer_loop(
     values: &mut [BaseField],
     i: usize,

src/core/poly/utils.rs

@@ -1,3 +1,4 @@
+use super::circle::CircleDomain;
 use crate::core::fields::{ExtensionOf, Field};
 
 /// Folds values recursively in `O(n)` by a hierarchical application of folding factors.
@@ -53,6 +54,17 @@ pub fn repeat_value<T: Copy>(values: &[T], duplicity: usize) -> Vec<T> {
     res
 }
 
+/// Computes the line twiddles for a [CircleDomain] from the precomputed twiddles tree.
+pub fn domain_line_twiddles_from_tree<T>(domain: CircleDomain, twiddle_buffer: &[T]) -> Vec<&[T]> {
+    (0..domain.half_coset.log_size())
+        .map(|i| {
+            let len = 1 << i;
+            &twiddle_buffer[twiddle_buffer.len() - len * 2..twiddle_buffer.len() - len]
+        })
+        .rev()
+        .collect()
+}
+
 #[cfg(test)]
 mod tests {
     use super::repeat_value;