feat(gbench): customize batch-sizes of gbench through cli args

gbench/Cargo.toml

@@ -17,10 +17,10 @@ generic-array = "0.14.4"
 log = "0.4.8"
 neptune = { path = "../", default-features = false, features=["gpu"] }
 rust-gpu-tools = { version = "0.3.0", optional = true }
+structopt = { version = "0.3", default-features = false }
 
 [features]
 default = ["pairing", "gpu"]
 gpu = ["neptune/gpu", "rust-gpu-tools"]
 pairing = ["neptune/pairing", "bellperson/pairing"]
 blst = ["neptune/blst", "bellperson/blst"]
-

gbench/README.md

@@ -2,8 +2,15 @@
 
 Benchmarking tool for Neptune.
 
+## Usage
+
+Running gbench with customized batch sizes on the default GPU:
+
+`RUST_LOG=info cargo run -- --max-tree-batch-size 700000 --max-column-batch-size 400000`
+
+
 ## Environment variables
 
  - `NEPTUNE_GBENCH_GPUS=<bus-id1>,<bus-id2>,...` allows you to select the GPUs you want to run gbench on given a comma-separated list of bus-ids.
- 
+
 (Bus-id is a decimal integer that can be found through `nvidia-smi`, `rocm-smi`, `lspci` and etc.)

gbench/src/main.rs

@@ -12,6 +12,7 @@ use rust_gpu_tools::opencl::GPUSelector;
 use std::result::Result;
 use std::thread;
 use std::time::Instant;
+use structopt::StructOpt;
 
 fn bench_column_building(
     log_prefix: &str,
@@ -96,17 +97,28 @@ fn bench_column_building(
     res[res.len() - 1]
 }
 
+#[derive(Debug, StructOpt, Clone, Copy)]
+#[structopt(name = "Neptune gbench", about = "Neptune benchmarking program")]
+struct Opts {
+    #[structopt(long = "max-tree-batch-size", default_value = "700000")]
+    max_tree_batch_size: usize,
+    #[structopt(long = "max-column-batch-size", default_value = "400000")]
+    max_column_batch_size: usize,
+}
+
 fn main() -> Result<(), Error> {
     #[cfg(all(feature = "gpu", target_os = "macos"))]
     unimplemented!("Running on macos is not recommended and may have bad consequences -- experiment at your own risk.");
     env_logger::init();
 
+    let opts = Opts::from_args();
+
     let kib = 1024 * 1024 * 4; // 4GiB
                                // let kib = 1024 * 512; // 512MiB
     let bytes = kib * 1024;
     let leaves = bytes / 32;
-    let max_column_batch_size = 400000;
-    let max_tree_batch_size = 700000;
+    let max_column_batch_size = opts.max_column_batch_size;
+    let max_tree_batch_size = opts.max_tree_batch_size;
 
     info!("KiB: {}", kib);
     info!("leaves: {}", leaves);