From d70dfcdf7ce50ce0ecc42432daaec64c8a0b1911 Mon Sep 17 00:00:00 2001
From: Jacob Trombetta <jacob.trombetta@spaceandtime.io>
Date: Fri, 12 Jul 2024 16:22:10 -0400
Subject: [PATCH] bench: add Jaeger benchmarks for Dory proof evaluation
 (PROOF-892) (#47)

# Rationale for this change
In order to run benchmarks easily on Dory proof evaluation performance,
this PR adds a new benchmark type to `jaeger_benches`. An argument is
added to the command line interface to define the benchmark type.

Previously, the Jaeger benchmarks were only run on the inner product
proof with the command:
`cargo bench -p proof-of-sql --bench jaeger_benches`
which has been updated to
`cargo bench -p proof-of-sql --bench jaeger_benches InnerProductProof`.

The Dory proof evaluation benchmark can be run with
`cargo bench -p proof-of-sql --bench jaeger_benches Dory
features="test"`.

# What changes are included in this PR?
- A Jaeger benchmark is added for Dory.
- A command line argument is added to specify which benchmark to run:
`InnerProductProof` or `Dory --features="test"`

# Are these changes tested?
Yes
---
 crates/proof-of-sql/benches/README.md         |  5 +-
 crates/proof-of-sql/benches/jaeger_benches.rs | 53 ++++++++++++++++---
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/crates/proof-of-sql/benches/README.md b/crates/proof-of-sql/benches/README.md
index bda601fb8..4419dce19 100644
--- a/crates/proof-of-sql/benches/README.md
+++ b/crates/proof-of-sql/benches/README.md
@@ -8,9 +8,10 @@ To run benchmarks with Jaeger, you need to do the following
     ```bash
     docker run --rm -d --name jaeger -p 6831:6831/udp -p 16686:16686 jaegertracing/all-in-one:latest
     ```
-2. Run the benchmark.
+2. Run a benchmark.
     ```bash
-    cargo bench -p proof-of-sql --bench jaeger_benches
+    cargo bench -p proof-of-sql --bench jaeger_benches InnerProductProof
+    cargo bench -p proof-of-sql --bench jaeger_benches Dory --features="test"
     ```
 3. Navigate to http://localhost:16686/ to see the results.
 4. To end the Jaeger service, run
diff --git a/crates/proof-of-sql/benches/jaeger_benches.rs b/crates/proof-of-sql/benches/jaeger_benches.rs
index 992bc998e..63313fa29 100644
--- a/crates/proof-of-sql/benches/jaeger_benches.rs
+++ b/crates/proof-of-sql/benches/jaeger_benches.rs
@@ -2,14 +2,21 @@
 //! To run, execute the following commands:
 //! ```bash
 //! docker run --rm -d --name jaeger -p 6831:6831/udp -p 16686:16686 jaegertracing/all-in-one:latest
-//! cargo bench -p proof-of-sql --bench jaeger_benches
+//! cargo bench -p proof-of-sql --bench jaeger_benches InnerProductProof
+//! cargo bench -p proof-of-sql --bench jaeger_benches Dory --features="test"
 //! ```
 //! Then, navigate to http://localhost:16686 to view the traces.
 
 use blitzar::{compute::init_backend, proof::InnerProductProof};
+#[cfg(feature = "test")]
+use proof_of_sql::proof_primitive::dory::{
+    DoryEvaluationProof, DoryProverPublicSetup, DoryVerifierPublicSetup, ProverSetup,
+    PublicParameters, VerifierSetup,
+};
 mod scaffold;
 use crate::scaffold::querys::QUERIES;
 use scaffold::jaeger_scaffold;
+use std::env;
 
 const SIZE: usize = 1_000_000;
 
@@ -25,13 +32,47 @@ fn main() {
         .with(opentelemetry)
         .try_init()
         .unwrap();
-    {
-        // Run 3 times to ensure that warm-up of the GPU has occured.
-        for _ in 0..3 {
-            for (title, query, columns) in QUERIES.iter() {
-                jaeger_scaffold::<InnerProductProof>(title, query, columns, SIZE, &(), &());
+
+    // Check for command-line arguments to select the benchmark type.
+    let args: Vec<String> = env::args().collect();
+    let benchmark_type = args
+        .get(1)
+        .expect("Please specify the benchmark type: InnerProductProof or Dory");
+
+    match benchmark_type.as_str() {
+        "InnerProductProof" => {
+            // Run 3 times to ensure that warm-up of the GPU has occurred.
+            for _ in 0..3 {
+                for (title, query, columns) in QUERIES.iter() {
+                    jaeger_scaffold::<InnerProductProof>(title, query, columns, SIZE, &(), &());
+                }
+            }
+        }
+        #[cfg(feature = "test")]
+        "Dory" => {
+            // Run 3 times to ensure that warm-up of the GPU has occurred.
+            let pp =
+                PublicParameters::rand(10, &mut proof_of_sql::proof_primitive::dory::test_rng());
+            let ps = ProverSetup::from(&pp);
+            let prover_setup = DoryProverPublicSetup::new(&ps, 10);
+            let vs = VerifierSetup::from(&pp);
+            let verifier_setup = DoryVerifierPublicSetup::new(&vs, 10);
+
+            for _ in 0..3 {
+                for (title, query, columns) in QUERIES.iter() {
+                    jaeger_scaffold::<DoryEvaluationProof>(
+                        title,
+                        query,
+                        columns,
+                        SIZE,
+                        &prover_setup,
+                        &verifier_setup,
+                    );
+                }
             }
         }
+        _ => panic!("Invalid benchmark type specified."),
     }
+
     opentelemetry::global::shutdown_tracer_provider();
 }