chore: update gkr2 correctness test (M31 only)

Author: GhostOfGauss

gkr/src/tests/gkr_correctness.rs

@@ -289,35 +289,52 @@ pub fn gkr_square_test_circuit<C: GKRConfig>() -> Circuit<C> {
     circuit
 }
 
+// CMD: RUSTFLAGS="-C target-cpu=native" mpiexec -n 4 cargo test --package gkr --lib -- tests::gkr_correctness::gkr_square_correctness_test --exact --show-output
 #[test]
-fn gkr_square_correctness() {
-    type GkrConfigType = config::M31ExtConfigSha2;
-
+fn gkr_square_correctness_test() {
     env_logger::init();
+    let mpi_config = MPIConfig::new();
+    assert!(gkr_square_correctness::<config::M31ExtConfigSha2>(
+        mpi_config.clone()
+    ));
+    assert!(gkr_square_correctness::<config::M31ExtConfigKeccak>(
+        mpi_config
+    ));
+    MPIConfig::finalize();
+}
 
-    let mut circuit = gkr_square_test_circuit::<GkrConfigType>();
+fn gkr_square_correctness<C: GKRConfig>(mpi_config: MPIConfig) -> bool {
+    let config = Config::<C>::new(GKRScheme::GkrSquare, mpi_config);
+
+    let mut circuit = gkr_square_test_circuit::<C>();
     // Set input layers with N_2_0 = 3, N_2_1 = 5, N_2_2 = 7,
-    // and N_2_3 varying from 0 to 15
-    let final_vals = (0..16).map(|x| x.into()).collect::<Vec<_>>();
-    let final_vals = <GkrConfigType as GKRConfig>::SimdCircuitField::pack(&final_vals);
+    // and N_2_3 varying from 0 to SIMD packing size
+    let mut final_vals = (0..16).map(|x| x.into()).collect::<Vec<_>>();
+    // Add variety for MPI participants
+    final_vals[0] += C::CircuitField::from(config.mpi_config.world_rank as u32);
+    let final_vals = C::SimdCircuitField::pack(&final_vals);
     circuit.layers[0].input_vals = vec![2.into(), 3.into(), 5.into(), final_vals];
     // Set public input PI[0] = 13
     circuit.public_input = vec![13.into()];
 
-    let config = Config::<GkrConfigType>::new(GKRScheme::GkrSquare, MPIConfig::default());
-    do_prove_verify(config, &mut circuit);
+    let result = do_prove_verify(config, &mut circuit);
+    result
 }
 
-fn do_prove_verify<C: GKRConfig>(config: Config<C>, circuit: &mut Circuit<C>) {
+fn do_prove_verify<C: GKRConfig>(config: Config<C>, circuit: &mut Circuit<C>) -> bool {
     circuit.evaluate();
 
     // Prove
     let mut prover = Prover::new(&config);
     prover.prepare_mem(&circuit);
     let (mut claimed_v, proof) = prover.prove(circuit);
 
-    // Verify
-    let verifier = Verifier::new(&config);
-    let public_input = circuit.public_input.clone();
-    assert!(verifier.verify(circuit, &public_input, &mut claimed_v, &proof))
+    // Verify if root process
+    if config.mpi_config.is_root() {
+        let verifier = Verifier::new(&config);
+        let public_input = circuit.public_input.clone();
+        return verifier.verify(circuit, &public_input, &mut claimed_v, &proof);
+    }
+    // Non-root processes return true
+    return true;
 }