From 33d9e7fd57795238e42e273ede16e9e4c63ada78 Mon Sep 17 00:00:00 2001
From: erabinov <erabinovich92@gmail.com>
Date: Tue, 5 Nov 2024 12:29:42 -0800
Subject: [PATCH] working to do no precompile shapes too
---
.../find_minimal_large_recursion_shape.rs | 22 +++++++++++++++++--
crates/prover/src/shapes.rs | 14 ++++++++----
2 files changed, 30 insertions(+), 6 deletions(-)
diff --git a/crates/prover/scripts/find_minimal_large_recursion_shape.rs b/crates/prover/scripts/find_minimal_large_recursion_shape.rs
index 7f4b74becb..87a21143a7 100644
--- 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
index f045dfd69d..2fc5d24fae 100644
--- 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(),