Add Circuit and Witness serialization for GPU prover

Cargo.toml

@@ -12,6 +12,7 @@ members = [
     "config_macros", # proc macros used to declare a new config, this has to a separate crate due to rust compilation issues
     "gkr",
     "gkr_engine", # definitions of GKR engine and associated types
+    "gpu", # GPU support and circuit serialization
     "hasher", # definitions of FiatShamirFieldHasher, FiatShamirBytesHash, and associated types
     "poly_commit",
     "serdes", # serialization and deserialization of various data structures

bin/src/main.rs

@@ -8,7 +8,7 @@ use circuit::Circuit;
 use clap::Parser;
 use gkr::{
     BN254ConfigMIMC5KZG, BN254ConfigSha2Hyrax, BN254ConfigSha2Raw, GF2ExtConfigSha2Orion,
-    GF2ExtConfigSha2Raw, Goldilocksx8ConfigSha2Orion, Goldilocksx8ConfigSha2Raw,
+    GF2ExtConfigSha2Raw, Goldilocksx1ConfigSha2Raw, Goldilocksx8ConfigSha2Orion, Goldilocksx8ConfigSha2Raw,
     M31x1ConfigSha2RawVanilla, M31x16ConfigSha2OrionSquare, M31x16ConfigSha2OrionVanilla,
     M31x16ConfigSha2RawSquare, M31x16ConfigSha2RawVanilla, Prover,
     utils::{
@@ -69,7 +69,13 @@ fn main() {
 
         "m31ext3" => match pcs_type {
             PolynomialCommitmentType::Raw => match args.circuit.as_str() {
-                "keccak" => run_benchmark::<M31x16ConfigSha2RawVanilla>(&args, mpi_config.clone()),
+                "keccak" => {
+                    if std::env::var("EXPANDER_GPU").map_or(false, |v| v == "1") {
+                        run_benchmark::<M31x1ConfigSha2RawVanilla>(&args, mpi_config.clone())
+                    } else {
+                        run_benchmark::<M31x16ConfigSha2RawVanilla>(&args, mpi_config.clone())
+                    }
+                },
                 "poseidon" => run_benchmark::<M31x16ConfigSha2RawSquare>(&args, mpi_config.clone()),
                 _ => unreachable!(),
             },
@@ -112,7 +118,13 @@ fn main() {
         },
         "goldilocks" => match pcs_type {
             PolynomialCommitmentType::Raw => match args.circuit.as_str() {
-                "keccak" => run_benchmark::<Goldilocksx8ConfigSha2Raw>(&args, mpi_config.clone()),
+                "keccak" => {
+                    if std::env::var("EXPANDER_GPU").map_or(false, |v| v == "1") {
+                        run_benchmark::<Goldilocksx1ConfigSha2Raw>(&args, mpi_config.clone())
+                    } else {
+                        run_benchmark::<Goldilocksx8ConfigSha2Raw>(&args, mpi_config.clone())
+                    }
+                },
                 _ => unreachable!(),
             },
             PolynomialCommitmentType::Orion => match args.circuit.as_str() {
@@ -206,6 +218,7 @@ where
         (FieldType::M31x1, "keccak") => 2,
         (FieldType::M31x16, "keccak") => 2,
         (FieldType::M31x16, "poseidon") => 120,
+        (FieldType::Goldilocksx1, "keccak") => 2,
         (FieldType::Goldilocksx8, "keccak") => 2,
         (FieldType::BabyBearx16, "keccak") => 2,
         (FieldType::BN254, "keccak") => 2,

gkr/Cargo.toml

@@ -13,6 +13,7 @@ gf2_128 = { path = "../arith/gf2_128" }
 gkr_engine = { path = "../gkr_engine" }
 gkr_hashers = { path = "../hasher" }
 goldilocks = { path = "../arith/goldilocks" }
+gpu = { path = "../gpu" }
 mersenne31 = { path = "../arith/mersenne31" }
 poly_commit = { path = "../poly_commit" }
 polynomials = { path = "../arith/polynomials" }

gkr/src/prover/gkr_vanilla.rs

@@ -14,7 +14,8 @@ pub fn gkr_prove<F: FieldEngine>(
     sp: &mut ProverScratchPad<F>,
     transcript: &mut impl Transcript,
     mpi_config: &MPIConfig,
-) -> (F::ChallengeField, ExpanderDualVarChallenge<F>) {
+) -> (F::ChallengeField, ExpanderDualVarChallenge<F>) 
+where F::CircuitField: std::fmt::Debug, F::SimdCircuitField: std::fmt::Debug {
     let layer_num = circuit.layers.len();
 
     let mut challenge: ExpanderDualVarChallenge<F> =
@@ -36,6 +37,16 @@ pub fn gkr_prove<F: FieldEngine>(
         mpi_config,
     );
 
+    // Serialize circuit to file if EXPANDER_GPU environment variable is set to 1
+    if std::env::var("EXPANDER_GPU").map_or(false, |v| v == "1") {
+        if let Err(e) = gpu::serdes::serial_circuit_witness_as_plaintext(circuit, transcript, &challenge) {
+            println!("Failed to serialize circuit: {}", e);
+        }
+    }
+
+    let mut final_vx_claim = None;
+    let mut final_vy_claim = None;
+
     for i in (0..layer_num).rev() {
         let timer = Timer::new(
             &format!(
@@ -47,7 +58,7 @@ pub fn gkr_prove<F: FieldEngine>(
             mpi_config.is_root(),
         );
 
-        (_, _) = sumcheck_prove_gkr_layer(
+        let (vx_claim, vy_claim) = sumcheck_prove_gkr_layer(
             &circuit.layers[i],
             &mut challenge,
             alpha,
@@ -57,6 +68,12 @@ pub fn gkr_prove<F: FieldEngine>(
             i == layer_num - 1,
         );
 
+        // Store the final layer claims for later use
+        if i == 0 {
+            final_vx_claim = Some(vx_claim);
+            final_vy_claim = vy_claim;
+        }
+
         if challenge.rz_1.is_some() {
             // TODO: try broadcast beta.unwrap directly
             let mut tmp = transcript.generate_field_element::<F::ChallengeField>();
@@ -68,5 +85,14 @@ pub fn gkr_prove<F: FieldEngine>(
         timer.stop();
     }
 
+    // Print final claims if EXPANDER_GPU environment variable is set to 1
+    if std::env::var("EXPANDER_GPU").map_or(false, |v| v == "1") {
+        if let Some(vx) = final_vx_claim {
+            gpu::serdes::print_final_claims::<F>(&vx, &final_vy_claim);
+            println!("GKR final proof claims as shown above.");
+            std::process::exit(0);
+        }
+    }
+
     (claimed_v, challenge)
 }

gpu/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "gpu"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+circuit = { path = "../circuit" }
+gkr_engine = { path = "../gkr_engine" }
+gkr_hashers = { path = "../hasher" }
+transcript = { path = "../transcript" }
+thiserror.workspace = true 

gpu/src/lib.rs

@@ -0,0 +1 @@
+pub mod serdes;