Merge remote-tracking branch 'origin/tamir/v1.3.0-rc2' into scroll-de…

…v-0914

crates/core/machine/src/riscv/shape.rs

@@ -79,7 +79,6 @@ impl<F: PrimeField32> CoreShapeConfig<F> {
             })
             .collect();
 
-        tracing::info!("Found shape");
         let shape = CoreShape { inner: shape? };
         Some(shape)
     }
@@ -297,8 +296,8 @@ impl<F: PrimeField32> Default for CoreShapeConfig<F> {
         let included_shapes = vec![];
 
         // Preprocessed chip heights.
-        let program_heights = vec![Some(16), Some(20), Some(21), Some(22)];
-        let program_memory_heights = vec![Some(16), Some(20), Some(21), Some(22)];
+        let program_heights = vec![Some(16), Some(19), Some(20), Some(21), Some(22)];
+        let program_memory_heights = vec![Some(16), Some(19), Some(20), Some(21), Some(22)];
 
         let allowed_preprocessed_log_heights = HashMap::from([
             (RiscvAir::Program(ProgramChip::default()), program_heights),
@@ -314,7 +313,7 @@ impl<F: PrimeField32> Default for CoreShapeConfig<F> {
         let shift_right_heights = vec![None, Some(10), Some(16), Some(19)];
         let shift_left_heights = vec![None, Some(10), Some(16), Some(19)];
         let lt_heights = vec![None, Some(10), Some(16), Some(19)];
-        let memory_local_heights = vec![Some(16), Some(20)];
+        let memory_local_heights = vec![Some(16), Some(19)];
         let syscall_heights = vec![None, Some(19)];
 
         let short_allowed_log_heights = HashMap::from([

crates/core/machine/src/utils/prove.rs

@@ -327,11 +327,6 @@ where
         // Create the challenger and observe the verifying key.
         let mut challenger = prover.config().challenger();
         pk.observe_into(&mut challenger);
-        // challenger.observe(pk.preprocessed_commit());
-        // challenger.observe(pk.pc_start());
-        // for _ in 0..7 {
-        //     challenger.observe(Val::<SC>::zero());
-        // }
 
         // Spawn the phase 1 prover thread.
         let phase_1_prover_span = tracing::Span::current().clone();
@@ -508,16 +503,16 @@ where
                             // Let another worker update the state.
                             record_gen_sync.advance_turn();
 
-                            #[cfg(feature = "debug")]
-                            all_records_tx.send(records.clone()).unwrap();
-
                             // Fix the shape of the records.
                             if let Some(shape_config) = shape_config {
                                 for record in records.iter_mut() {
                                     shape_config.fix_shape(record).unwrap();
                                 }
                             }
 
+                            #[cfg(feature = "debug")]
+                            all_records_tx.send(records.clone()).unwrap();
+
                             // Generate the traces.
                             let mut local_traces = Vec::new();
                             tracing::debug_span!("generate local traces", index).in_scope(|| {

crates/prover/scripts/build_compress_vks.rs

@@ -1,21 +1,26 @@
-use std::{fs::File, path::PathBuf};
+use std::path::PathBuf;
 
 use clap::Parser;
-use p3_baby_bear::BabyBear;
-use sp1_core_machine::{riscv::CoreShapeConfig, utils::setup_logger};
-use sp1_prover::{utils::get_all_vk_digests, InnerSC, REDUCE_BATCH_SIZE};
-use sp1_recursion_circuit_v2::merkle_tree::MerkleTree;
-use sp1_recursion_core_v2::shape::RecursionShapeConfig;
+use sp1_core_machine::utils::setup_logger;
+use sp1_prover::{components::DefaultProverComponents, shapes::build_vk_map, REDUCE_BATCH_SIZE};
 
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
 struct Args {
     #[clap(short, long)]
     build_dir: PathBuf,
-    #[clap(short, long)]
+    #[clap(short, long, default_value_t = false)]
     dummy: bool,
     #[clap(short, long, default_value_t = REDUCE_BATCH_SIZE)]
     reduce_batch_size: usize,
+    #[clap(short, long, default_value_t = 1)]
+    num_compiler_workers: usize,
+    #[clap(short, long, default_value_t = 1)]
+    num_setup_workers: usize,
+    #[clap(short, long)]
+    start: Option<usize>,
+    #[clap(short, long)]
+    end: Option<usize>,
 }
 
 fn main() {
@@ -24,32 +29,19 @@ fn main() {
 
     let reduce_batch_size = args.reduce_batch_size;
     let build_dir = args.build_dir;
-
-    let core_shape_config = CoreShapeConfig::default();
-    let recursion_shape_config = RecursionShapeConfig::default();
-
-    std::fs::create_dir_all(&build_dir).expect("failed to create build directory");
-
-    tracing::info!("building compress vk map");
-    let vk_map = get_all_vk_digests(&core_shape_config, &recursion_shape_config, reduce_batch_size);
-    tracing::info!("compress vks generated, number of keys: {}", vk_map.len());
-
-    // Save the vk map to a file.
-    tracing::info!("saving vk map to file");
-    let vk_map_path = build_dir.join("vk_map.bin");
-    let mut vk_map_file = File::create(vk_map_path).unwrap();
-    bincode::serialize_into(&mut vk_map_file, &vk_map).unwrap();
-    tracing::info!("File saved successfully.");
-
-    // Build a merkle tree from the vk map.
-    tracing::info!("building merkle tree");
-    let (root, merkle_tree) =
-        MerkleTree::<BabyBear, InnerSC>::commit(vk_map.keys().cloned().collect());
-
-    // Saving merkle tree data to file.
-    tracing::info!("saving merkle tree to file");
-    let merkle_tree_path = build_dir.join("merkle_tree.bin");
-    let mut merkle_tree_file = File::create(merkle_tree_path).unwrap();
-    bincode::serialize_into(&mut merkle_tree_file, &(root, merkle_tree)).unwrap();
-    tracing::info!("File saved successfully.");
+    let dummy = args.dummy;
+    let num_compiler_workers = args.num_compiler_workers;
+    let num_setup_workers = args.num_setup_workers;
+    let range_start = args.start;
+    let range_end = args.end;
+
+    build_vk_map::<DefaultProverComponents>(
+        build_dir,
+        reduce_batch_size,
+        dummy,
+        num_compiler_workers,
+        num_setup_workers,
+        range_start,
+        range_end,
+    );
 }

crates/prover/src/lib.rs

@@ -13,6 +13,7 @@
 
 pub mod build;
 pub mod components;
+pub mod shapes;
 pub mod types;
 pub mod utils;
 pub mod verify;
@@ -74,8 +75,8 @@ use sp1_recursion_core_v2::{
 use sp1_recursion_circuit_v2::{
     hash::FieldHasher,
     machine::{
-        SP1CompressRootVerifier, SP1CompressRootVerifierWithVKey, SP1CompressShape,
-        SP1CompressWithVKeyVerifier, SP1CompressWithVKeyWitnessValues, SP1CompressWitnessValues,
+        SP1CompressRootVerifier, SP1CompressRootVerifierWithVKey, SP1CompressWithVKeyVerifier,
+        SP1CompressWithVKeyWitnessValues, SP1CompressWithVkeyShape, SP1CompressWitnessValues,
         SP1DeferredVerifier, SP1DeferredWitnessValues, SP1MerkleProofWitnessValues,
         SP1RecursionShape, SP1RecursionWitnessValues, SP1RecursiveVerifier,
     },
@@ -133,19 +134,22 @@ pub struct SP1Prover<C: SP1ProverComponents = DefaultProverComponents> {
 
     pub recursion_cache_misses: AtomicUsize,
 
-    pub compress_programs: Mutex<LruCache<SP1CompressShape, Arc<RecursionProgram<BabyBear>>>>,
+    pub compress_programs:
+        Mutex<LruCache<SP1CompressWithVkeyShape, Arc<RecursionProgram<BabyBear>>>>,
 
     pub compress_cache_misses: AtomicUsize,
 
     pub root: <InnerSC as FieldHasher<BabyBear>>::Digest,
 
     pub allowed_vk_map: BTreeMap<<InnerSC as FieldHasher<BabyBear>>::Digest, usize>,
 
-    pub merkle_tree: MerkleTree<BabyBear, InnerSC>,
+    pub vk_merkle_tree: MerkleTree<BabyBear, InnerSC>,
 
     pub core_shape_config: Option<CoreShapeConfig<BabyBear>>,
 
     pub recursion_shape_config: Option<RecursionShapeConfig<BabyBear, CompressAir<BabyBear>>>,
+
+    pub vk_verification: bool,
 }
 
 impl<C: SP1ProverComponents> SP1Prover<C> {
@@ -187,9 +191,6 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         )
         .expect("PROVER_COMPRESS_CACHE_SIZE must be a non-zero usize");
 
-        let core_shape_config = CoreShapeConfig::default();
-        let recursion_shape_config = RecursionShapeConfig::default();
-
         let allowed_vk_map =
             bincode::deserialize(VK_MAP_BYTES).expect("failed to deserialize vk map");
 
@@ -199,12 +200,15 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         let core_shape_config = env::var("FIX_CORE_SHAPES")
             .map(|v| v.eq_ignore_ascii_case("true"))
             .unwrap_or(true)
-            .then_some(core_shape_config);
+            .then_some(CoreShapeConfig::default());
 
         let recursion_shape_config = env::var("FIX_RECURSION_SHAPES")
             .map(|v| v.eq_ignore_ascii_case("true"))
             .unwrap_or(true)
-            .then_some(recursion_shape_config);
+            .then_some(RecursionShapeConfig::default());
+
+        let vk_verification =
+            env::var("VERIFY_VK").map(|v| v.eq_ignore_ascii_case("true")).unwrap_or(false);
 
         Self {
             core_prover,
@@ -216,10 +220,11 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
             compress_programs: Mutex::new(LruCache::new(compress_cache_size)),
             compress_cache_misses: AtomicUsize::new(0),
             root,
-            merkle_tree,
+            vk_merkle_tree: merkle_tree,
             allowed_vk_map,
             core_shape_config,
             recursion_shape_config,
+            vk_verification,
         }
     }
 
@@ -309,12 +314,19 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         let mut cache = self.recursion_programs.lock().unwrap();
         cache
             .get_or_insert(input.shape(), || {
-                tracing::info!("Proof shape: {:?}", input.shape());
                 let misses = self.recursion_cache_misses.fetch_add(1, Ordering::Relaxed);
-                tracing::debug!("Core cache miss, misses: {}", misses);
+                tracing::debug!("core cache miss, misses: {}", misses);
                 // Get the operations.
                 let builder_span = tracing::debug_span!("build recursion program").entered();
                 let mut builder = Builder::<InnerConfig>::default();
+
+                // // TODO: remove comment or make a test flag.
+                // let dummy_input = SP1RecursionWitnessValues::<CoreSC>::dummy(
+                //     self.core_prover.machine(),
+                //     &input.shape(),
+                // );
+                // let input = dummy_input.read(&mut builder);
+
                 let input = input.read(&mut builder);
                 SP1RecursiveVerifier::verify(&mut builder, self.core_prover.machine(), input);
                 let operations = builder.into_operations();
@@ -346,11 +358,19 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
                 // Get the operations.
                 let builder_span = tracing::debug_span!("build compress program").entered();
                 let mut builder = Builder::<InnerConfig>::default();
+
+                // TODO: remove comment or make a test flag.
+                // let dummy_input = SP1CompressWithVKeyWitnessValues::<CoreSC>::dummy(
+                //     self.compress_prover.machine(),
+                //     &input.shape(),
+                // );
+                // let input = dummy_input.read(&mut builder);
                 let input = input.read(&mut builder);
                 SP1CompressWithVKeyVerifier::verify(
                     &mut builder,
                     self.compress_prover.machine(),
                     input,
+                    self.vk_verification,
                 );
                 let operations = builder.into_operations();
                 builder_span.exit();
@@ -381,6 +401,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
             &mut builder,
             self.compress_prover.machine(),
             input,
+            self.vk_verification,
         );
         let operations = builder.into_operations();
         builder_span.exit();
@@ -506,7 +527,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
                 vks_and_proofs,
                 start_reconstruct_deferred_digest: deferred_digest,
                 is_complete: false,
-                sp1_vk: vk.clone(),
+                sp1_vk_digest: vk.hash_babybear(),
                 end_pc: Val::<InnerSC>::zero(),
                 end_shard: last_proof_pv.shard + BabyBear::one(),
                 end_execution_shard: last_proof_pv.execution_shard,
@@ -1117,7 +1138,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         let proofs = vk_indices
             .iter()
             .map(|index| {
-                let (_, proof) = MerkleTree::open(&self.merkle_tree, *index);
+                let (_, proof) = MerkleTree::open(&self.vk_merkle_tree, *index);
                 proof
             })
             .collect();