[feat] support custom segmentation strategy

benchmarks/src/bin/fib_e2e.rs

@@ -24,10 +24,7 @@ async fn main() -> Result<()> {
     // Must be larger than RangeTupleCheckerAir.height == 524288
     let max_segment_length = args.max_segment_length.unwrap_or(1_000_000);
 
-    let app_config = args.app_config(Rv32ImConfig::with_public_values_and_segment_len(
-        NUM_PUBLIC_VALUES,
-        max_segment_length,
-    ));
+    let app_config = args.app_config(Rv32ImConfig::with_public_values(NUM_PUBLIC_VALUES));
     let agg_config = args.agg_config();
 
     let sdk = Sdk;
@@ -58,6 +55,7 @@ async fn main() -> Result<()> {
     run_with_metric_collection("OUTPUT_PATH", || {
         let mut e2e_prover =
             ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk);
+        e2e_prover.set_max_segment_len(max_segment_length);
         e2e_prover.set_program_name("fib_e2e");
         let _proof = e2e_prover.generate_proof_for_evm(stdin);
     });

crates/sdk/examples/sdk_evm.rs

@@ -105,6 +105,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         app_pk,
         app_committed_exe,
         agg_pk,
+        None,
         stdin,
     )?;
 

crates/sdk/src/lib.rs

@@ -190,13 +190,17 @@ impl Sdk {
         app_pk: Arc<AppProvingKey<VC>>,
         app_exe: Arc<NonRootCommittedExe>,
         agg_pk: AggProvingKey,
+        max_segment_len: Option<usize>,
         inputs: StdIn,
     ) -> Result<EvmProof>
     where
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let e2e_prover = ContinuationProver::new(reader, app_pk, app_exe, agg_pk);
+        let mut e2e_prover = ContinuationProver::new(reader, app_pk, app_exe, agg_pk);
+        if let Some(max_segment_len) = max_segment_len {
+            e2e_prover.set_max_segment_len(max_segment_len);
+        }
         let proof = e2e_prover.generate_proof_for_evm(inputs);
         Ok(proof)
     }

crates/sdk/src/prover/app.rs

@@ -17,8 +17,7 @@ use crate::{
 #[derive(Getters)]
 pub struct AppProver<VC> {
     pub program_name: Option<String>,
-    #[getset(get = "pub")]
-    app_prover: VmLocalProver<SC, VC, BabyBearPoseidon2Engine>,
+    pub app_prover: VmLocalProver<SC, VC, BabyBearPoseidon2Engine>,
 }
 
 impl<VC> AppProver<VC> {

crates/sdk/src/prover/mod.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use openvm_circuit::arch::VmConfig;
+use openvm_circuit::arch::{DefaultSegmentationStrategy, VmConfig};
 use openvm_native_recursion::halo2::EvmProof;
 use openvm_stark_sdk::openvm_stark_backend::Chip;
 
@@ -29,8 +29,8 @@ use crate::{
 };
 
 pub struct ContinuationProver<VC> {
-    stark_prover: StarkProver<VC>,
-    halo2_prover: Halo2Prover,
+    pub stark_prover: StarkProver<VC>,
+    pub halo2_prover: Halo2Prover,
 }
 
 impl<VC> ContinuationProver<VC> {
@@ -59,6 +59,16 @@ impl<VC> ContinuationProver<VC> {
         self
     }
 
+    pub fn set_max_segment_len(&mut self, max_segment_len: usize) -> &mut Self {
+        self.stark_prover
+            .app_prover
+            .app_prover
+            .set_custom_segmentation_strategy(Arc::new(
+                DefaultSegmentationStrategy::new_with_max_segment_len(max_segment_len),
+            ));
+        self
+    }
+
     pub fn generate_proof_for_evm(&self, input: StdIn) -> EvmProof
     where
         VC: VmConfig<F>,

crates/sdk/src/prover/stark.rs

@@ -10,8 +10,8 @@ use crate::{
 };
 
 pub struct StarkProver<VC> {
-    app_prover: AppProver<VC>,
-    agg_prover: AggStarkProver,
+    pub app_prover: AppProver<VC>,
+    pub agg_prover: AggStarkProver,
 }
 impl<VC> StarkProver<VC> {
     pub fn new(

crates/sdk/src/prover/vm/local.rs

@@ -3,8 +3,8 @@ use std::{marker::PhantomData, sync::Arc};
 use async_trait::async_trait;
 use openvm_circuit::{
     arch::{
-        hasher::poseidon2::vm_poseidon2_hasher, SingleSegmentVmExecutor, Streams, VirtualMachine,
-        VmComplexTraceHeights, VmConfig,
+        hasher::poseidon2::vm_poseidon2_hasher, segment::SegmentationStrategy,
+        SingleSegmentVmExecutor, Streams, VirtualMachine, VmComplexTraceHeights, VmConfig,
     },
     system::{memory::tree::public_values::UserPublicValuesProof, program::trace::VmCommittedExe},
 };
@@ -25,6 +25,7 @@ pub struct VmLocalProver<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> {
     pub pk: Arc<VmProvingKey<SC, VC>>,
     pub committed_exe: Arc<VmCommittedExe<SC>>,
     overridden_heights: Option<VmComplexTraceHeights>,
+    segmentation_strategy: Option<Arc<dyn SegmentationStrategy>>,
     _marker: PhantomData<E>,
 }
 
@@ -34,6 +35,7 @@ impl<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> VmLocalProver<SC, VC, E>
             pk,
             committed_exe,
             overridden_heights: None,
+            segmentation_strategy: None,
             _marker: PhantomData,
         }
     }
@@ -47,6 +49,7 @@ impl<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> VmLocalProver<SC, VC, E>
             pk,
             committed_exe,
             overridden_heights,
+            segmentation_strategy: None,
             _marker: PhantomData,
         }
     }
@@ -55,6 +58,10 @@ impl<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> VmLocalProver<SC, VC, E>
         self.overridden_heights = Some(overridden_heights);
     }
 
