feat(sdk)!: configurable aggregation tree shape in SDK (#1570)

Introduce `AggregationTreeConfig` struct for simple configuration of
aggregation tree.

These were already in the provers, but not exposed to be changeable. 

Breaking API changes in SDK:
- `AggStarkProver::new`
- `ContinuationProver::new`
- `Sdk` now internally contains `AggregationTreeConfig`, although
`Sdk::new()` is unchanged and just sets a default

Towards INT-3302

Author: jonathanpwang

Cargo.lock

@@ -63,6 +63,7 @@ async fn main() -> Result<()> {
             app_pk,
             app_committed_exe,
             full_agg_pk,
+            args.agg_tree_config,
         );
         e2e_prover.set_program_name("fib_e2e");
         let _proof = e2e_prover.generate_proof_for_evm(stdin);

benchmarks/prove/src/bin/fib_e2e.rs

@@ -85,6 +85,7 @@ fn main() -> Result<()> {
             app_pk,
             app_committed_exe,
             full_agg_pk,
+            args.agg_tree_config,
         );
         prover.set_program_name("kitchen_sink");
         let stdin = StdIn::default();

benchmarks/prove/src/bin/kitchen_sink.rs

@@ -9,13 +9,12 @@ use openvm_native_compiler::conversion::CompilerOptions;
 use openvm_sdk::{
     commit::commit_app_exe,
     config::{
-        AggConfig, AggStarkConfig, AppConfig, Halo2Config, DEFAULT_APP_LOG_BLOWUP,
-        DEFAULT_INTERNAL_LOG_BLOWUP, DEFAULT_LEAF_LOG_BLOWUP, DEFAULT_ROOT_LOG_BLOWUP,
+        AggConfig, AggStarkConfig, AggregationTreeConfig, AppConfig, Halo2Config,
+        DEFAULT_APP_LOG_BLOWUP, DEFAULT_INTERNAL_LOG_BLOWUP, DEFAULT_LEAF_LOG_BLOWUP,
+        DEFAULT_ROOT_LOG_BLOWUP,
     },
     keygen::{leaf_keygen, AppProvingKey},
-    prover::{
-        vm::local::VmLocalProver, AppProver, LeafProvingController, DEFAULT_NUM_CHILDREN_LEAF,
-    },
+    prover::{vm::local::VmLocalProver, AppProver, LeafProvingController},
     Sdk, StdIn,
 };
 use openvm_stark_backend::utils::metrics_span;
@@ -66,6 +65,10 @@ pub struct BenchmarkCli {
     #[arg(short, long, alias = "max_segment_length")]
     pub max_segment_length: Option<usize>,
 
+    /// Controls the arity (num_children) of the aggregation tree
+    #[command(flatten)]
+    pub agg_tree_config: AggregationTreeConfig,
+
     /// Whether to execute with additional profiling metric collection
     #[arg(long)]
     pub profiling: bool,
@@ -219,7 +222,7 @@ where
         let leaf_prover =
             VmLocalProver::<SC, NativeConfig, E>::new(leaf_vm_pk, app_pk.leaf_committed_exe);
         let leaf_controller = LeafProvingController {
-            num_children: DEFAULT_NUM_CHILDREN_LEAF,
+            num_children: AggregationTreeConfig::default().num_children_leaf,
         };
         leaf_controller.generate_proof(&leaf_prover, &app_proof);
     }

benchmarks/prove/src/util.rs

@@ -5,7 +5,7 @@ use eyre::Result;
 use openvm_native_recursion::halo2::utils::CacheHalo2ParamsReader;
 use openvm_sdk::{
     commit::AppExecutionCommit,
-    config::SdkVmConfig,
+    config::{AggregationTreeConfig, SdkVmConfig},
     fs::{
         read_agg_pk_from_file, read_app_pk_from_file, read_exe_from_file, write_app_proof_to_file,
         write_evm_proof_to_file,
@@ -56,12 +56,15 @@ enum ProveSubCommand {
 
         #[arg(long, action, help = "Path to output proof", default_value = DEFAULT_EVM_PROOF_PATH)]
         output: PathBuf,
+
+        #[command(flatten)]
+        agg_tree_config: AggregationTreeConfig,
     },
 }
 
 impl ProveCmd {
     pub fn run(&self) -> Result<()> {
-        let sdk = Sdk::new();
+        let mut sdk = Sdk::new();
         match &self.command {
             ProveSubCommand::App {
                 app_pk,
@@ -79,7 +82,9 @@ impl ProveCmd {
                 exe,
                 input,
                 output,
+                agg_tree_config,
             } => {
+                sdk.set_agg_tree_config(*agg_tree_config);
                 let params_reader = CacheHalo2ParamsReader::new(DEFAULT_PARAMS_DIR);
                 let (app_pk, committed_exe, input) =
                     Self::prepare_execution(&sdk, app_pk, exe, input)?;

crates/cli/src/commands/prove.rs

@@ -1,3 +1,4 @@
+use clap::Args;
 use openvm_circuit::arch::instructions::program::DEFAULT_MAX_NUM_PUBLIC_VALUES;
 use openvm_continuations::verifier::{
     common::types::VmVerifierPvs, internal::types::InternalVmVerifierPvs,
@@ -15,6 +16,11 @@ pub const DEFAULT_LEAF_LOG_BLOWUP: usize = 1;
 pub const DEFAULT_INTERNAL_LOG_BLOWUP: usize = 2;
 pub const DEFAULT_ROOT_LOG_BLOWUP: usize = 3;
 
+// Aggregation Tree Defaults
+const DEFAULT_NUM_CHILDREN_LEAF: usize = 1;
+const DEFAULT_NUM_CHILDREN_INTERNAL: usize = 4;
+const DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS: usize = 4;
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct AppConfig<VC> {
     #[serde(default)]
@@ -59,6 +65,23 @@ pub struct Halo2Config {
     pub profiling: bool,
 }
 
+#[derive(Clone, Copy, Debug, Serialize, Deserialize, Args)]
+pub struct AggregationTreeConfig {
+    /// Each leaf verifier circuit will aggregate this many App VM proofs.
+    #[arg(long, default_value_t = DEFAULT_NUM_CHILDREN_LEAF)]
+    pub num_children_leaf: usize,
+    /// Each internal verifier circuit will aggregate this many proofs,
+    /// where each proof may be of either leaf or internal verifier (self) circuit.
+    #[arg(long, default_value_t = DEFAULT_NUM_CHILDREN_INTERNAL)]
+    pub num_children_internal: usize,
+    /// Safety threshold: how many times to do 1-to-1 aggregation of the "last" internal
+    /// verifier proof before it is small enough for the root verifier circuit.
+    /// Note: almost always no wrapping is needed.
+    #[arg(long, default_value_t = DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS)]
+    pub max_internal_wrapper_layers: usize,
+    // root currently always has 1 child for now
+}
+
 impl<VC> AppConfig<VC> {
     pub fn new(app_fri_params: FriParameters, app_vm_config: VC) -> Self {
         Self {
@@ -187,3 +210,13 @@ impl AggStarkConfig {
         config
     }
 }
+
+impl Default for AggregationTreeConfig {
+    fn default() -> Self {
+        Self {
+            num_children_leaf: DEFAULT_NUM_CHILDREN_LEAF,
+            num_children_internal: DEFAULT_NUM_CHILDREN_INTERNAL,
+            max_internal_wrapper_layers: DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS,
+        }
+    }
+}

crates/sdk/src/config/mod.rs

@@ -10,7 +10,7 @@ use std::{
 use alloy_primitives::{Bytes, FixedBytes};
 use alloy_sol_types::{sol, SolCall, SolValue};
 use commit::commit_app_exe;
-use config::AppConfig;
+use config::{AggregationTreeConfig, AppConfig};
 use eyre::{Context, Result};
 use keygen::{AppProvingKey, AppVerifyingKey};
 use openvm_build::{
@@ -101,12 +101,14 @@ pub struct VerifiedContinuationVmPayload {
 }
 
 pub struct GenericSdk<E: StarkFriEngine<SC>> {
+    agg_tree_config: AggregationTreeConfig,
     _phantom: PhantomData<E>,
 }
 
 impl<E: StarkFriEngine<SC>> Default for GenericSdk<E> {
     fn default() -> Self {
         Self {
+            agg_tree_config: AggregationTreeConfig::default(),
             _phantom: PhantomData,
         }
     }
@@ -119,6 +121,14 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         Self::default()
     }
 
+    pub fn agg_tree_config(&self) -> &AggregationTreeConfig {
+        &self.agg_tree_config
+    }
+
+    pub fn set_agg_tree_config(&mut self, agg_tree_config: AggregationTreeConfig) {
+        self.agg_tree_config = agg_tree_config;
+    }
+
     pub fn build<P: AsRef<Path>>(
         &self,
         guest_opts: GuestOptions,
@@ -264,7 +274,8 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let stark_prover = StarkProver::<VC, E>::new(app_pk, app_exe, agg_stark_pk);
+        let stark_prover =
+            StarkProver::<VC, E>::new(app_pk, app_exe, agg_stark_pk, self.agg_tree_config);
         let proof = stark_prover.generate_root_verifier_input(inputs);
         Ok(proof)
     }
@@ -281,7 +292,8 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let e2e_prover = ContinuationProver::<VC, E>::new(reader, app_pk, app_exe, agg_pk);
+        let e2e_prover =
+            ContinuationProver::<VC, E>::new(reader, app_pk, app_exe, agg_pk, self.agg_tree_config);
         let proof = e2e_prover.generate_proof_for_evm(inputs);
         Ok(proof)
     }

crates/sdk/src/lib.rs

@@ -11,6 +11,7 @@ use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::proof::Proo
 use tracing::info_span;
 
 use crate::{
+    config::AggregationTreeConfig,
     keygen::AggStarkProvingKey,
     prover::{
         vm::{local::VmLocalProver, SingleSegmentVmProver},
@@ -19,10 +20,6 @@ use crate::{
     NonRootCommittedExe, RootSC, F, SC,
 };
 
-pub const DEFAULT_NUM_CHILDREN_LEAF: usize = 1;
-const DEFAULT_NUM_CHILDREN_INTERNAL: usize = 2;
-const DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS: usize = 4;
-
 pub struct AggStarkProver<E: StarkFriEngine<SC>> {
     leaf_prover: VmLocalProver<SC, NativeConfig, E>,
     leaf_controller: LeafProvingController,
@@ -43,11 +40,12 @@ impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
     pub fn new(
         agg_stark_pk: AggStarkProvingKey,
         leaf_committed_exe: Arc<NonRootCommittedExe>,
+        tree_config: AggregationTreeConfig,
     ) -> Self {
         let leaf_prover =
             VmLocalProver::<SC, NativeConfig, E>::new(agg_stark_pk.leaf_vm_pk, leaf_committed_exe);
         let leaf_controller = LeafProvingController {
-            num_children: DEFAULT_NUM_CHILDREN_LEAF,
+            num_children: tree_config.num_children_leaf,
         };
         let internal_prover = VmLocalProver::<SC, NativeConfig, E>::new(
             agg_stark_pk.internal_vm_pk,
@@ -59,8 +57,8 @@ impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
             leaf_controller,
             internal_prover,
             root_prover,
-            num_children_internal: DEFAULT_NUM_CHILDREN_INTERNAL,
-            max_internal_wrapper_layers: DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS,
+            num_children_internal: tree_config.num_children_internal,
+            max_internal_wrapper_layers: tree_config.max_internal_wrapper_layers,
         }
     }
 

crates/sdk/src/prover/agg.rs

@@ -3,7 +3,9 @@ use std::sync::Arc;
 use openvm_circuit::arch::VmConfig;
 use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::Chip};
 
-use crate::{keygen::AppProvingKey, stdin::StdIn, NonRootCommittedExe, F, SC};
+use crate::{
+    config::AggregationTreeConfig, keygen::AppProvingKey, stdin::StdIn, NonRootCommittedExe, F, SC,
+};
 
 mod agg;
 pub use agg::*;
@@ -25,8 +27,8 @@ pub use stark::*;
 use crate::{keygen::AggProvingKey, prover::halo2::Halo2Prover, types::EvmProof};
 
 pub struct ContinuationProver<VC, E: StarkFriEngine<SC>> {
-    stark_prover: StarkProver<VC, E>,
-    halo2_prover: Halo2Prover,
+    pub stark_prover: StarkProver<VC, E>,
+    pub halo2_prover: Halo2Prover,
 }
 
 impl<VC, E: StarkFriEngine<SC>> ContinuationProver<VC, E> {
@@ -35,6 +37,7 @@ impl<VC, E: StarkFriEngine<SC>> ContinuationProver<VC, E> {
         app_pk: Arc<AppProvingKey<VC>>,
         app_committed_exe: Arc<NonRootCommittedExe>,
         agg_pk: AggProvingKey,
+        agg_tree_config: AggregationTreeConfig,
     ) -> Self
     where
         VC: VmConfig<F>,
@@ -43,7 +46,8 @@ impl<VC, E: StarkFriEngine<SC>> ContinuationProver<VC, E> {
             agg_stark_pk,
             halo2_pk,
         } = agg_pk;
-        let stark_prover = StarkProver::new(app_pk, app_committed_exe, agg_stark_pk);
+        let stark_prover =
+            StarkProver::new(app_pk, app_committed_exe, agg_stark_pk, agg_tree_config);
         Self {
             stark_prover,
             halo2_prover: Halo2Prover::new(reader, halo2_pk),

crates/sdk/src/prover/mod.rs

@@ -6,20 +6,22 @@ use openvm_stark_backend::{proof::Proof, Chip};
 use openvm_stark_sdk::engine::StarkFriEngine;
 
 use crate::{
+    config::AggregationTreeConfig,
     keygen::{AggStarkProvingKey, AppProvingKey},
     prover::{agg::AggStarkProver, app::AppProver},
     NonRootCommittedExe, RootSC, StdIn, F, SC,
 };
 
 pub struct StarkProver<VC, E: StarkFriEngine<SC>> {
-    app_prover: AppProver<VC, E>,
-    agg_prover: AggStarkProver<E>,
+    pub app_prover: AppProver<VC, E>,
+    pub agg_prover: AggStarkProver<E>,
 }
 impl<VC, E: StarkFriEngine<SC>> StarkProver<VC, E> {
     pub fn new(
         app_pk: Arc<AppProvingKey<VC>>,
         app_committed_exe: Arc<NonRootCommittedExe>,
         agg_stark_pk: AggStarkProvingKey,
+        agg_tree_config: AggregationTreeConfig,
     ) -> Self
     where
         VC: VmConfig<F>,
@@ -36,7 +38,11 @@ impl<VC, E: StarkFriEngine<SC>> StarkProver<VC, E> {
 
         Self {
             app_prover: AppProver::new(app_pk.app_vm_pk.clone(), app_committed_exe),
-            agg_prover: AggStarkProver::new(agg_stark_pk, app_pk.leaf_committed_exe.clone()),
+            agg_prover: AggStarkProver::new(
+                agg_stark_pk,
+                app_pk.leaf_committed_exe.clone(),
+                agg_tree_config,
+            ),
         }
     }
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {