feat(cli): add SegmentationArgs to prove command

benchmarks/prove/src/util.rs

@@ -66,7 +66,7 @@ pub struct BenchmarkCli {
     #[arg(long, alias = "max_segment_length")]
     pub max_segment_length: Option<u32>,
 
-    /// Max cells per chip in segment for continuations
+    /// Total cells used in all chips in segment for continuations
     #[arg(long)]
     pub segment_max_cells: Option<usize>,
 

crates/cli/src/commands/prove.rs

@@ -2,9 +2,14 @@ use std::path::PathBuf;
 
 use clap::Parser;
 use eyre::Result;
-use openvm_circuit::arch::instructions::exe::VmExe;
+use openvm_circuit::arch::{
+    execution_mode::metered::segment_ctx::{
+        SegmentationLimits, DEFAULT_MAX_CELLS, DEFAULT_MAX_TRACE_HEIGHT,
+    },
+    instructions::exe::VmExe,
+};
 use openvm_sdk::{
-    config::{AggregationTreeConfig, SdkVmConfig},
+    config::{AggregationTreeConfig, AppConfig, SdkVmConfig},
     fs::{encode_to_file, read_object_from_file, write_to_file_json},
     keygen::AppProvingKey,
     types::VersionedVmStarkProof,
@@ -52,6 +57,9 @@ enum ProveSubCommand {
 
         #[command(flatten)]
         cargo_args: RunCargoArgs,
+
+        #[command(flatten)]
+        segmentation_args: SegmentationArgs,
     },
     Stark {
         #[arg(
@@ -76,6 +84,9 @@ enum ProveSubCommand {
         #[command(flatten)]
         cargo_args: RunCargoArgs,
 
+        #[command(flatten)]
+        segmentation_args: SegmentationArgs,
+
         #[command(flatten)]
         agg_tree_config: AggregationTreeConfig,
     },
@@ -103,11 +114,26 @@ enum ProveSubCommand {
         #[command(flatten)]
         cargo_args: RunCargoArgs,
 
+        #[command(flatten)]
+        segmentation_args: SegmentationArgs,
+
         #[command(flatten)]
         agg_tree_config: AggregationTreeConfig,
     },
 }
 
+#[derive(Clone, Copy, Parser)]
+pub struct SegmentationArgs {
+    /// Trace height threshold per chip for triggering segmentation for continuations in the app
+    /// proof. Note that these thresholds are not absolute limits.
+    #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT)]
+    pub segment_max_height: u32,
+    /// Total cells used across all chips for triggering segmentation for continuations in the app
+    /// proof. Note that these thresholds are not absolute limits.
+    #[arg(long, default_value_t = DEFAULT_MAX_CELLS)]
+    pub segment_max_cells: usize,
+}
+
 impl ProveCmd {
     pub fn run(&self) -> Result<()> {
         match &self.command {
@@ -116,9 +142,11 @@ impl ProveCmd {
                 proof,
                 run_args,
                 cargo_args,
+                segmentation_args,
             } => {
                 let app_pk = load_app_pk(app_pk, cargo_args)?;
-                let sdk = Sdk::new(app_pk.app_config())?.with_app_pk(app_pk);
+                let app_config = get_app_config(&app_pk, segmentation_args);
+                let sdk = Sdk::new(app_config)?.with_app_pk(app_pk);
                 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
 
                 let app_proof = sdk
@@ -141,6 +169,7 @@ impl ProveCmd {
                 proof,
                 run_args,
                 cargo_args,
+                segmentation_args,
                 agg_tree_config,
             } => {
                 let app_pk = load_app_pk(app_pk, cargo_args)?;
@@ -149,7 +178,8 @@ impl ProveCmd {
                 let agg_pk = read_object_from_file(default_agg_stark_pk_path()).map_err(|e| {
                     eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
                 })?;
-                let sdk = Sdk::new(app_pk.app_config())?
+                let app_config = get_app_config(&app_pk, segmentation_args);
+                let sdk = Sdk::new(app_config)?
                     .with_agg_tree_config(*agg_tree_config)
                     .with_app_pk(app_pk)
                     .with_agg_pk(agg_pk);
@@ -178,6 +208,7 @@ impl ProveCmd {
                 proof,
                 run_args,
                 cargo_args,
+                segmentation_args,
                 agg_tree_config,
             } => {
                 let app_pk = load_app_pk(app_pk, cargo_args)?;
@@ -187,7 +218,8 @@ impl ProveCmd {
                 let (agg_pk, halo2_pk) = read_default_agg_and_halo2_pk().map_err(|e| {
                     eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
                 })?;
-                let sdk = Sdk::new(app_pk.app_config())?
+                let app_config = get_app_config(&app_pk, segmentation_args);
+                let sdk = Sdk::new(app_config)?
                     .with_agg_tree_config(*agg_tree_config)
                     .with_app_pk(app_pk)
                     .with_agg_pk(agg_pk)
@@ -253,3 +285,26 @@ pub(crate) fn load_or_build_exe(
         exe_path.file_stem().unwrap().to_string_lossy().into_owned(),
     ))
 }
+
+fn get_app_config(
+    app_pk: &AppProvingKey<SdkVmConfig>,
+    segmentation_args: &SegmentationArgs,
+) -> AppConfig<SdkVmConfig> {
+    let mut app_config = app_pk.app_config();
+    app_config
+        .app_vm_config
+        .system
+        .config
+        .set_segmentation_limits((*segmentation_args).into());
+    app_config
+}
+
+impl From<SegmentationArgs> for SegmentationLimits {
+    fn from(args: SegmentationArgs) -> Self {
+        SegmentationLimits {
+            max_trace_height: args.segment_max_height,
+            max_cells: args.segment_max_cells,
+            ..Default::default()
+        }
+    }
+}

crates/vm/src/arch/execution_mode/metered/segment_ctx.rs

@@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize};
 
 pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
 
-pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << 23;
-pub const DEFAULT_MAX_CELLS: usize = 2_000_000_000; // 2B
+pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << 22;
+pub const DEFAULT_MAX_CELLS: usize = 1_200_000_000; // 1.2B
 const DEFAULT_MAX_INTERACTIONS: usize = BabyBear::ORDER_U32 as usize;
 
 #[derive(derive_new::new, Clone, Debug, Serialize, Deserialize)]