Vectorization for prime fields

Semi-honest AdditiveShare can hold a vector of sharings instead of just one sharing. The semi-honest multiply can operate on these vectors.
private-attribution · Jan 26, 2024 · d8cd1d1 · d8cd1d1
1 parent b54f8aa
commit d8cd1d1
Show file tree

Hide file tree

Showing 24 changed files with 1,148 additions and 179 deletions.
diff --git a/ipa-core/benches/ct/arithmetic_circuit.rs b/ipa-core/benches/ct/arithmetic_circuit.rs
@@ -1,8 +1,38 @@
 use criterion::{
-    black_box, criterion_group, criterion_main, BenchmarkId, Criterion, SamplingMode, Throughput,
+    black_box, criterion_group, criterion_main, measurement::Measurement, BenchmarkGroup,
+    BenchmarkId, Criterion, SamplingMode, Throughput,
 };
-use ipa_core::{ff::Fp31, test_fixture::circuit};
-use tokio::runtime::Builder;
+use ipa_core::{
+    ff::{Field, Fp31, Fp32BitPrime},
+    protocol::{basics::SecureMul, context::SemiHonestContext},
+    secret_sharing::{replicated::semi_honest::AdditiveShare as Replicated, FieldSimd, IntoShares},
+    test_fixture::circuit,
+};
+use rand::distributions::{Distribution, Standard};
+use tokio::runtime::{Builder, Runtime};
+
+fn do_benchmark<M, F, const N: usize>(
+    rt: &Runtime,
+    group: &mut BenchmarkGroup<M>,
+    width: u32,
+    depth: u16,
+) where
+    M: Measurement,
+    F: Field + FieldSimd<N>,
+    for<'a> Replicated<F, N>: SecureMul<SemiHonestContext<'a>>,
+    [F; N]: IntoShares<Replicated<F, N>>,
+    Standard: Distribution<F>,
+{
+    group.throughput(Throughput::Elements((width * depth as u32) as u64));
+    group.bench_with_input(
+        BenchmarkId::new("circuit", format!("{width}:{depth}:{}x{}", F::NAME, N)),
+        &(width, depth),
+        |b, &(width, depth)| {
+            b.to_async(rt)
+                .iter(|| circuit::arithmetic::<F, N>(black_box(width), black_box(depth)));
+        },
+    );
+}
 
 pub fn criterion_benchmark(c: &mut Criterion) {
     let rt = Builder::new_multi_thread()
@@ -16,19 +46,14 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     group.sample_size(10);
     group.sampling_mode(SamplingMode::Flat);
 
-    for width in [5_000u32, 50_000, 500_000, 1_000_000] {
-        for depth in [1u8, 10, 64] {
-            group.throughput(Throughput::Elements((width * depth as u32) as u64));
-            group.bench_with_input(
-                BenchmarkId::new("circuit", format!("{width}:{depth}")),
-                &(width, depth),
-                |b, &(width, depth)| {
-                    b.to_async(&rt)
-                        .iter(|| circuit::arithmetic::<Fp31>(black_box(width), black_box(depth)));
-                },
-            );
-        }
-    }
+    do_benchmark::<_, Fp31, 1>(&rt, &mut group, 512_000, 1);
+    do_benchmark::<_, Fp31, 1>(&rt, &mut group, 51_200, 10);
+    do_benchmark::<_, Fp31, 1>(&rt, &mut group, 8_000, 64);
+
+    do_benchmark::<_, Fp32BitPrime, 1>(&rt, &mut group, 25_600, 10);
+    do_benchmark::<_, Fp32BitPrime, 1>(&rt, &mut group, 2_560, 100);
+    do_benchmark::<_, Fp32BitPrime, 32>(&rt, &mut group, 4_000, 64);
+    do_benchmark::<_, Fp32BitPrime, 32>(&rt, &mut group, 250, 1_024);
 }
 
 criterion_group!(benches, criterion_benchmark);

diff --git a/ipa-core/benches/iai/arithmetic_circuit.rs b/ipa-core/benches/iai/arithmetic_circuit.rs
@@ -10,10 +10,10 @@ pub fn iai_benchmark() {
         .expect("Creating runtime failed");
 
     const CIRCUIT_WIDTH: u32 = 500_000;
-    const CIRCUIT_DEPTH: u8 = 1;
+    const CIRCUIT_DEPTH: u16 = 1;
 
     rt.block_on(async {
-        circuit::arithmetic::<Fp31>(black_box(CIRCUIT_WIDTH), black_box(CIRCUIT_DEPTH)).await;
+        circuit::arithmetic::<Fp31, 1>(black_box(CIRCUIT_WIDTH), black_box(CIRCUIT_DEPTH)).await;
     })
 }
 

