Vectorization

andyleiserson · Dec 21, 2023 · 4f9957c · 4f9957c
1 parent c9fd7bf
commit 4f9957c
Show file tree

Hide file tree

Showing 20 changed files with 1,049 additions and 174 deletions.
diff --git a/ipa-core/src/ff/boolean.rs b/ipa-core/src/ff/boolean.rs
@@ -5,7 +5,10 @@ use super::Gf32Bit;
 use crate::{
     ff::{Field, Serializable},
     protocol::prss::FromRandomU128,
-    secret_sharing::{replicated::malicious::ExtendableField, Block, SharedValue},
+    secret_sharing::{
+        replicated::malicious::ExtendableField, Block, FieldVectorizable, SharedValue, StdArray,
+        Vectorizable,
+    },
 };
 
 impl Block for bool {
@@ -30,6 +33,14 @@ impl SharedValue for Boolean {
     const ZERO: Self = Self(false);
 }
 
+impl Vectorizable<1> for Boolean {
+    type Array = StdArray<Boolean, 1>;
+}
+
+impl FieldVectorizable<1> for Boolean {
+    type T = StdArray<Boolean, 1>;
+}
+
 ///conversion to Scalar struct of `curve25519_dalek`
 impl From<Boolean> for bool {
     fn from(s: Boolean) -> Self {

diff --git a/ipa-core/src/ff/boolean_array.rs b/ipa-core/src/ff/boolean_array.rs
@@ -8,7 +8,7 @@ use typenum::{U14, U2, U32, U8};
 use crate::{
     ff::{boolean::Boolean, ArrayAccess, Field, Serializable},
     protocol::prss::{FromRandom, FromRandomU128},
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, FieldVectorizable, SharedValue, StdArray, Vectorizable},
 };
 
 /// The implementation below cannot be constrained without breaking Rust's
@@ -152,6 +152,10 @@ macro_rules! boolean_array_impl_small {
                 Field::truncate_from(src)
             }
         }
+
+        impl FieldVectorizable<1> for $name {
+            type T = StdArray<$name, 1>;
+        }
     };
 }
 
@@ -272,6 +276,10 @@ macro_rules! boolean_array_impl {
                 }
             }
 
