working to do no precompile shapes too

succinctlabs · Nov 5, 2024 · 33d9e7f · 33d9e7f
1 parent 8ff3fa7
commit 33d9e7f
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/crates/prover/scripts/find_minimal_large_recursion_shape.rs b/crates/prover/scripts/find_minimal_large_recursion_shape.rs
@@ -46,7 +46,7 @@ fn main() {
 
     prover.recursion_shape_config = Some(RecursionShapeConfig::from_hash_map(&candidate));
 
-    assert!(check_shapes(reduce_batch_size, true, num_compiler_workers, &prover,));
+    assert!(check_shapes(reduce_batch_size, true, false, num_compiler_workers, &prover,));
 
     let mut answer = candidate.clone();
 
@@ -58,12 +58,29 @@ fn main() {
                 new_val -= 1;
                 answer.insert(key.clone(), new_val);
                 prover.recursion_shape_config = Some(RecursionShapeConfig::from_hash_map(&answer));
-                done = !check_shapes(reduce_batch_size, true, num_compiler_workers, &prover);
+                done = !check_shapes(reduce_batch_size, true, false, num_compiler_workers, &prover);
             }
             answer.insert(key.clone(), new_val + 1);
         }
     }
 
+    let mut no_precompile_answer = answer.clone();
+
+    for (key, value) in answer.iter() {
+        if key != "PublicValues" {
+            let mut done = false;
+            let mut new_val = *value;
+            while !done {
+                new_val -= 1;
+                no_precompile_answer.insert(key.clone(), new_val);
+                prover.recursion_shape_config =
+                    Some(RecursionShapeConfig::from_hash_map(&no_precompile_answer));
+                done = !check_shapes(reduce_batch_size, true, true, num_compiler_workers, &prover);
+            }
+            no_precompile_answer.insert(key.clone(), new_val + 1);
+        }
+    }
+
     // Check that the shrink shape is compatible with the final compress shape choice.
     let mut shrink_shape = ShrinkAir::<BabyBear>::shrink_shape().inner;
 
@@ -109,5 +126,6 @@ fn main() {
     }
 
     println!("Final compress shape: {:?}", answer);
+    println!("Final compress shape with no precompiles: {:?}", no_precompile_answer);
     println!("Final shrink shape: {:?}", shrink_shape);
 }
diff --git a/crates/prover/src/shapes.rs b/crates/prover/src/shapes.rs
@@ -61,6 +61,7 @@ pub enum VkBuildError {
 pub fn check_shapes<C: SP1ProverComponents>(
     reduce_batch_size: usize,
     only_maximal: bool,
+    no_precompiles: bool,
     num_compiler_workers: usize,
     prover: &SP1Prover<C>,
 ) -> bool {
@@ -78,6 +79,7 @@ pub fn check_shapes<C: SP1ProverComponents>(
             core_shape_config,
             recursion_shape_config,
             reduce_batch_size,
+            no_precompiles,
         )
         .collect::<BTreeSet<SP1ProofShape>>()
     } else {
@@ -173,6 +175,7 @@ pub fn build_vk_map<C: SP1ProverComponents>(
                 core_shape_config,
                 recursion_shape_config,
                 reduce_batch_size,
+                false,
             )
             .collect::<BTreeSet<_>>()
         } else {
@@ -340,11 +343,14 @@ impl SP1ProofShape {
         core_shape_config: &'a CoreShapeConfig<BabyBear>,
         recursion_shape_config: &'a RecursionShapeConfig<BabyBear, CompressAir<BabyBear>>,
         reduce_batch_size: usize,
+        no_precompiles: bool,
     ) -> impl Iterator<Item = Self> + 'a {
-        core_shape_config
-            // .maximal_core_plus_precompile_shapes()
-            .maximal_core_shapes()
-            .into_iter()
+        let core_shape_iter = if no_precompiles {
+            core_shape_config.maximal_core_shapes().into_iter()
+        } else {
+            core_shape_config.maximal_core_plus_precompile_shapes().into_iter()
+        };
+        core_shape_iter
             .map(|core_shape| {
                 Self::Recursion(ProofShape {
                     chip_information: core_shape.inner.into_iter().collect(),