diff --git a/ipa-core/benches/oneshot/arithmetic_circuit.rs b/ipa-core/benches/oneshot/arithmetic_circuit.rs
@@ -14,7 +14,7 @@ pub struct CircuitArgs {
     pub width: u32,
 
     #[arg(short, long, help = "depth of the circuit", default_value_t = 10)]
-    pub depth: u8,
+    pub depth: u16,
 
     /// Cargo passes the bench argument
     /// https://doc.rust-lang.org/cargo/commands/cargo-bench.html
@@ -34,7 +34,7 @@ pub async fn main() {
     }
 
     let start = Instant::now();
-    circuit::arithmetic::<Fp31>(args.width, args.depth).await;
+    circuit::arithmetic::<Fp31, 1>(args.width, args.depth).await;
     let duration = start.elapsed().as_secs_f32();
 
     println!("benchmark complete after {duration}s");

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 {
@@ -40,6 +43,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 ArrayAlias = StdArray<Boolean, 1>;
+}
+
 ///conversion to Scalar struct of `curve25519_dalek`
 impl From<Boolean> for bool {
     fn from(s: Boolean) -> Self {
@@ -146,6 +157,8 @@ impl From<bool> for Boolean {
 
 ///implement Field because required by PRSS
 impl Field for Boolean {
+    const NAME: &'static str = "Boolean";
+
     const ONE: Boolean = Boolean(true);
 
     fn as_u128(&self) -> u128 {

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
@@ -42,6 +42,12 @@ impl<'a> Iterator for BAIterator<'a> {
     }
 }
 
+impl<'a> ExactSizeIterator for BAIterator<'a> {
+    fn len(&self) -> usize {
+        self.iterator.len()
+    }
+}
+
 /// A value of ONE has a one in the first element of the bit array, followed by `$bits-1` zeros.
 /// This macro uses a bit of recursive repetition to produce those zeros.
 ///
@@ -95,6 +101,8 @@ macro_rules! boolean_array_impl_small {
 
         // TODO(812): remove this impl; BAs are not field elements.
         impl Field for $name {
+            const NAME: &'static str = stringify!($name);
+
             const ONE: Self = Self(bitarr_one!($bits));
 
             fn as_u128(&self) -> u128 {
@@ -153,6 +161,10 @@ macro_rules! boolean_array_impl_small {
                 Field::truncate_from(src)
             }
         }
+
+        impl FieldVectorizable<1> for $name {
+            type ArrayAlias = StdArray<$name, 1>;
+        }
     };
 }
 
@@ -358,6 +370,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 {
@@ -394,7 +410,7 @@ macro_rules! boolean_array_impl {
             #[allow(clippy::into_iter_without_iter)]
             impl<'a> IntoIterator for &'a AdditiveShare<$name> {
                 type Item = AdditiveShare<Boolean>;
-                type IntoIter = ASIterator<BAIterator<'a>>;
+                type IntoIter = ASIterator<'a, $name>;
 
                 fn into_iter(self) -> Self::IntoIter {
                     self.iter()

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>;
+}
+
 #[derive(thiserror::Error, Debug)]
 #[error("{0:?} is not the canonical encoding of a Ristretto point.")]
 pub struct NonCanonicalEncoding(CompressedRistretto);

diff --git a/ipa-core/src/ff/ec_prime_field.rs b/ipa-core/src/ff/ec_prime_field.rs
@@ -9,7 +9,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 {
@@ -176,8 +176,18 @@ 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 ArrayAlias = StdArray<Self, 1>;
+}
+
 ///implement Field because required by PRSS
 impl Field for Fp25519 {
+    const NAME: &'static str = "Fp25519";
+
     const ONE: Fp25519 = Fp25519::ONE;
 
     ///both following methods are based on hashing and do not allow to actually convert elements in Fp25519

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,7 +29,12 @@ pub trait Field:
     + FromRandomU128
     + TryFrom<u128, Error = error::Error>
     + Into<Self::Storage>
+    + Vectorizable<1>
+    + FieldVectorizable<1, ArrayAlias = <Self as Vectorizable<1>>::Array>
 {
+    // Name of the field
+    const NAME: &'static str;
+
     /// Multiplicative identity element
     const ONE: Self;
 

diff --git a/ipa-core/src/ff/galois_field.rs b/ipa-core/src/ff/galois_field.rs
@@ -15,7 +15,7 @@ use crate::{
     ff::{boolean_array::NonZeroPadding, Field, Serializable},
     impl_serializable_trait,
     protocol::prss::FromRandomU128,
-    secret_sharing::{Block, SharedValue},
+    secret_sharing::{Block, FieldVectorizable, SharedValue, Vectorizable},
 };
 
 /// Trait for data types storing arbitrary number of bits.
@@ -148,6 +148,12 @@ impl<'a> Iterator for BoolIterator<'a> {
     }
 }
 
+impl<'a> ExactSizeIterator for BoolIterator<'a> {
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+}
+
 macro_rules! bit_array_impl {
     ( $modname:ident, $name:ident, $store:ty, $bits:expr, $one:expr, $polynomial:expr, $deser_type: tt, $({$($extra:item)*})? ) => {
         #[allow(clippy::suspicious_arithmetic_impl)]
@@ -169,7 +175,17 @@ 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 ArrayAlias = crate::secret_sharing::StdArray<$name, 1>;
+            }
+
             impl Field for $name {
+                const NAME: &'static str = stringify!($field);
+
                 const ONE: Self = Self($one);
 
                 fn as_u128(&self) -> u128 {
@@ -693,5 +709,25 @@ bit_array_impl!(
                 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/mod.rs b/ipa-core/src/ff/mod.rs
@@ -82,7 +82,7 @@ pub trait Serializable: Sized {
 
 pub trait ArrayAccess {
     type Output;
-    type Iter<'a>: Iterator<Item = Self::Output> + Send
+    type Iter<'a>: Iterator<Item = Self::Output> + ExactSizeIterator + Send
     where
         Self: 'a;
 

diff --git a/ipa-core/src/ff/prime_field.rs b/ipa-core/src/ff/prime_field.rs
@@ -6,7 +6,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 {
@@ -33,7 +33,17 @@ 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 ArrayAlias = StdArray<$field, 1>;
+        }
+
         impl Field for $field {
+            const NAME: &'static str = stringify!($field);
+
             const ONE: Self = $field(1);
 
             fn as_u128(&self) -> u128 {
@@ -317,6 +327,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 ArrayAlias = StdArray<Fp32BitPrime, 32>;
+    }
+
     #[cfg(all(test, unit_test))]
     mod specialized_tests {
         use super::*;