From 20a76bf117688c79648b48a4263c2a873d4b428d Mon Sep 17 00:00:00 2001
From: Ohad <137686240+ohad-starkware@users.noreply.github.com>
Date: Mon, 18 Mar 2024 09:12:36 +0200
Subject: [PATCH] integrated batch inverse to cpu twiddles (#469)

---
 src/core/backend/cpu/circle.rs | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/core/backend/cpu/circle.rs b/src/core/backend/cpu/circle.rs
index 9d6f9360d..5bc11eb77 100644
--- a/src/core/backend/cpu/circle.rs
+++ b/src/core/backend/cpu/circle.rs
@@ -130,6 +130,9 @@ impl PolyOps for CPUBackend {
     }
 
     fn precompute_twiddles(mut coset: Coset) -> TwiddleTree<Self> {
+        const CHUNK_LOG_SIZE: usize = 12;
+        const CHUNK_SIZE: usize = 1 << CHUNK_LOG_SIZE;
+
         let root_coset = coset;
         let mut twiddles = Vec::with_capacity(coset.size());
         for _ in 0..coset.log_size() {
@@ -146,8 +149,24 @@ impl PolyOps for CPUBackend {
         }
         twiddles.push(1.into());
 
-        // TODO(spapini): Batch inverse.
-        let itwiddles = twiddles.iter().map(|&t| t.inverse()).collect();
+        // Inverse twiddles.
+        // Fallback to the non-chunked version if the domain is not big enough.
+        if CHUNK_SIZE > coset.size() {
+            let itwiddles = twiddles.iter().map(|&t| t.inverse()).collect();
+            return TwiddleTree {
+                root_coset,
+                twiddles,
+                itwiddles,
+            };
+        }
+
+        let mut itwiddles = vec![BaseField::zero(); twiddles.len()];
+        twiddles
+            .array_chunks::<CHUNK_SIZE>()
+            .zip(itwiddles.array_chunks_mut::<CHUNK_SIZE>())
+            .for_each(|(src, dst)| {
+                BaseField::batch_inverse(src, dst);
+            });
 
         TwiddleTree {
             root_coset,