Merge pull request #2 from FoodChain1028/extend_ffi

Fix the issues mentioned in the PR to mopro/main

mopro-core/Cargo.toml

@@ -13,6 +13,7 @@ wasmer = { git = "https://github.com/oskarth/wasmer.git", rev = "09c7070" }
 [features]
 default = []
 dylib = ["wasmer/dylib"]
+gpu-benchmarks = ["jemalloc-ctl", "jemallocator"]
 
 [dependencies]
 ark-circom = { git = "https://github.com/arkworks-rs/circom-compat.git" }
@@ -44,8 +45,10 @@ ark-zkey = { path = "../ark-zkey" }
 thiserror = "=1.0.39"
 color-eyre = "=0.6.2"
 criterion = "=0.3.6"
-jemalloc-ctl = "0.5.4"
-jemallocator = "0.5.4"
+
+# GPU benchmarks
+jemalloc-ctl = { version = "0.5.4", optional = true}
+jemallocator = { version = "0.5.4", optional = true}
 
 [build-dependencies]
 color-eyre = "0.6"
@@ -54,4 +57,4 @@ wasmer = { git = "https://github.com/oskarth/wasmer.git", rev = "09c7070" }
 
 [[bin]]
 name = "generate_benchmark_report"
-path = "src/middleware/gpu_exploration/bin/generate_benchmark_report.rs"
+path = "src/middleware/gpu_explorations/bin/generate_benchmark_report.rs"

mopro-core/benchmarks/gpu_explorations/msm_bench.csv

@@ -0,0 +1,22 @@
+num_msm,avg_processing_time(ms),total_processing_time(ms),memory_allocated(Bytes)
+1,1.318041,1.318041,52224
+500,0.358647444,179.323722,52224
+1000,0.303445498,303.445498,52224
+1500,0.30560589400000004,458.408841,52224
+2000,0.307506974,615.013948,52224
+2500,0.30087727480000004,752.1931870000001,52224
+3000,0.2974580376666667,892.3741130000001,52224
+3500,0.29932860400000005,1047.650114,52224
+4000,0.29971050875,1198.842035,52224
+4500,0.29770848288888885,1339.6881729999998,52224
+5000,0.29796124,1489.8062,52224
+5500,0.29821722199999995,1640.1947209999998,52224
+6000,0.29704238,1782.25428,52224
+6500,0.2977782895384615,1935.558882,52224
+7000,0.2997174805714286,2098.0223640000004,52224
+7500,0.29769822346666663,2232.736676,52224
+8000,0.29741228725,2379.298298,52224
+8500,0.29773156623529407,2530.718313,52224
+9000,0.2976449374444445,2678.8044370000002,52224
+9500,0.29733350284210525,2824.6682769999998,52224
+10000,0.29742018890000005,2974.201889,52224

mopro-core/src/middleware/gpu_exploration/bin/generate_benchmark_report.rs → mopro-core/src/middleware/gpu_explorations/bin/generate_benchmark_report.rs

@@ -1,22 +1,24 @@
 // This file is used to generate the benchmark report for the GPU exploration middleware.
 
-use std::{env, fs::File, io::Write};
-
-use mopro_core::middleware::gpu_exploration::run_msm_benchmark;
+#[cfg(feature = "gpu-benchmarks")]
+use {
+    mopro_core::middleware::gpu_explorations::run_msm_benchmark,
+    std::{cmp, env, fs::File, io::Write},
+};
 
