Fix bellman memory measurements

StefanosChaliasos · StefanosChaliasos · commit 12508a9b77a1 · 2023-06-18T16:07:42.000+03:00
diff --git a/.gitignore b/.gitignore
@@ -114,3 +114,5 @@ gnark/template
 .idea/
 
 tags/
+
+*.*.swp
diff --git a/_scripts/parsers/csv_parser_rust.py b/_scripts/parsers/csv_parser_rust.py
@@ -41,7 +41,9 @@ def parse_csv(csv_filename, memory_folder, circuit):
         csv_reader = csv.DictReader(file)
         for row in csv_reader:
             if row['circuit'] == circuit:
-                memory_file = os.path.join(memory_folder, f'halo2_{circuit}_memory_{row["operation"]}.txt')
+                files =  os.listdir(memory_folder)
+                memory_filename = next(f for f in files if row["operation"] in f)
+                memory_file = os.path.join(memory_folder, memory_filename)
                 ram = extract_ram_from_file(memory_file)
                 row['ram(mb)'] = ram
             csv_rows.append(row)
diff --git a/_scripts/reader/helper.py b/_scripts/reader/helper.py
@@ -8,6 +8,7 @@
 # BELLMAN PATHS
 BELLMAN = os.path.join(MAIN_DIR, "bellman_circuits")
 BELLMAN_BENCH = os.path.join(BENCHMARKS_DIR, "bellman")
+BELLMAN_BENCH_MEMORY = os.path.join(BELLMAN_BENCH, "memory")
 BELLMAN_BENCH_JSON = os.path.join(BELLMAN_BENCH, "jsons")
 # BELLMAN_CE PATHS
 BELLMAN_CE = os.path.join(MAIN_DIR, "bellman_ce_circuits")
diff --git a/_scripts/reader/process_circuit.py b/_scripts/reader/process_circuit.py
@@ -136,34 +136,43 @@ def build_command_bellman(payload, count):
     for circuit, input_path in payload.circuit.items():
         for inp in helper.get_all_input_files(input_path):
             commands.append(f"cd {helper.BELLMAN}; ")
-            output_mem_size = os.path.join(
-                helper.BELLMAN_BENCH_JSON,
-                circuit + "_" + os.path.basename(inp)
-            )
             output_bench = os.path.join(
                 helper.BELLMAN_BENCH_JSON,
                 circuit + "_bench_" + os.path.basename(inp)
             )
             input_file = os.path.join("..", inp)
-            command_mem_size: str = "RUSTFLAGS=-Awarnings cargo run --bin {binary} --release -- --input {input_file} --output {output}; ".format(
-                binary=circuit,
-                input_file=input_file,
-                output=output_mem_size
-            )
-            commands.append(command_mem_size)
             command_bench: str = "RUSTFLAGS=-Awarnings INPUT_FILE={input_file} CIRCUIT={circuit} cargo criterion --message-format=json --bench {bench} 1> {output}; ".format(
                 circuit=circuit,
                 input_file=input_file,
                 bench="benchmark_circuit",
                 output=output_bench
             )
             commands.append(command_bench)
+            # Memory commands
+            os.makedirs(f"{helper.BELLMAN_BENCH_MEMORY}/{inp}", exist_ok=True)
+            # Altough each operation need only a subset of the arguments we pass
+            # all of them for simplicity
+            os.makedirs(os.path.join(helper.BELLMAN, "tmp"), exist_ok=True)
+            for op in payload.operation:
+                cargo_cmd = "cargo run --bin {circuit} --release -- --input {inp} --phase {phase} --params {params} --proof {proof}".format(
+                    circuit=circuit,
+                    inp=input_file,
+                    phase=op,
+                    params=os.path.join("tmp", "params"),
+                    proof=os.path.join("tmp", "proof"),
+                )
+                commands.append(
+                    "RUSTFLAGS=-Awarnings {memory_cmd} -h -l {cargo} 2> {time_file} > /dev/null; ".format(
+                        memory_cmd=helper.MEMORY_CMD,
+                        cargo=cargo_cmd,
+                        time_file=f"{helper.BELLMAN_BENCH_MEMORY}/{inp}/bellman_{circuit}_memory_{op}.txt"
+                    )
+                )
             commands.append("cd ..; ")
             out = os.path.join(
                 helper.BELLMAN_BENCH,
                 "bellman_bls12_381_" + circuit + ".csv"
             )