+    pub fn set_custom_segmentation_strategy(&mut self, strategy: Arc<dyn SegmentationStrategy>) {
+        self.segmentation_strategy = Some(strategy);
+    }
+
     pub fn vm_config(&self) -> &VC {
         &self.pk.vm_config
     }
@@ -74,11 +81,15 @@ where
     fn prove(&self, input: impl Into<Streams<Val<SC>>>) -> ContinuationVmProof<SC> {
         assert!(self.pk.vm_config.system().continuation_enabled);
         let e = E::new(self.pk.fri_params);
-        let vm = VirtualMachine::new_with_overridden_trace_heights(
+        let mut vm = VirtualMachine::new_with_overridden_trace_heights(
             e,
             self.pk.vm_config.clone(),
             self.overridden_heights.clone(),
         );
+        if let Some(segmentation_strategy) = self.segmentation_strategy.clone() {
+            vm.executor
+                .set_custom_segmentation_strategy(segmentation_strategy);
+        }
         let results = vm
             .execute_and_generate_with_cached_program(self.committed_exe.clone(), input)
             .unwrap();

crates/sdk/tests/integration_test.rs

@@ -3,8 +3,8 @@ use std::{borrow::Borrow, path::PathBuf, sync::Arc};
 use openvm_build::GuestOptions;
 use openvm_circuit::{
     arch::{
-        hasher::poseidon2::vm_poseidon2_hasher, ExecutionError, SingleSegmentVmExecutor,
-        SystemConfig, VmConfig, VmExecutor,
+        hasher::poseidon2::vm_poseidon2_hasher, DefaultSegmentationStrategy, ExecutionError,
+        SingleSegmentVmExecutor, SystemConfig, VmConfig, VmExecutor,
     },
     system::{memory::tree::public_values::UserPublicValuesProof, program::trace::VmCommittedExe},
 };
@@ -127,7 +127,6 @@ fn small_test_app_config(app_log_blowup: usize) -> AppConfig<NativeConfig> {
             .into(),
         app_vm_config: NativeConfig::new(
             SystemConfig::default()
-                .with_max_segment_len(200)
                 .with_continuations()
                 .with_public_values(NUM_PUB_VALUES),
             Native,
@@ -154,7 +153,11 @@ fn test_public_values_and_leaf_verification() {
     let leaf_committed_exe = app_pk.leaf_committed_exe.clone();
 
     let app_engine = BabyBearPoseidon2Engine::new(app_pk.app_vm_pk.fri_params);
-    let app_vm = VmExecutor::new(app_pk.app_vm_pk.vm_config.clone());
+
+    let segmentation_strategy =
+        Arc::new(DefaultSegmentationStrategy::new_with_max_segment_len(200));
+    let mut app_vm = VmExecutor::new(app_pk.app_vm_pk.vm_config.clone());
+    app_vm.set_custom_segmentation_strategy(segmentation_strategy);
     let app_vm_result = app_vm
         .execute_and_generate_with_cached_program(app_committed_exe.clone(), vec![])
         .unwrap();
@@ -350,6 +353,7 @@ fn test_static_verifier_custom_pv_handler() {
             Arc::new(app_pk),
             app_committed_exe,
             agg_pk,
+            Some(200),
             StdIn::default(),
         )
         .unwrap();
@@ -379,6 +383,7 @@ fn test_e2e_proof_generation_and_verification() {
             Arc::new(app_pk),
             app_committed_exe_for_test(app_log_blowup),
             agg_pk,
+            Some(200),
             StdIn::default(),
         )
         .unwrap();

crates/vm/src/arch/config.rs

@@ -18,7 +18,6 @@ use super::{
 };
 use crate::system::memory::BOUNDARY_AIR_OFFSET;
 
-const DEFAULT_MAX_SEGMENT_LEN: usize = (1 << 22) - 100;
 // sbox is decomposed to have this max degree for Poseidon2. We set to 3 so quotient_degree = 2
 // allows log_blowup = 1
 const DEFAULT_POSEIDON2_MAX_CONSTRAINT_DEGREE: usize = 3;
@@ -86,8 +85,6 @@ pub struct SystemConfig {
     /// cannot read public values directly, but they can decommit the public values from the memory
     /// merkle root.
     pub num_public_values: usize,
-    /// When continuations are enabled, a heuristic used to determine when to segment execution.
-    pub max_segment_len: usize,
     /// Whether to collect detailed profiling metrics.
     /// **Warning**: this slows down the runtime.
     pub profiling: bool,
@@ -110,7 +107,6 @@ impl SystemConfig {
             continuation_enabled: false,
             memory_config,
             num_public_values,
-            max_segment_len: DEFAULT_MAX_SEGMENT_LEN,
             profiling: false,
         }
     }
@@ -135,11 +131,6 @@ impl SystemConfig {
         self
     }
 
-    pub fn with_max_segment_len(mut self, max_segment_len: usize) -> Self {
-        self.max_segment_len = max_segment_len;
-        self
-    }
-
     pub fn with_profiling(mut self) -> Self {
         self.profiling = true;
         self