specific simd batch inverse functions

starkware-libs · Jan 13, 2025 · eeeccf7 · eeeccf7
1 parent e304b42
commit eeeccf7
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git 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_M31_BATCH_INVERSE_CHUNK_SIZE: usize = 1 << 10;
+        batch_inverse_chunked(column, OPTIMAL_PACKED_M31_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
@@ -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
@@ -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 {