Karatsuba algorithm for multiplication

crates/prover/benches/field.rs

@@ -4,6 +4,7 @@ use rand::{Rng, SeedableRng};
 use stwo_prover::core::fields::cm31::CM31;
 use stwo_prover::core::fields::m31::M31;
 use stwo_prover::core::fields::qm31::SecureField;
+use stwo_prover::core::fields::FieldExpOps;
 
 pub const N_ELEMENTS: usize = 1 << 16;
 pub const N_STATE_ELEMENTS: usize = 8;
@@ -55,6 +56,18 @@ pub fn cm31_operations_bench(c: &mut criterion::Criterion) {
         })
     });
 
+    c.bench_function("CM31 square", |b| {
+        b.iter(|| {
+            for _ in &elements {
+                for _ in 0..128 {
+                    for state_elem in &mut state {
+                        *state_elem = state_elem.square();
+                    }
+                }
+            }
+        })
+    });
+
     c.bench_function("CM31 add", |b| {
         b.iter(|| {
             for elem in &elements {

crates/prover/src/core/fields/cm31.rs

@@ -52,7 +52,13 @@ impl Mul for CM31 {
     }
 }
 
-impl FieldExpOps for CM31 {
+impl FieldExpOps for CM31 {    
+    fn square(&self) -> Self {
+        // Complex squaring algorithm taken from https://eprint.iacr.org/2006/471.pdf
+        let a0a1 = self.0 * self.1;
+        Self((self.0 + self.1) * (self.0 - self.1), a0a1 + a0a1)    
+    }
+
     fn inverse(&self) -> Self {
         assert!(!self.is_zero(), "0 has no inverse");
         // 1 / (a + bi) = (a - bi) / (a^2 + b^2).

crates/prover/src/core/fields/qm31.rs

@@ -58,11 +58,11 @@ impl Mul for QM31 {
     type Output = Self;
 
     fn mul(self, rhs: Self) -> Self::Output {
-        // (a + bu) * (c + du) = (ac + rbd) + (ad + bc)u.
-        Self(
-            self.0 * rhs.0 + R * self.1 * rhs.1,
-            self.0 * rhs.1 + self.1 * rhs.0,
-        )
+        // Karatsuba algorithm taken from https://eprint.iacr.org/2006/471.pdf
+        let v0 = self.0 * rhs.0;
+        let v1 = self.1 * rhs.1;
+
+        Self(v0 + R * v1, (self.0 + self.1) * (rhs.0 + rhs.1) - v0 - v1)
     }
 }