back to original scripts

succinctlabs · Oct 8, 2024 · dd6b43e · dd6b43e
1 parent a4abe54
commit dd6b43e
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 55 deletions.
diff --git a/crates/prover/scripts/build_compress_vks.rs b/crates/prover/scripts/build_compress_vks.rs
@@ -1,10 +1,6 @@
-use std::{
-    collections::{BTreeMap, BTreeSet},
-    path::PathBuf,
-};
+use std::path::PathBuf;
 
 use clap::Parser;
-use p3_baby_bear::BabyBear;
 use sp1_core_machine::utils::setup_logger;
 use sp1_prover::{
     components::DefaultProverComponents, shapes::build_vk_map_to_file, REDUCE_BATCH_SIZE,
@@ -41,34 +37,14 @@ fn main() {
     let range_start = args.start;
     let range_end = args.end;
 
-    let old_map_file = std::fs::File::open(build_dir.join("allowed_vk_map.bin")).unwrap();
-    let old_map: BTreeMap<[BabyBear; 8], usize> = bincode::deserialize_from(old_map_file).unwrap();
-
-    println!("old_set: {:?}", old_map.len());
-
-    let shrink_set_file = std::fs::File::open(build_dir.join("shrink/vk_map.bin")).unwrap();
-    let shrink_set: BTreeMap<[BabyBear; 8], usize> =
-        bincode::deserialize_from(shrink_set_file).unwrap();
-
-    println!("shrink_set: {:?}", shrink_set.len());
-
-    let new_set = old_map.into_keys().chain(shrink_set.into_keys()).collect::<BTreeSet<_>>();
-
-    println!("new_set: {:?}", new_set.len());
-
-    let new_map =
-        new_set.into_iter().enumerate().map(|(i, vk)| (vk, i)).collect::<BTreeMap<_, _>>();
-
-    let mut file = std::fs::File::create(build_dir.join("allowed_vk_map_new.bin")).unwrap();
-    bincode::serialize_into(&mut file, &new_map).unwrap();
-    // build_vk_map_to_file::<DefaultProverComponents>(
-    //     build_dir,
-    //     reduce_batch_size,
-    //     dummy,
-    //     num_compiler_workers,
-    //     num_setup_workers,
-    //     range_start,
-    //     range_end,
-    // )
-    // .unwrap();
+    build_vk_map_to_file::<DefaultProverComponents>(
+        build_dir,
+        reduce_batch_size,
+        dummy,
+        num_compiler_workers,
+        num_setup_workers,
+        range_start,
+        range_end,
+    )
+    .unwrap();
 }
diff --git a/crates/prover/src/shapes.rs b/crates/prover/src/shapes.rs
@@ -68,8 +68,7 @@ pub fn build_vk_map<C: SP1ProverComponents>(
         )
         .into_keys()
         .collect::<BTreeSet<_>>();
-        // let height = dummy_set.len().next_power_of_two().ilog2() as usize;
-        let height = 21;
+        let height = dummy_set.len().next_power_of_two().ilog2() as usize;
         (dummy_set, vec![], height)
     } else {
         let (vk_tx, vk_rx) = std::sync::mpsc::channel();
@@ -87,8 +86,7 @@ pub fn build_vk_map<C: SP1ProverComponents>(
                 .collect::<BTreeSet<_>>();
         let num_shapes = all_shapes.len();
 
-        // let height = num_shapes.next_power_of_two().ilog2() as usize;
-        let height = 21;
+        let height = num_shapes.next_power_of_two().ilog2() as usize;
         let chunk_size = indices_set.as_ref().map(|indices| indices.len()).unwrap_or(num_shapes);
 
         std::thread::scope(|s| {
@@ -214,22 +212,22 @@ impl SP1ProofShape {
         recursion_shape_config: &'a RecursionShapeConfig<BabyBear, CompressAir<BabyBear>>,
         reduce_batch_size: usize,
     ) -> impl Iterator<Item = Self> + 'a {
-        // core_shape_config
-        //     .generate_all_allowed_shapes()
-        //     .map(Self::Recursion)
-        //     .chain((1..=reduce_batch_size).flat_map(|batch_size| {
-        //         recursion_shape_config.get_all_shape_combinations(batch_size).map(Self::Compress)
-        //     }))
-        //     .chain(
-        //         recursion_shape_config
-        //             .get_all_shape_combinations(1)
-        //             .map(|mut x| Self::Deferred(x.pop().unwrap())),
-        //     )
-        //     .chain(
-        recursion_shape_config
-            .get_all_shape_combinations(1)
-            .map(|mut x| Self::Shrink(x.pop().unwrap()))
-        // )
+        core_shape_config
+            .generate_all_allowed_shapes()
+            .map(Self::Recursion)
+            .chain((1..=reduce_batch_size).flat_map(|batch_size| {
+                recursion_shape_config.get_all_shape_combinations(batch_size).map(Self::Compress)
+            }))
+            .chain(
+                recursion_shape_config
+                    .get_all_shape_combinations(1)
+                    .map(|mut x| Self::Deferred(x.pop().unwrap())),
+            )
+            .chain(
+                recursion_shape_config
+                    .get_all_shape_combinations(1)
+                    .map(|mut x| Self::Shrink(x.pop().unwrap())),
+            )
     }
 
     pub fn dummy_vk_map<'a>(