From a0b349610ef682cb8de1f45c9e9b4148a8e648ec Mon Sep 17 00:00:00 2001
From: Alex Koshelev <koshelev@meta.com>
Date: Fri, 5 Jul 2024 12:25:29 -0700
Subject: [PATCH] Add PRSS criterion benchmark

---
 ipa-core/Cargo.toml         |  6 ++++
 ipa-core/benches/ct/prss.rs | 68 +++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 ipa-core/benches/ct/prss.rs

diff --git a/ipa-core/Cargo.toml b/ipa-core/Cargo.toml
index 06fc1b262..cd7ed5a41 100644
--- a/ipa-core/Cargo.toml
+++ b/ipa-core/Cargo.toml
@@ -199,6 +199,12 @@ path = "benches/ct/arithmetic_circuit.rs"
 harness = false
 required-features = ["enable-benches"]
 
+[[bench]]
+name = "criterion_prss"
+path = "benches/ct/prss.rs"
+harness = false
+required-features = ["enable-benches"]
+
 [[bench]]
 name = "iai_arithmetic"
 path = "benches/iai/arithmetic_circuit.rs"
diff --git a/ipa-core/benches/ct/prss.rs b/ipa-core/benches/ct/prss.rs
new file mode 100644
index 000000000..4e14a485d
--- /dev/null
+++ b/ipa-core/benches/ct/prss.rs
@@ -0,0 +1,68 @@
+use std::sync::Arc;
+
+use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use ipa_core::{
+    ff::boolean_array::{BA256, BA64, BA8},
+    protocol::{
+        prss::{Endpoint, IndexedSharedRandomness, KeyExchange, SharedRandomness},
+        Gate,
+    },
+};
+use rand::thread_rng;
+
+fn make_prss() -> Arc<IndexedSharedRandomness> {
+    let prss_setup = Endpoint::prepare(&mut thread_rng());
+    let left_peer = KeyExchange::new(&mut thread_rng());
+    let right_peer = KeyExchange::new(&mut thread_rng());
+
+    let prss = prss_setup
+        .setup(&left_peer.public_key(), &right_peer.public_key())
+        .indexed(&Gate::default());
+
+    prss
+}
+
+fn prss_benchmark(c: &mut Criterion) {
+    if cfg!(debug_assertions) {
+        panic!(
+            "Debug assertions enable PRSS generators to track used indices. This \
+                will make this benchmark very slow. Turn them off."
+        )
+    }
+
+    // Setup PRSS outside measured code.
+    let mut group = c.benchmark_group("prss_one_chunk");
+    let prss = make_prss();
+
+    // PRSS generates a pair of 16 byte values
+    group.throughput(Throughput::Bytes(32));
+    let mut index = 0_u32;
+    group.bench_function("BA8", |b| {
+        b.iter(|| {
+            let data = prss.generate::<(BA8, _), _>(index);
+            index += 1;
+            black_box(data);
+        })
+    });
+    group.bench_function("BA64", |b| {
+        b.iter(|| {
+            let data = prss.generate::<(BA64, _), _>(index);
+            index += 1;
+            black_box(data);
+        })
+    });
+    group.finish();
+
+    let mut group = c.benchmark_group("prss_two_chunks");
+    group.throughput(Throughput::Bytes(64));
+    group.bench_function("BA256", |b| {
+        b.iter(|| {
+            let data = prss.generate::<(BA256, _), _>(index);
+            index += 1;
+            black_box(data);
+        })
+    });
+}
+
+criterion_group!(benches, prss_benchmark);
+criterion_main!(benches);