Distributed prover

Author: Champii

Cargo.lock

@@ -79,6 +79,7 @@ p3-uni-stark = "0.1.3-succinct"
 p3-maybe-rayon = "0.1.3-succinct"
 p3-bn254-fr = "0.1.3-succinct"
 
+
 # For local development. 
 
 # p3-air = { path = "../Plonky3/air" }
@@ -101,3 +102,12 @@ p3-bn254-fr = "0.1.3-succinct"
 # p3-uni-stark = { path = "../Plonky3/uni-stark" }
 # p3-maybe-rayon = { path = "../Plonky3/maybe-rayon" }
 # p3-bn254-fr = { path = "../Plonky3/bn254-fr" }
+
+# Patch Plonky3 for Serialize and Deserialize of DuplexChallenger
+[patch.crates-io]
+p3-field = { git = "https://github.com/Champii/Plonky3.git", branch = "serde_patch" }
+p3-challenger = { git = "https://github.com/Champii/Plonky3.git", branch = "serde_patch" }
+p3-poseidon2 = { git = "https://github.com/Champii/Plonky3.git", branch = "serde_patch" }
+p3-baby-bear = { git = "https://github.com/Champii/Plonky3.git", branch = "serde_patch" }
+p3-symmetric = { git = "https://github.com/Champii/Plonky3.git", branch = "serde_patch" }
+

Cargo.toml

@@ -6,6 +6,7 @@ mod options;
 #[cfg(any(test, feature = "programs"))]
 mod programs;
 mod prove;
+pub mod prove_distributed;
 mod serde;
 mod tracer;
 

core/src/utils/mod.rs

@@ -0,0 +1,154 @@
+use std::sync::{mpsc::SyncSender, Arc};
+
+pub use crate::{air::PublicValues, runtime::Program, stark::RiscvAir};
+
+use crate::{
+    runtime::{ExecutionRecord, NoOpSubproofVerifier, Runtime},
+    stark::{MachineProver, MachineRecord},
+    utils::{baby_bear_poseidon2::Val, BabyBearPoseidon2, SP1CoreOpts},
+};
+
+use super::Checkpoint;
+
+fn trace_checkpoint(
+    program: Program,
+    checkpoint: Checkpoint,
+    opts: SP1CoreOpts,
+) -> (Vec<ExecutionRecord>, Checkpoint) {
+    let mut runtime = Runtime::recover(program, checkpoint, opts);
+
+    runtime.subproof_verifier = Arc::new(NoOpSubproofVerifier);
+
+    let (events, _) =
+        tracing::debug_span!("runtime.trace").in_scope(|| runtime.execute_record().unwrap());
+
+    let state = runtime.state.clone();
+
+    (events, state)
+}
+
+pub fn process<P: MachineProver<BabyBearPoseidon2, RiscvAir<Val>>>(
+    prover: &P,
+    program: &Program,
+    checkpoint: Checkpoint,
+    nb_checkpoints: usize,
+    state: PublicValues<u32, u32>,
+    opts: SP1CoreOpts,
+    records_tx: SyncSender<Vec<ExecutionRecord>>,
+    deferred: &mut ExecutionRecord,
+    is_deferred: bool,
+) {
+    if is_deferred {
+        process_deferred(program, checkpoint, state, opts, records_tx, deferred);
+    } else {
+        process_regular(
+            prover,
+            program,
+            checkpoint,
+            nb_checkpoints,
+            state,
+            opts,
+            records_tx,
+            deferred,
+        );
+    }
+}
+
+fn process_regular<P: MachineProver<BabyBearPoseidon2, RiscvAir<Val>>>(
+    prover: &P,
+    program: &Program,
+    mut checkpoint: Checkpoint,
+    nb_checkpoints: usize,
+    mut state: PublicValues<u32, u32>,
+    opts: SP1CoreOpts,
+    records_tx: SyncSender<Vec<ExecutionRecord>>,
+    deferred: &mut ExecutionRecord,
+) {
+    tracing::debug_span!("phase 1 record generator").in_scope(|| {
+        let mut processed_checkpoints = 0;
+
+        while processed_checkpoints < nb_checkpoints {
+            log::info!(
+                "Processing checkpoint {}/{}",
+                processed_checkpoints + 1,
+                nb_checkpoints
+            );
+            // Trace the checkpoint and reconstruct the execution records.
+            let (mut records, new_checkpoint) = tracing::debug_span!("trace checkpoint")
+                .in_scope(|| trace_checkpoint(program.clone(), checkpoint, opts));
+
+            checkpoint = new_checkpoint;
+
+            // Update the public values & prover state for the shards which contain "cpu events".
+            for record in records.iter_mut() {
+                state.shard += 1;
+                state.execution_shard = record.public_values.execution_shard;
+                state.start_pc = record.public_values.start_pc;
+                state.next_pc = record.public_values.next_pc;
+                record.public_values = state;
+            }
+
+            // Generate the dependencies.
+            tracing::debug_span!("generate dependencies")
+                .in_scope(|| prover.machine().generate_dependencies(&mut records, &opts));
+
+            // Defer events that are too expensive to include in every shard.
+            for record in records.iter_mut() {
+                deferred.append(&mut record.defer());
+            }
+
+            // See if any deferred shards are ready to be commited to.
+            let mut _deferred = deferred.split(false, opts.split_opts);
+
+            // Update the public values & prover state for the shards which do not contain "cpu events"
+            // before committing to them.
+            state.execution_shard += 1;
+
+            records_tx.send(records).unwrap();
+
+            processed_checkpoints += 1;
+        }
+    });
+}
+
+fn process_deferred(
+    program: &Program,
+    checkpoint: Checkpoint,
+    mut state: PublicValues<u32, u32>,
+    opts: SP1CoreOpts,
+    records_tx: SyncSender<Vec<ExecutionRecord>>,
+    deferred: &mut ExecutionRecord,
+) {
+    tracing::debug_span!("phase 1 record generator").in_scope(|| {
+        // Trace the checkpoint and reconstruct the execution records.
+        let (mut records, _) = tracing::debug_span!("trace checkpoint")
+            .in_scope(|| trace_checkpoint(program.clone(), checkpoint, opts));
+
+        // Update the public values & prover state for the shards which contain "cpu events".
+        for record in records.iter_mut() {
+            // state.shard += 1;
+            state.execution_shard = record.public_values.execution_shard;
+            state.start_pc = record.public_values.start_pc;
+            state.next_pc = record.public_values.next_pc;
+            record.public_values = state;
+        }
+
+        // See if any deferred shards are ready to be commited to.
+        let mut deferred = deferred.split(true, opts.split_opts);
+
+        // Update the public values & prover state for the shards which do not contain "cpu events"
+        // before committing to them.
+
+        for record in deferred.iter_mut() {
+            state.shard += 1;
+            state.previous_init_addr_bits = record.public_values.previous_init_addr_bits;
+            state.last_init_addr_bits = record.public_values.last_init_addr_bits;
+            state.previous_finalize_addr_bits = record.public_values.previous_finalize_addr_bits;
+            state.last_finalize_addr_bits = record.public_values.last_finalize_addr_bits;
+            state.start_pc = state.next_pc;
+            record.public_values = state;
+        }
+
+        records_tx.send(deferred).unwrap();
+    });
+}