fix(new-execution): replace segmentation strategy with segmentation limits (#1901)

Resolves INT-4316

Author: shuklaayush

benchmarks/execute/benches/execute.rs

@@ -245,16 +245,10 @@ fn benchmark_execute_metered(bencher: Bencher, program: &str) {
             let vm_config = ExecuteConfig::default();
             let exe = load_program_executable(program).expect("Failed to load program executable");
 
-            let segmentation_strategy = &vm_config.as_ref().segmentation_strategy;
-
             let (ctx, executor_idx_to_air_idx) = metering_setup();
-            let ctx = ctx
-                .clone()
-                .with_max_trace_height(segmentation_strategy.max_trace_height() as u32)
-                .with_max_cells(segmentation_strategy.max_cells());
             let interpreter = InterpretedInstance::new(vm_config, exe).unwrap();
 
-            (interpreter, vec![], ctx, executor_idx_to_air_idx)
+            (interpreter, vec![], ctx.clone(), executor_idx_to_air_idx)
         })
         .bench_values(|(interpreter, input, ctx, executor_idx_to_air_idx)| {
             interpreter

benchmarks/prove/src/util.rs

@@ -1,12 +1,12 @@
-use std::{path::PathBuf, sync::Arc};
+use std::path::PathBuf;
 
 use clap::{command, Parser};
 use eyre::Result;
 use openvm_benchmarks_utils::{build_elf, get_programs_dir};
 use openvm_circuit::arch::{
-    instructions::exe::VmExe, verify_single, DefaultSegmentationStrategy, InsExecutorE1,
-    InsExecutorE2, InstructionExecutor, MatrixRecordArena, SystemConfig, VmBuilder, VmConfig,
-    VmExecutionConfig,
+    execution_mode::metered::segment_ctx::SegmentationLimits, instructions::exe::VmExe,
+    verify_single, InsExecutorE1, InsExecutorE2, InstructionExecutor, MatrixRecordArena,
+    SystemConfig, VmBuilder, VmConfig, VmExecutionConfig,
 };
 use openvm_native_circuit::{NativeConfig, NativeCpuBuilder};
 use openvm_native_compiler::conversion::CompilerOptions;
@@ -86,9 +86,9 @@ impl BenchmarkCli {
 
         app_vm_config.as_mut().profiling = self.profiling;
         if let Some(max_segment_length) = self.max_segment_length {
-            app_vm_config.as_mut().set_segmentation_strategy(Arc::new(
-                DefaultSegmentationStrategy::new_with_max_segment_len(max_segment_length),
-            ));
+            app_vm_config.as_mut().set_segmentation_limits(
+                SegmentationLimits::default().with_max_trace_height(max_segment_length as u32),
+            );
         }
         AppConfig {
             app_fri_params: FriParameters::standard_with_100_bits_conjectured_security(

crates/sdk/guest/fib/src/main.rs

@@ -3,16 +3,31 @@
 
 openvm::entry!(main);
 
-pub fn main() {
-    // arbitrary n that results in more than 1 segment
-    let n = core::hint::black_box((1 << 5) + 4);
+fn fibonacci(n: u32) -> (u32, u32) {
+    if n <= 1 {
+        return (0, n);
+    }
     let mut a: u32 = 0;
     let mut b: u32 = 1;
-    for _ in 1..n {
+    for _ in 2..=n {
         let sum = a + b;
         a = b;
         b = sum;
     }
+    (a, b)
+}
+
+pub fn main() {
+    // arbitrary n that results in more than 1 segment
+    let n = core::hint::black_box(1 << 5);
+
+    let mut a = 0;
+    let mut b = 0;
+    // calculate nth fibonacci number n times
+    for _ in 0..n {
+        (a, b) = fibonacci(n);
+    }
+
     if a == 0 {
         panic!();
     }

crates/vm/src/arch/config.rs

@@ -1,6 +1,7 @@
-use std::{fs::File, io::Write, path::Path, sync::Arc};
+use std::{fs::File, io::Write, path::Path};
 
 use derive_new::new;
+use getset::{Setters, WithSetters};
 use openvm_instructions::NATIVE_AS;
 use openvm_poseidon2_air::Poseidon2Config;
 use openvm_stark_backend::{
@@ -10,14 +11,11 @@ use openvm_stark_backend::{
 };
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
-use super::{
-    segmentation_strategy::{DefaultSegmentationStrategy, SegmentationStrategy},
-    AnyEnum, VmChipComplex, PUBLIC_VALUES_AIR_ID,
-};
+use super::{AnyEnum, VmChipComplex, PUBLIC_VALUES_AIR_ID};
 use crate::{
     arch::{
-        AirInventory, AirInventoryError, Arena, ChipInventoryError, ExecutorInventory,
-        ExecutorInventoryError,
+        execution_mode::metered::segment_ctx::SegmentationLimits, AirInventory, AirInventoryError,
+        Arena, ChipInventoryError, ExecutorInventory, ExecutorInventoryError,
     },
     system::{
         memory::{merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, CHUNK},
@@ -172,9 +170,10 @@ impl MemoryConfig {
 
 /// System-level configuration for the virtual machine. Contains all configuration parameters that
 /// are managed by the architecture, including configuration for continuations support.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, Setters, WithSetters)]
 pub struct SystemConfig {
     /// The maximum constraint degree any chip is allowed to use.
+    #[getset(set_with = "pub")]
     pub max_constraint_degree: usize,
     /// True if the VM is in continuation mode. In this mode, an execution could be segmented and
     /// each segment is proved by a proof. Each proof commits the before and after state of the
@@ -195,15 +194,12 @@ pub struct SystemConfig {
     /// Whether to collect detailed profiling metrics.
     /// **Warning**: this slows down the runtime.
     pub profiling: bool,
-    /// Segmentation strategy
+    /// Segmentation limits
     /// This field is skipped in serde as it's only used in execution and
     /// not needed after any serialize/deserialize.
-    #[serde(skip, default = "get_default_segmentation_strategy")]
-    pub segmentation_strategy: Arc<dyn SegmentationStrategy>,
-}
-
-pub fn get_default_segmentation_strategy() -> Arc<DefaultSegmentationStrategy> {
-    Arc::new(DefaultSegmentationStrategy::default())
+    #[serde(skip, default = "SegmentationLimits::default")]
+    #[getset(set = "pub")]
+    pub segmentation_limits: SegmentationLimits,
 }
 
 impl SystemConfig {
@@ -212,7 +208,6 @@ impl SystemConfig {
         mut memory_config: MemoryConfig,
         num_public_values: usize,
     ) -> Self {
-        let segmentation_strategy = get_default_segmentation_strategy();
         assert!(
             memory_config.clk_max_bits <= 29,
             "Timestamp max bits must be <= 29 for LessThan to work in 31-bit field"
@@ -223,8 +218,8 @@ impl SystemConfig {
             continuation_enabled: false,
             memory_config,
             num_public_values,
-            segmentation_strategy,
             profiling: false,
+            segmentation_limits: SegmentationLimits::default(),
         }
     }
 
@@ -236,11 +231,6 @@ impl SystemConfig {
         )
     }
 
-    pub fn with_max_constraint_degree(mut self, max_constraint_degree: usize) -> Self {
-        self.max_constraint_degree = max_constraint_degree;
-        self
-    }
-
     pub fn with_continuations(mut self) -> Self {
         self.continuation_enabled = true;
         self
@@ -258,16 +248,10 @@ impl SystemConfig {
     }
 
     pub fn with_max_segment_len(mut self, max_segment_len: usize) -> Self {
-        self.segmentation_strategy = Arc::new(
-            DefaultSegmentationStrategy::new_with_max_segment_len(max_segment_len),
-        );
+        self.segmentation_limits.max_trace_height = max_segment_len as u32;
         self
     }
 
-    pub fn set_segmentation_strategy(&mut self, strategy: Arc<dyn SegmentationStrategy>) {
-        self.segmentation_strategy = strategy;
-    }
-
     pub fn with_profiling(mut self) -> Self {
         self.profiling = true;
         self

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

@@ -9,14 +9,16 @@ use super::{
 };
 use crate::{
     arch::{
-        execution_mode::{E1ExecutionCtx, E2ExecutionCtx},
+        execution_mode::{
+            metered::segment_ctx::SegmentationLimits, E1ExecutionCtx, E2ExecutionCtx,
+        },
         VmSegmentState,
     },
     system::memory::{dimensions::MemoryDimensions, online::GuestMemory},
 };
 
 pub const DEFAULT_PAGE_BITS: usize = 6;
-const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
+pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
 
 #[derive(Clone, Debug, WithSetters)]
 pub struct MeteredCtx<const PAGE_BITS: usize = DEFAULT_PAGE_BITS> {
@@ -45,6 +47,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         air_names: Vec<String>,
         widths: Vec<usize>,
         interactions: Vec<usize>,
+        segmentation_limits: SegmentationLimits,
     ) -> Self {
         let (trace_heights, is_trace_height_constant): (Vec<u32>, Vec<bool>) =
             constant_trace_heights
@@ -84,7 +87,8 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             air_names[memory_ctx.adapter_offset]
         );
 
-        let segmentation_ctx = SegmentationCtx::new(air_names, widths, interactions);
+        let segmentation_ctx =
+            SegmentationCtx::new(air_names, widths, interactions, segmentation_limits);
 
         let mut ctx = Self {
             trace_heights,

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

@@ -1,3 +1,4 @@
+use getset::WithSetters;
 use openvm_stark_backend::p3_field::PrimeField32;
 use p3_baby_bear::BabyBear;
 use serde::{Deserialize, Serialize};
@@ -13,10 +14,13 @@ pub struct Segment {
     pub trace_heights: Vec<u32>,
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, WithSetters)]
 pub struct SegmentationLimits {
+    #[getset(set_with = "pub")]
     pub max_trace_height: u32,
+    #[getset(set_with = "pub")]
     pub max_cells: usize,
+    #[getset(set_with = "pub")]
     pub max_interactions: usize,
 }
 
@@ -40,7 +44,26 @@ pub struct SegmentationCtx {
 }
 
 impl SegmentationCtx {
-    pub fn new(air_names: Vec<String>, widths: Vec<usize>, interactions: Vec<usize>) -> Self {
+    pub fn new(
+        air_names: Vec<String>,
+        widths: Vec<usize>,
+        interactions: Vec<usize>,
+        segmentation_limits: SegmentationLimits,
+    ) -> Self {
+        Self {
+            segments: Vec::new(),
+            air_names,
+            widths,
+            interactions,
+            segmentation_limits,
+        }
+    }
+
+    pub fn new_with_default_segmentation_limits(
+        air_names: Vec<String>,
+        widths: Vec<usize>,
+        interactions: Vec<usize>,
+    ) -> Self {
         Self {
             segments: Vec::new(),
             air_names,

crates/vm/src/arch/mod.rs

@@ -14,8 +14,6 @@ mod record_arena;
 /// Runtime execution and segmentation
 // TODO: rename this module
 pub mod segment;
-/// Strategy for determining when to segment VM execution
-pub mod segmentation_strategy;
 /// Top level [VmExecutor] and [VirtualMachine] constructor and API.
 pub mod vm;
 
@@ -34,5 +32,4 @@ pub use integration_api::*;
 pub use openvm_instructions as instructions;
 pub use record_arena::*;
 pub use segment::*;
-pub use segmentation_strategy::*;
 pub use vm::*;