+#[cfg(feature = "gpu-benchmarks")]
 fn main() {
     let path = env::current_dir()
         .unwrap()
-        .join("src/middleware/gpu_exploration/msm_bench.csv");
+        .join("benchmarks/gpu_explorations/msm_bench.csv");
     let mut file = File::create(path).unwrap();
     writeln!(
         file,
-        "num_msm,avg_processing_time(sec),total_processing_time(sec),memory_allocated(MiB)"
+        "num_msm,avg_processing_time(ms),total_processing_time(ms),memory_allocated(Bytes)"
     )
     .unwrap();
-    // generate 30 figures to run (range from 1 to 1000)
-    let trials = (1..1000).step_by(30);
-
+    // generate trials = [1, 500, 1_000, 1_500, ..., 10_000]
+    let trials: Vec<u32> = (0..21).map(|i| cmp::max(i * 500, 1)).collect();
     for each in trials {
         let bench_data = run_msm_benchmark(Some(each)).unwrap();
         writeln!(

mopro-core/src/middleware/gpu_exploration/mod.rs → mopro-core/src/middleware/gpu_explorations/mod.rs

@@ -11,12 +11,9 @@ pub struct BenchmarkResult {
     pub num_msm: u32,
     pub avg_processing_time: f64,
     pub total_processing_time: f64,
-    pub allocated_memory: f64,
+    pub allocated_memory: u32,
 }
 
-#[global_allocator]
-static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
-
 fn single_msm() -> Result<(), Box<dyn Error>> {
     let mut rng = ark_std::test_rng();
 
@@ -48,15 +45,13 @@ pub fn run_msm_benchmark(num_msm: Option<u32>) -> Result<BenchmarkResult, Box<dy
     }
     mem_epoch.advance().unwrap(); // Update msm memory usage
 
-    let allocated_size = allocated.read().unwrap() as f64 / usize::pow(1_024, 2) as f64; // Convert to MiB
-
-    let msm_avg = total_msm / num_msm.try_into().unwrap();
-    let msm_avg = msm_avg.subsec_nanos() as f64 / 1_000_000_000f64 + (msm_avg.as_secs() as f64);
+    let allocated_size = allocated.read().unwrap() as u32; // in Bytes
+    let msm_avg = (total_msm.as_secs_f64() / num_msm as f64) * 1_000.0; // in ms
 
     Ok(BenchmarkResult {
         num_msm,
         avg_processing_time: msm_avg,
-        total_processing_time: total_msm.as_secs_f64(),
+        total_processing_time: total_msm.as_secs_f64() * 1_000.0,
         allocated_memory: allocated_size,
     })
 }
@@ -75,11 +70,11 @@ mod tests {
         let benchmarks = run_msm_benchmark(None).unwrap();
         println!("\nBenchmarking {:?} msm on BN254 curve", benchmarks.num_msm);
         println!(
-            "└─ Average msm time: {:.5} seconds\n└─ Overall processing time: {:.5} seconds",
+            "└─ Average msm time: {:.5} ms\n└─ Overall processing time: {:.5} ms",
             benchmarks.avg_processing_time, benchmarks.total_processing_time
         );
         println!(
-            "└─ Memory allocated: {:.5} MiB",
+            "└─ Memory allocated: {:.5} Bytes",
             benchmarks.allocated_memory,
         );
     }

mopro-core/src/middleware/mod.rs

@@ -1,2 +1,4 @@
 pub mod circom;
-pub mod gpu_exploration;
+
+#[cfg(feature = "gpu-benchmarks")]
+pub mod gpu_explorations;

mopro-ffi/Cargo.toml

@@ -16,8 +16,9 @@ path = "uniffi-bindgen.rs"
 [features]
 default = []
 
-# If we enable dylib here it should be enabled in mopro-core as well
+# If we enable them here, they should be enabled in mopro-core as well
 dylib = ["mopro-core/dylib"]
+gpu-benchmarks = ["mopro-core/gpu-benchmarks"]
 
 [patch.crates-io]
 # NOTE: Forked wasmer to work around memory limits

mopro-ffi/src/lib.rs

@@ -1,7 +1,9 @@
 use mopro_core::middleware::circom;
-use mopro_core::middleware::gpu_exploration::{self, BenchmarkResult};
 use mopro_core::MoproError;
 
+#[cfg(feature = "gpu-benchmarks")]
+use mopro_core::middleware::gpu_explorations::{self, BenchmarkResult};
+
 use num_bigint::BigInt;
 use std::collections::HashMap;
 use std::path::Path;
@@ -27,6 +29,16 @@ pub struct SetupResult {
     pub provingKey: Vec<u8>,
 }
 
+// NOTE: Need to hardcode the types here, otherwise UniFFI will complain if the gpu-benchmarks feature is not enabled
+#[derive(Debug, Clone)]
+#[cfg(not(feature = "gpu-benchmarks"))]
+pub struct BenchmarkResult {
+    pub num_msm: u32,
+    pub avg_processing_time: f64,
+    pub total_processing_time: f64,
+    pub allocated_memory: u32,
+}
+
 //     pub inputs: Vec<u8>,
 
 impl From<mopro_core::MoproError> for FFIError {
@@ -159,11 +171,18 @@ impl MoproCircom {
     }
 }
 
+#[cfg(feature = "gpu-benchmarks")]
 pub fn run_msm_benchmark(num_msm: Option<u32>) -> Result<BenchmarkResult, MoproError> {
-    let benchmarks = gpu_exploration::run_msm_benchmark(num_msm).unwrap();
+    let benchmarks = gpu_explorations::run_msm_benchmark(num_msm).unwrap();
     Ok(benchmarks)
 }
 
+#[cfg(not(feature = "gpu-benchmarks"))]
+pub fn run_msm_benchmark(num_msm: Option<u32>) -> Result<BenchmarkResult, MoproError> {
+    println!("gpu-benchmarks feature not enabled!");
+    panic!("gpu-benchmarks feature not enabled!");
+}
+
 fn add(a: u32, b: u32) -> u32 {
     a + b
 }
@@ -300,15 +319,16 @@ mod tests {
     }
 
     #[test]
+    #[cfg(feature = "gpu-benchmarks")]
     fn test_run_msm_benchmark() -> Result<(), MoproError> {
         let benchmarks = run_msm_benchmark(None).unwrap();
         println!("\nBenchmarking {:?} msm on BN254 curve", benchmarks.num_msm);
         println!(
-            "└─ Average msm time: {:.5} seconds\n└─ Overall processing time: {:.5} seconds",
+            "└─ Average msm time: {:.5} ms\n└─ Overall processing time: {:.5} ms",
             benchmarks.avg_processing_time, benchmarks.total_processing_time
         );
         println!(
-            "└─ Memory allocated: {:.5} MiB",
+            "└─ Memory allocated: {:.5} Bytes",
             benchmarks.allocated_memory,
         );
         Ok(())

mopro-ffi/src/mopro.udl

@@ -31,7 +31,7 @@ dictionary BenchmarkResult {
   u32 num_msm;
   double avg_processing_time;
   double total_processing_time;
-  double allocated_memory;
+  u32 allocated_memory;
 };
 
 [Error]