From 2e2ddba6c4f2c4da90338217a0acf634d0b0cc7e Mon Sep 17 00:00:00 2001
From: ohad-starkware <ohad.a@starkware.co>
Date: Mon, 13 Jan 2025 10:50:59 +0200
Subject: [PATCH] chunked batch inverse

---
 crates/prover/src/core/fields/mod.rs | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/crates/prover/src/core/fields/mod.rs b/crates/prover/src/core/fields/mod.rs
index df95525d8..b33a53ee1 100644
--- a/crates/prover/src/core/fields/mod.rs
+++ b/crates/prover/src/core/fields/mod.rs
@@ -95,13 +95,22 @@ pub fn batch_inverse_in_place<F: FieldExpOps>(column: &[F], dst: &mut [F]) {
     dst[0..WIDTH].clone_from_slice(&tail_inverses);
 }
 
-// TODO(Ohad): chunks, parallelize.
 pub fn batch_inverse<F: FieldExpOps>(column: &[F]) -> Vec<F> {
     let mut dst = vec![unsafe { std::mem::zeroed() }; column.len()];
     batch_inverse_in_place(column, &mut dst);
     dst
 }
 
+// TODO(Ohad): parallelize.
+pub fn batch_inverse_chunked<T: FieldExpOps>(column: &[T], chunk_size: usize) -> Vec<T> {
+    let mut dst = vec![unsafe { std::mem::zeroed() }; column.len()];
+    let iter = dst.chunks_mut(chunk_size).zip(column.chunks(chunk_size));
+    iter.for_each(|(dst, column)| {
+        batch_inverse_in_place(column, dst);
+    });
+    dst
+}
+
 pub trait Field:
     NumAssign
     + Neg<Output = Self>
@@ -472,8 +481,8 @@ mod tests {
     use rand::{Rng, SeedableRng};
 
     use super::batch_inverse_in_place;
-    use crate::core::fields::batch_inverse;
     use crate::core::fields::m31::M31;
+    use crate::core::fields::{batch_inverse, batch_inverse_chunked};
 
     #[test]
     fn test_batch_inverse() {
@@ -495,4 +504,16 @@ mod tests {
 
         batch_inverse_in_place(&elements, &mut dst);
     }
+
+    #[test]
+    fn test_batch_inverse_chunked() {
+        let mut rng = SmallRng::seed_from_u64(0);
+        let elements: [M31; 16] = rng.gen();
+        let chunk_size = 4;
+        let expected = batch_inverse(&elements);
+
+        let result = batch_inverse_chunked(&elements, chunk_size);
+
+        assert_eq!(expected, result);
+    }
 }