From 80513ec67145db08b4198f9b5573fd39cb16e56d Mon Sep 17 00:00:00 2001
From: ohad-starkware <ohad.a@starkware.co>
Date: Mon, 13 Jan 2025 14:29:25 +0200
Subject: [PATCH] specific simd batch inverse functions

---
 crates/prover/src/core/backend/simd/cm31.rs | 7 ++++++-
 crates/prover/src/core/backend/simd/m31.rs  | 7 ++++++-
 crates/prover/src/core/backend/simd/qm31.rs | 7 ++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/crates/prover/src/core/backend/simd/cm31.rs b/crates/prover/src/core/backend/simd/cm31.rs
index 2155e8ff1..07398a613 100644
--- a/crates/prover/src/core/backend/simd/cm31.rs
+++ b/crates/prover/src/core/backend/simd/cm31.rs
@@ -6,7 +6,7 @@ use num_traits::{One, Zero};
 
 use super::m31::{PackedM31, N_LANES};
 use crate::core::fields::cm31::CM31;
-use crate::core::fields::FieldExpOps;
+use crate::core::fields::{batch_inverse_chunked, FieldExpOps};
 
 /// SIMD implementation of [`CM31`].
 #[derive(Copy, Clone, Debug)]
@@ -132,6 +132,11 @@ impl FieldExpOps for PackedCM31 {
         // 1 / (a + bi) = (a - bi) / (a^2 + b^2).
         Self([self.a(), -self.b()]) * (self.a().square() + self.b().square()).inverse()
     }
+
+    fn invert_many(column: &[Self]) -> Vec<Self> {
+        const OPTIMAL_PACKED_CM31_BATCH_INVERSE_CHUNK_SIZE: usize = 1 << 10;
+        batch_inverse_chunked(column, OPTIMAL_PACKED_CM31_BATCH_INVERSE_CHUNK_SIZE)
+    }
 }
 
 impl Add<PackedM31> for PackedCM31 {
diff --git a/crates/prover/src/core/backend/simd/m31.rs b/crates/prover/src/core/backend/simd/m31.rs
index 3d10be8c0..173d1005d 100644
--- a/crates/prover/src/core/backend/simd/m31.rs
+++ b/crates/prover/src/core/backend/simd/m31.rs
@@ -12,7 +12,7 @@ use rand::distributions::{Distribution, Standard};
 use super::qm31::PackedQM31;
 use crate::core::fields::m31::{pow2147483645, BaseField, M31, P};
 use crate::core::fields::qm31::QM31;
-use crate::core::fields::FieldExpOps;
+use crate::core::fields::{batch_inverse_chunked, FieldExpOps};
 
 pub const LOG_N_LANES: u32 = 4;
 
@@ -251,6 +251,11 @@ impl FieldExpOps for PackedM31 {
         assert!(!self.is_zero(), "0 has no inverse");
         pow2147483645(*self)
     }
+
+    fn invert_many(column: &[Self]) -> Vec<Self> {
+        const OPTIMAL_PACKED_M31_BATCH_INVERSE_CHUNK_SIZE: usize = 1 << 9;
+        batch_inverse_chunked(column, OPTIMAL_PACKED_M31_BATCH_INVERSE_CHUNK_SIZE)
+    }
 }
 
 unsafe impl Pod for PackedM31 {}
diff --git a/crates/prover/src/core/backend/simd/qm31.rs b/crates/prover/src/core/backend/simd/qm31.rs
index ce7231d0a..077837df7 100644
--- a/crates/prover/src/core/backend/simd/qm31.rs
+++ b/crates/prover/src/core/backend/simd/qm31.rs
@@ -10,7 +10,7 @@ use super::cm31::PackedCM31;
 use super::m31::{PackedM31, N_LANES};
 use crate::core::fields::m31::M31;
 use crate::core::fields::qm31::QM31;
-use crate::core::fields::FieldExpOps;
+use crate::core::fields::{batch_inverse_chunked, FieldExpOps};
 
 pub type PackedSecureField = PackedQM31;
 
@@ -171,6 +171,11 @@ impl FieldExpOps for PackedQM31 {
         let denom_inverse = denom.inverse();
         Self([self.a() * denom_inverse, -self.b() * denom_inverse])
     }
+
+    fn invert_many(column: &[Self]) -> Vec<Self> {
+        const OPTIMAL_PACKED_QM31_BATCH_INVERSE_CHUNK_SIZE: usize = 1 << 11;
+        batch_inverse_chunked(column, OPTIMAL_PACKED_QM31_BATCH_INVERSE_CHUNK_SIZE)
+    }
 }
 
 impl Add<PackedM31> for PackedQM31 {