+            impl Vectorizable<1> for $name {
+                type Array = StdArray<$name, 1>;
+            }
+
             impl std::ops::Mul for $name {
                 type Output = Self;
                 fn mul(self, rhs: Self) -> Self::Output {

diff --git a/ipa-core/src/ff/curve_points.rs b/ipa-core/src/ff/curve_points.rs
@@ -7,7 +7,7 @@ use typenum::U32;
 
 use crate::{
     ff::{ec_prime_field::Fp25519, Serializable},
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, SharedValue, StdArray, Vectorizable},
 };
 
 impl Block for CompressedRistretto {
@@ -35,6 +35,10 @@ impl SharedValue for RP25519 {
     const ZERO: Self = Self(CompressedRistretto([0_u8; 32]));
 }
 
+impl Vectorizable<1> for RP25519 {
+    type Array = StdArray<Self, 1>;
+}
+
 impl Serializable for RP25519 {
     type Size = <<RP25519 as SharedValue>::Storage as Block>::Size;
 

diff --git a/ipa-core/src/ff/ec_prime_field.rs b/ipa-core/src/ff/ec_prime_field.rs
@@ -7,7 +7,7 @@ use typenum::U32;
 use crate::{
     ff::{boolean_array::BA256, Field, Serializable},
     protocol::prss::FromRandomU128,
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, FieldVectorizable, SharedValue, StdArray, Vectorizable},
 };
 
 impl Block for Scalar {
@@ -172,6 +172,14 @@ macro_rules! sc_hash_impl {
 #[cfg(test)]
 sc_hash_impl!(u64);
 
+impl Vectorizable<1> for Fp25519 {
+    type Array = StdArray<Self, 1>;
+}
+
+impl FieldVectorizable<1> for Fp25519 {
+    type T = StdArray<Self, 1>;
+}
+
 ///implement Field because required by PRSS
 impl Field for Fp25519 {
     const ONE: Fp25519 = Fp25519::ONE;

diff --git a/ipa-core/src/ff/field.rs b/ipa-core/src/ff/field.rs
@@ -8,7 +8,7 @@ use typenum::{U1, U4};
 use crate::{
     error,
     protocol::prss::FromRandomU128,
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, FieldVectorizable, SharedValue, Vectorizable},
 };
 
 impl Block for u8 {
@@ -29,6 +29,8 @@ pub trait Field:
     + FromRandomU128
     + TryFrom<u128, Error = error::Error>
     + Into<Self::Storage>
+    + Vectorizable<1>
+    + FieldVectorizable<1, T = <Self as Vectorizable<1>>::Array>
 {
     /// Multiplicative identity element
     const ONE: Self;

diff --git a/ipa-core/src/ff/galois_field.rs b/ipa-core/src/ff/galois_field.rs
@@ -14,7 +14,7 @@ use super::ArrayAccess;
 use crate::{
     ff::{Field, Serializable},
     protocol::prss::FromRandomU128,
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, FieldVectorizable, SharedValue, Vectorizable},
 };
 
 /// Trait for data types storing arbitrary number of bits.
@@ -168,6 +168,14 @@ macro_rules! bit_array_impl {
                 const ZERO: Self = Self(<$store>::ZERO);
             }
 
+            impl Vectorizable<1> for $name {
+                type Array = crate::secret_sharing::StdArray<$name, 1>;
+            }
+
+            impl FieldVectorizable<1> for $name {
+                type T = crate::secret_sharing::StdArray<$name, 1>;
+            }
+
             impl Field for $name {
                 const ONE: Self = Self($one);
 
@@ -685,5 +693,31 @@ bit_array_impl!(
                 v
             }
         }
+
+        impl From<Gf2> for bool {
+            fn from(value: Gf2) -> Self {
+                value != Gf2::ZERO
+            }
+        }
+
+        impl From<crate::ff::boolean::Boolean> for Gf2 {
+            fn from(value: crate::ff::boolean::Boolean) -> Self {
+                bool::from(value).into()
+            }
+        }
+
+        impl From<Gf2> for crate::ff::boolean::Boolean {
+            fn from(value: Gf2) -> Self {
+                bool::from(value).into()
+            }
+        }
+
+        impl std::ops::Not for Gf2 {
+            type Output = Self;
+
+            fn not(self) -> Self {
+                (!bool::from(self)).into()
+            }
+        }
     }
 );
diff --git a/ipa-core/src/ff/prime_field.rs b/ipa-core/src/ff/prime_field.rs
@@ -4,7 +4,7 @@ use super::Field;
 use crate::{
     ff::Serializable,
     protocol::prss::FromRandomU128,
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, FieldVectorizable, SharedValue, StdArray, Vectorizable},
 };
 
 pub trait PrimeField: Field {
@@ -43,6 +43,14 @@ macro_rules! field_impl {
             const ZERO: Self = $field(0);
         }
 
+        impl Vectorizable<1> for $field {
+            type Array = StdArray<$field, 1>;
+        }
+
+        impl FieldVectorizable<1> for $field {
+            type T = StdArray<$field, 1>;
+        }
+
         impl Field for $field {
             const ONE: Self = $field(1);
 
@@ -270,6 +278,14 @@ mod fp31 {
 mod fp32bit {
     field_impl! { Fp32BitPrime, u32, 32, 4_294_967_291 }
 
+    impl Vectorizable<32> for Fp32BitPrime {
+        type Array = StdArray<Fp32BitPrime, 32>;
+    }
+
+    impl FieldVectorizable<32> for Fp32BitPrime {
+        type T = StdArray<Fp32BitPrime, 32>;
+    }
+
     #[cfg(all(test, unit_test))]
     mod specialized_tests {
         use super::*;

diff --git a/ipa-core/src/protocol/basics/mul/mod.rs b/ipa-core/src/protocol/basics/mul/mod.rs
@@ -7,9 +7,12 @@ use crate::{
         context::{Context, UpgradedMaliciousContext},
         RecordId,
     },
-    secret_sharing::replicated::{
-        malicious::{AdditiveShare as MaliciousReplicated, ExtendableField},
-        semi_honest::AdditiveShare as Replicated,
+    secret_sharing::{
+        replicated::{
+            malicious::{AdditiveShare as MaliciousReplicated, ExtendableField},
+            semi_honest::AdditiveShare as Replicated,
+        },
+        FieldSimd,
     },
 };
 
@@ -52,7 +55,11 @@ use {malicious::multiply as malicious_mul, semi_honest::multiply as semi_honest_
 
 /// Implement secure multiplication for semi-honest contexts with replicated secret sharing.
 #[async_trait]
-impl<C: Context, F: Field> SecureMul<C> for Replicated<F> {
+impl<C, F, const N: usize> SecureMul<C> for Replicated<F, N>
+where
+    C: Context,
+    F: Field + FieldSimd<N>,
+{
     async fn multiply_sparse<'fut>(
         &self,
         rhs: &Self,