Merge pull request #852 from andyleiserson/vector-mult-4

Vectorization for prime fields

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);

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;
     })
 }
 

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");

ipa-core/src/error.rs

@@ -1,4 +1,8 @@
-use std::{backtrace::Backtrace, convert::Infallible, fmt::Debug};
+use std::{
+    backtrace::Backtrace,
+    convert::Infallible,
+    fmt::{Debug, Display},
+};
 
 use thiserror::Error;
 
@@ -60,6 +64,8 @@ pub enum Error {
     Unsupported(String),
     #[error("Decompressing invalid elliptic curve point: {0}")]
     DecompressingInvalidCurvePoint(String),
+    #[error(transparent)]
+    LengthError(#[from] LengthError),
 }
 
 impl Default for Error {
@@ -85,6 +91,20 @@ pub type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
 
 pub type Res<T> = Result<T, Error>;
 
+#[derive(Debug)]
+pub struct LengthError {
+    pub expected: usize,
+    pub actual: usize,
+}
+
+impl Display for LengthError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "expected {} items, got {}", self.expected, self.actual)
+    }
+}
+
+impl std::error::Error for LengthError {}
+
 /// Set up a global panic hook that dumps the panic information to our tracing subsystem if it is
 /// available and duplicates that to standard error output.
 ///

ipa-core/src/ff/boolean.rs

@@ -4,8 +4,12 @@ use typenum::U1;
 use super::Gf32Bit;
 use crate::{
     ff::{Field, Serializable},
+    impl_shared_value_common,
     protocol::prss::FromRandomU128,
-    secret_sharing::{replicated::malicious::ExtendableField, Block, SharedValue},
+    secret_sharing::{
+        replicated::malicious::ExtendableField, Block, FieldVectorizable, SharedValue, StdArray,
+        Vectorizable,
+    },
 };
 
 impl Block for bool {
@@ -38,6 +42,16 @@ impl SharedValue for Boolean {
     type Storage = bool;
     const BITS: u32 = 1;
     const ZERO: Self = Self(false);
+
+    impl_shared_value_common!();
+}
+
+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`
@@ -146,6 +160,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 {

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>;
+        }
     };
 }
 
@@ -260,6 +272,7 @@ macro_rules! boolean_array_impl {
             use super::*;
             use crate::{
                 ff::{boolean::Boolean, ArrayAccess, Expand, Serializable},
+                impl_shared_value_common,
                 secret_sharing::{
                     replicated::semi_honest::{ASIterator, AdditiveShare},
                     SharedValue,
@@ -272,6 +285,11 @@ macro_rules! boolean_array_impl {
             #[derive(Clone, Copy, PartialEq, Eq, Debug)]
             pub struct $name(pub(super) Store);
 
+            impl $name {
+                #[cfg(all(test, unit_test))]
+                const STORE_LEN: usize = bitvec::mem::elts::<u8>($bits);
+            }
+
             impl ArrayAccess for $name {
                 type Output = Boolean;
                 type Iter<'a> = BAIterator<'a>;
@@ -300,6 +318,8 @@ macro_rules! boolean_array_impl {
                 type Storage = Store;
                 const BITS: u32 = $bits;
                 const ZERO: Self = Self(<Store>::ZERO);
+
+                impl_shared_value_common!();
             }
 
             impl_serializable_trait!($name, $bits, Store, $deser_type);
@@ -358,6 +378,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 +418,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()
@@ -411,12 +435,26 @@ macro_rules! boolean_array_impl {
 
             #[cfg(all(test, unit_test))]
             mod tests {
+                use proptest::{
+                    prelude::{prop, Arbitrary, Strategy},
+                    proptest,
+                };
                 use rand::{thread_rng, Rng};
 
                 use super::*;
 
-                // Only small BAs expose this via `Field`.
-                const ONE: $name = $name(bitarr_one!($bits));
+                impl Arbitrary for $name {
+                    type Parameters = <[u8; $name::STORE_LEN] as Arbitrary>::Parameters;
+                    type Strategy = prop::strategy::Map<
+                        <[u8; $name::STORE_LEN] as Arbitrary>::Strategy,
+                        fn([u8; $name::STORE_LEN]) -> Self,
+                    >;
+
+                    fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
+                        <[u8; $name::STORE_LEN]>::arbitrary_with(args)
+                            .prop_map(|arr| $name(Store::from(arr)))
+                    }
+                }
 
                 #[test]
                 fn set_boolean_array() {
@@ -428,29 +466,27 @@ macro_rules! boolean_array_impl {
                     assert_eq!(ba.get(i), Some(a));
                 }
 
-                #[test]
-                fn iterate_boolean_array() {
-                    let bits = ONE;
-                    let iter = bits.iter();
-                    for (i, j) in iter.enumerate() {
-                        if i == 0 {
-                            assert_eq!(j, Boolean::ONE);
-                        } else {
-                            assert_eq!(j, Boolean::ZERO);
+                proptest! {
+                    #[test]
+                    fn iterate_boolean_array(a: $name) {
+                        let mut iter = a.iter().enumerate();
+                        assert_eq!(iter.len(), $bits);
+                        while let Some((i, b)) = iter.next() {
+                            assert_eq!(bool::from(b), a.0[i]);
+                            assert_eq!(iter.len(), $bits - 1 - i);
                         }
                     }
-                }
 
-                #[test]
-                fn iterate_secret_shared_boolean_array() {
+                    #[test]
+                    fn iterate_secret_shared_boolean_array(a: AdditiveShare<$name>) {
                     use crate::secret_sharing::replicated::ReplicatedSecretSharing;
-                    let bits = AdditiveShare::new(ONE, ONE);
-                    let iter = bits.into_iter();
-                    for (i, j) in iter.enumerate() {
-                        if i == 0 {
-                            assert_eq!(j, AdditiveShare::new(Boolean::ONE, Boolean::ONE));
-                        } else {
-                            assert_eq!(j, AdditiveShare::<Boolean>::ZERO);
+                        let mut iter = a.iter().enumerate();
+                        assert_eq!(iter.len(), $bits);
+                        while let Some((i, sb)) = iter.next() {
+                            let left = Boolean::from(a.left().0[i]);
+                            let right = Boolean::from(a.right().0[i]);
+                            assert_eq!(sb, AdditiveShare::new(left, right));
+                            assert_eq!(iter.len(), $bits - 1 - i);
                         }
                     }
                 }

ipa-core/src/ff/curve_points.rs

@@ -7,7 +7,8 @@ use typenum::U32;
 
 use crate::{
     ff::{ec_prime_field::Fp25519, Serializable},
-    secret_sharing::{Block, SharedValue},
+    impl_shared_value_common,
+    secret_sharing::{Block, SharedValue, StdArray, Vectorizable},
 };
 
 impl Block for CompressedRistretto {
@@ -33,6 +34,12 @@ impl SharedValue for RP25519 {
     type Storage = CompressedRistretto;
     const BITS: u32 = 256;
     const ZERO: Self = Self(CompressedRistretto([0_u8; 32]));
+
+    impl_shared_value_common!();
+}
+
+impl Vectorizable<1> for RP25519 {
+    type Array = StdArray<Self, 1>;
 }
 
 #[derive(thiserror::Error, Debug)]