-            
             python_command = "python3"
             try:
                 subprocess.run([python_command, "--version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
@@ -175,14 +184,20 @@ def build_command_bellman(payload, count):
                     print("Neither Python nor Python3 are installed or accessible. Please install or check your path settings.")
                     sys.exit(1)
 
-            transform_command = "{python} _scripts/parsers/criterion_rust_parser.py --framework bellman --category circuit --backend bellman --curve bls12_381 --input {inp} --criterion_json {bench} --mem_proof_json {mem} --output_csv {out}; ".format(
+            transform_command = "{python} _scripts/parsers/criterion_rust_parser.py --framework bellman --category circuit --backend bellman --curve bls12_381 --input {inp} --criterion_json {bench} --proof {proof} --output_csv {out}; ".format(
                 python=python_command,
                 inp=inp,
                 bench=output_bench,
-                mem=output_mem_size,
+                proof=os.path.join(helper.BELLMAN, "tmp", "proof"),
                 out=out
             )
             commands.append(transform_command)
+            time_merge = "python3 _scripts/parsers/csv_parser_rust.py --memory_folder {memory_folder} --time_filename {time_filename} --circuit {circuit}; ".format(
+                memory_folder=os.path.join(helper.BELLMAN_BENCH_MEMORY, inp),
+                time_filename=out,
+                circuit=circuit
+            )
+            commands.append(time_merge)
 
     # Join the commands into a single string
     command = "".join(commands)
diff --git a/bellman_circuits/.gitignore b/bellman_circuits/.gitignore
@@ -0,0 +1,3 @@
+params
+proof
+tmp/
diff --git a/bellman_circuits/bellman_utils/Cargo.toml b/bellman_circuits/bellman_utils/Cargo.toml
@@ -14,4 +14,5 @@ psutil = "3.2.2"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 bellman = "0.14.0"
-bls12_381 = "0.8.0"
+bls12_381 = "0.8.0"
+clap = { version = "4.2.7", features = ["derive"] }
diff --git a/bellman_circuits/bellman_utils/src/lib.rs b/bellman_circuits/bellman_utils/src/lib.rs
@@ -1,7 +1,27 @@
 use std::{fs::File};
 use std::fs;
-use bellman::groth16::Proof;
-use bls12_381::Bls12;
+use bellman::{Circuit, groth16};
+use bellman::groth16::{Proof, Parameters};
+use bls12_381::{Bls12, Scalar};
+use clap::Parser;
+use rand::rngs::OsRng;
+
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+pub struct BinaryArgs {
+    #[arg(short, long)]
+    pub input: String,
+
+    #[arg(short, long)]
+    pub phase: String,
+
+    #[arg(short, long)]
+    pub params: Option<String>,
+
+    #[arg(short, long)]
+    pub proof: Option<String>,
+}
 
 pub fn measure_size_in_bytes(proof: &Proof<Bls12>) -> usize {
     // TODO: Should we serialize the proof in another format?
@@ -21,4 +41,44 @@ pub fn measure_size_in_bytes(proof: &Proof<Bls12>) -> usize {
     fs::remove_file(&temp_file_path).expect("Cannot remove temp file");
 
     return size_in_mb;
+}
+
+pub fn save_params(params_file: String, params: Parameters<Bls12>) {
+    let mut file = File::create(&params_file).expect("Failed to create file");
+    // Write the init_params to the file
+    params.write(&mut file).expect("Failed to write params to file");
+}
+
+pub fn load_params(params_file: String) -> Parameters<Bls12> {
+    let mut file = File::open(&params_file).expect("Failed to open file");
+    Parameters::read(&mut file, true).expect("Failed to read params from file")
+}
+
+pub fn save_proof(proof_file: String, proof: Proof<Bls12>) {
+    let mut file = File::create(&proof_file).expect("Failed to create file");
+    // Write the proof to the file
+    proof.write(&mut file).expect("Failed to write proof to file");
+}
+
+pub fn load_proof(proof_file: String) -> Proof<Bls12> {
+    let mut file = File::open(&proof_file).expect("Failed to open file");
+    Proof::read(&mut file).expect("Failed to read proof from file")
+}
+
+pub fn f_setup<C: Circuit<Scalar> + Clone>(circuit: C, params_file: String) {
+    let params = groth16::generate_random_parameters::<Bls12, _, _>(circuit.clone(), &mut OsRng).unwrap();
+    save_params(params_file, params);
+}
+
+pub fn f_prove<C: Circuit<Scalar> + Clone>(circuit: C, params_file: String, proof_file: String) {
+    let params = load_params(params_file);
+    let proof = groth16::create_random_proof(circuit, &params, &mut OsRng).unwrap();
+    save_proof(proof_file, proof);
+}
+
+pub fn f_verify(params_file: String, proof_file: String, public_input: Vec<Scalar>) {
+    let params = load_params(params_file);
+    let pvk = groth16::prepare_verifying_key(&params.vk);
+    let proof = load_proof(proof_file);
+    assert!(groth16::verify_proof(&pvk, &proof, &public_input).is_ok());
 }
diff --git a/bellman_circuits/src/bin/exponentiate.rs b/bellman_circuits/src/bin/exponentiate.rs
@@ -1,87 +1,55 @@
-// use bellman_circuits::benches::benchmark_circuit; // Assuming this is the path to the bench_proof function
 use bellman_circuits::circuits::exponentiate;
 use clap::{Parser};
 use rust_utils::{
-    get_memory,
     read_file_contents,
-    save_results,
 };
-use bellman_utils::measure_size_in_bytes;
-use bellman::groth16;
+use bellman_utils::{BinaryArgs, f_setup, f_verify, f_prove};
 use bellman::gadgets::multipack;
-use bls12_381::{Bls12, Scalar};
-use rand::rngs::OsRng;
+use bls12_381::Scalar;
 use ff::PrimeField;
 
-#[derive(Parser, Debug)]
-#[clap(
-    name = "MemoryBenchExponentiate",
-    about = "MemoryBenchExponentiate CLI is a CLI Application to Benchmark memory consumption of Exponentiate",
-    version = "0.0.1"
-)]
-
-struct Args {
-    #[arg(short, long)]
-    input: String,
-
-    #[arg(short, long)]
-    output: String,
-}
-
 fn main() {
     // Parse command line arguments
-    let args = Args::parse();
+    let args = BinaryArgs::parse();
 
     // Read and parse input from the specified JSON file
     let input_str = read_file_contents(args.input);
 
-     // Get data from config
-     let (x_64, e, y_64) = exponentiate::get_exponentiate_data(input_str);
-
-     // Create Scalar from some values
-     let x = Scalar::from(x_64);
-     let y = Scalar::from(y_64);
-
-      // Public inputs are x and y
-    let x_bits = multipack::bytes_to_bits_le(&x.to_repr().as_ref());
-    let y_bits = multipack::bytes_to_bits_le(&y.to_repr().as_ref());
-    let inputs = [multipack::compute_multipacking(&x_bits), multipack::compute_multipacking(&y_bits)].concat();
-
-
-    // Define the circuit
-    let circuit = exponentiate::ExponentiationCircuit {
-        x: Some(x),
-        e: e,
-        y: Some(y),
-    };
-
-    // Get the initial memory usage
-    let initial_rss = get_memory();
-
-    // Generate Parameters
-    let params = groth16::generate_random_parameters::<Bls12, _, _>(circuit.clone(), &mut OsRng).unwrap();
-
-    // Prepare the verification key
-    let pvk = groth16::prepare_verifying_key(&params.vk);
-
-    // Get the memory usage after setup
-    let setup_rss = get_memory();
-
-    // Create a Groth16 proof with our parameters
-    let proof = groth16::create_random_proof(circuit, &params, &mut OsRng).unwrap();
-
-    // Get the memory usage after proof generation
-    let proof_rss = get_memory();
-
-    // Verify the proof
-    let _ = groth16::verify_proof(&pvk, &proof, &inputs);
-
-    // Get the memory usage after proof verification
-    let verify_rss = get_memory();
+    // Get data from config
+    let (x_64, e, y_64) = exponentiate::get_exponentiate_data(input_str);
+
+    // Create Scalar from some values
+    let x = Scalar::from(x_64);
+    let y = Scalar::from(y_64);
+
+    if args.phase == "setup" {
+        let circuit = exponentiate::ExponentiationCircuit {
+            x: Some(x),
+            e: e,
+            y: Some(y),
+        };
+        let params_file = args.params.expect("Missing params argument");
+        f_setup(circuit, params_file);
+    } else if args.phase == "prove" {
+        let circuit = exponentiate::ExponentiationCircuit {
+            x: Some(x),
+            e: e,
+            y: Some(y),
+        };
+        let params_file = args.params.expect("Missing params argument");
+        let proof_file = args.proof.expect("Missing proof argument");
+        f_prove(circuit, params_file, proof_file);
+    } else if args.phase == "verify" {
+        // Public inputs are x and y
+        let x_bits = multipack::bytes_to_bits_le(&x.to_repr().as_ref());
+        let y_bits = multipack::bytes_to_bits_le(&y.to_repr().as_ref());
+        let inputs: Vec<Scalar> = [multipack::compute_multipacking(&x_bits), multipack::compute_multipacking(&y_bits)].concat();
+        let params_file = args.params.expect("Missing params argument");
+        let proof_file = args.proof.expect("Missing proof argument");
+        f_verify(params_file, proof_file, inputs)
+    } else {
+        panic!("Invalid phase (should be setup, prove, or verify)");
+    }
 
-    // Measure the proof size
-    let proof_size = measure_size_in_bytes(&proof);
 
-    // Save the results
-    save_results(initial_rss, setup_rss, proof_rss, verify_rss, proof_size, args.output);
 }
diff --git a/bellman_circuits/src/bin/sha256.rs b/bellman_circuits/src/bin/sha256.rs

-Original file line number
+Diff line change
 .idea/
 tags/
++
 +*.*.swp