Merge pull request #1458 from andyleiserson/dzkp-bench

End-to-end DZKP benchmark

ipa-core/Cargo.toml

@@ -250,7 +250,8 @@ harness = false
 required-features = ["enable-benches"]
 
 [[bench]]
-name = "dzkp_convert_prover"
+name = "dzkp"
+path = "benches/ct/dzkp.rs"
 harness = false
 required-features = ["enable-benches"]
 

ipa-core/benches/ct/dzkp.rs

@@ -0,0 +1,118 @@
+//! Benchmarks for DZKPs.
+
+use std::iter::{repeat_with, zip};
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, SamplingMode};
+use futures::{stream::iter, TryStreamExt};
+use ipa_core::{
+    ff::boolean_array::BA256,
+    helpers::TotalRecords,
+    protocol::{
+        basics::BooleanArrayMul,
+        context::{
+            dzkp_validator::{DZKPValidator, MultiplicationInputsBlock, TARGET_PROOF_SIZE},
+            malicious::TEST_DZKP_STEPS,
+            Context, DZKPUpgradedMaliciousContext, UpgradableContext,
+        },
+        RecordId,
+    },
+    secret_sharing::{replicated::semi_honest::AdditiveShare as Replicated, SharedValue},
+    sharding::NotSharded,
+    test_fixture::{Runner, TestWorld},
+    utils::non_zero_prev_power_of_two,
+};
+use rand::{thread_rng, Rng};
+use tokio::runtime::Builder;
+
+/// Benchmark for the table_indices_prover function in dzkp_field.rs.
+fn benchmark_table_indices_prover(c: &mut Criterion) {
+    let mut group = c.benchmark_group("benches");
+    group.bench_function("table_indices_prover", |b| {
+        b.iter_batched_ref(
+            || thread_rng().gen(),
+            |input: &mut MultiplicationInputsBlock| input.table_indices_prover(),
+            BatchSize::SmallInput,
+        )
+    });
+    group.finish();
+}
+
+/// Benchmark for end-to-end proof.
+///
+/// This benchmark focuses on proof performance by evaluating one of the simplest and
+/// most performant MPC circuits possible: 64 million AND gates in parallel.
+fn benchmark_proof(c: &mut Criterion) {
+    let rt = Builder::new_multi_thread()
+        .worker_threads(3)
+        .thread_name("helper-worker")
+        .enable_time()
+        .build()
+        .expect("Creating runtime failed");
+
+    type BA = BA256;
+    const COUNT: usize = 64 * 1024 * 1024 / BA::BITS as usize;
+
+    let mut group = c.benchmark_group("proof");
+    group.sample_size(10);
+    group.sampling_mode(SamplingMode::Flat);
+    group.bench_function("proof", |b| {
+        b.to_async(&rt).iter_batched(
+            || {
+                let mut rng = thread_rng();
+
+                let a = repeat_with(|| rng.gen()).take(COUNT).collect::<Vec<BA>>();
+                let b = repeat_with(|| rng.gen()).take(COUNT).collect::<Vec<BA>>();
+
+                (a, b)
+            },
+            |(a, b): (Vec<BA>, Vec<BA>)| async move {
+                TestWorld::default()
+                    .malicious((a.into_iter(), b.into_iter()), |ctx, (a, b)| async move {
+                        let batch_size = non_zero_prev_power_of_two(
+                            TARGET_PROOF_SIZE / usize::try_from(BA::BITS).unwrap(),
+                        );
+                        let v = ctx
+                            .set_total_records(TotalRecords::specified(COUNT)?)
+                            .dzkp_validator(TEST_DZKP_STEPS, batch_size);
+                        let m_ctx = v.context();
+
+                        v.validated_seq_join(iter(zip(a, b).enumerate().map(
+                            |(i, (a_malicious, b_malicious))| {
+                                let m_ctx = m_ctx.clone();
+                                let a_vec = <Replicated<BA> as BooleanArrayMul<
+                                    DZKPUpgradedMaliciousContext<NotSharded>,
+                                >>::Vectorized::from(
+                                    a_malicious
+                                );
+                                let b_vec = <Replicated<BA> as BooleanArrayMul<
+                                    DZKPUpgradedMaliciousContext<NotSharded>,
+                                >>::Vectorized::from(
+                                    b_malicious
+                                );
+                                async move {
+                                    <Replicated<BA> as BooleanArrayMul<_>>::multiply(
+                                        m_ctx,
+                                        RecordId::from(i),
+                                        &a_vec,
+                                        &b_vec,
+                                    )
+                                    .await
+                                    .map(<Replicated<BA>>::from)
+                                }
+                            },
+                        )))
+                        .try_collect::<Vec<_>>()
+                        .await
+                    })
+                    .await
+                    .map(Result::unwrap);
+            },
+            BatchSize::PerIteration,
+        )
+    });
+    group.finish();
+}
+
+criterion_group!(benches, benchmark_table_indices_prover);
+criterion_group!(proof, benchmark_proof);
+criterion_main!(benches, proof);

ipa-core/src/lib.rs

@@ -31,7 +31,7 @@ mod app;
 mod seq_join;
 mod serde;
 pub mod sharding;
-mod utils;
+pub mod utils;
 pub use app::{AppConfig, HelperApp, Setup as AppSetup};
 pub use utils::NonZeroU32PowerOfTwo;
 

ipa-core/src/protocol/context/malicious.rs

@@ -43,13 +43,11 @@ pub struct MaliciousProtocolSteps<'a, S: Step + ?Sized> {
 }
 
 #[cfg(all(feature = "in-memory-infra", any(test, feature = "test-fixture")))]
-pub(crate) const TEST_DZKP_STEPS: MaliciousProtocolSteps<
-    'static,
-    super::step::MaliciousProtocolStep,
-> = MaliciousProtocolSteps {
-    protocol: &super::step::MaliciousProtocolStep::MaliciousProtocol,
-    validate: &super::step::MaliciousProtocolStep::Validate,
-};
+pub const TEST_DZKP_STEPS: MaliciousProtocolSteps<'static, super::step::MaliciousProtocolStep> =
+    MaliciousProtocolSteps {
+        protocol: &super::step::MaliciousProtocolStep::MaliciousProtocol,
+        validate: &super::step::MaliciousProtocolStep::Validate,
+    };
 
 #[derive(Clone)]
 pub struct Context<'a, B: ShardBinding> {

ipa-core/src/protocol/context/step.rs

@@ -8,7 +8,7 @@ pub(crate) struct UpgradeStep;
 /// Steps used by the validation component of malicious protocol execution.
 /// In addition to these, an implicit step is used to initialize the value of `r`.
 #[derive(CompactStep)]
-pub(crate) enum MaliciousProtocolStep {
+pub enum MaliciousProtocolStep {
     /// For the execution of the malicious protocol.
     #[step(child = crate::protocol::ipa_prf::step::PrfStep)]
     MaliciousProtocol,

ipa-core/src/utils/power_of_two.rs

@@ -71,6 +71,7 @@ impl NonZeroU32PowerOfTwo {
 /// Returns the largest power of two less than or equal to `target`.
 ///
 /// Returns 1 if `target` is zero.
+#[must_use]
 pub fn non_zero_prev_power_of_two(target: usize) -> usize {
     let bits = usize::BITS - target.leading_zeros();