[feat] support custom segmentation (#1254)

* support custom segmentation

* address comments

Author: luffykai

crates/vm/src/arch/config.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use derive_new::new;
 use openvm_circuit::system::memory::MemoryTraceHeights;
 use openvm_instructions::program::DEFAULT_MAX_NUM_PUBLIC_VALUES;
@@ -13,12 +15,12 @@ pub use super::testing::{
     POSEIDON2_DIRECT_BUS, RANGE_TUPLE_CHECKER_BUS, READ_INSTRUCTION_BUS,
 };
 use super::{
+    segment::{DefaultSegmentationStrategy, SegmentationStrategy},
     AnyEnum, InstructionExecutor, SystemComplex, SystemExecutor, SystemPeriphery, VmChipComplex,
     VmInventoryError, PUBLIC_VALUES_AIR_ID,
 };
 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,11 +88,18 @@ 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,
+    /// Segmentation strategy
+    /// 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<dyn SegmentationStrategy> {
+    Arc::new(DefaultSegmentationStrategy::default())
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -105,12 +114,13 @@ impl SystemConfig {
         memory_config: MemoryConfig,
         num_public_values: usize,
     ) -> Self {
+        let segmentation_strategy = get_default_segmentation_strategy();
         Self {
             max_constraint_degree,
             continuation_enabled: false,
             memory_config,
             num_public_values,
-            max_segment_len: DEFAULT_MAX_SEGMENT_LEN,
+            segmentation_strategy,
             profiling: false,
         }
     }
@@ -136,10 +146,16 @@ impl SystemConfig {
     }
 
     pub fn with_max_segment_len(mut self, max_segment_len: usize) -> Self {
-        self.max_segment_len = max_segment_len;
+        self.segmentation_strategy = Arc::new(
+            DefaultSegmentationStrategy::new_with_max_segment_len(max_segment_len),
+        );
         self
     }
 
+    pub fn set_segmentation_strategy<S: SegmentationStrategy + 'static>(&mut self, strategy: S) {
+        self.segmentation_strategy = Arc::new(strategy);
+    }
+
     pub fn with_profiling(mut self) -> Self {
         self.profiling = true;
         self

crates/vm/src/arch/segment.rs

@@ -27,6 +27,84 @@ use crate::{
 /// Check segment every 100 instructions.
 const SEGMENT_CHECK_INTERVAL: usize = 100;
 
+const DEFAULT_MAX_SEGMENT_LEN: usize = (1 << 22) - 100;
+// a heuristic number for the maximum number of cells per chip in a segment
+// a few reasons for this number:
+//  1. `VmAirWrapper<Rv32BaseAluAdapterAir, BaseAluCoreAir<4, 8>` is
+//    the chip with the most cells in a segment from the reth-benchmark.
+//  2. `VmAirWrapper<Rv32BaseAluAdapterAir, BaseAluCoreAir<4, 8>`:
+//    its trace width is 36 and its after challenge trace width is 80.
+const DEFAULT_MAX_CELLS_PER_CHIP_IN_SEGMENT: usize = DEFAULT_MAX_SEGMENT_LEN * 120;
+
+pub trait SegmentationStrategy:
+    std::fmt::Debug + Send + Sync + std::panic::UnwindSafe + std::panic::RefUnwindSafe
+{
+    fn should_segment(
+        &self,
+        air_names: &[String],
+        trace_heights: &[usize],
+        trace_cells: &[usize],
+    ) -> bool;
+}
+
+/// Default segmentation strategy: segment if any chip's height or cells exceed the limits.
+#[derive(Debug)]
+pub struct DefaultSegmentationStrategy {
+    max_segment_len: usize,
+    max_cells_per_chip_in_segment: usize,
+}
+
+impl Default for DefaultSegmentationStrategy {
+    fn default() -> Self {
+        Self {
+            max_segment_len: DEFAULT_MAX_SEGMENT_LEN,
+            max_cells_per_chip_in_segment: DEFAULT_MAX_CELLS_PER_CHIP_IN_SEGMENT,
+        }
+    }
+}
+
+impl DefaultSegmentationStrategy {
+    pub fn new_with_max_segment_len(max_segment_len: usize) -> Self {
+        Self {
+            max_segment_len,
+            max_cells_per_chip_in_segment: max_segment_len * 120,
+        }
+    }
+}
+
+impl SegmentationStrategy for DefaultSegmentationStrategy {
+    fn should_segment(
+        &self,
+        air_names: &[String],
+        trace_heights: &[usize],
+        trace_cells: &[usize],
+    ) -> bool {
+        for (i, &height) in trace_heights.iter().enumerate() {
+            if height > self.max_segment_len {
+                tracing::info!(
+                    "Should segment because chip {} (name: {}) has height {}",
+                    i,
+                    air_names[i],
+                    height
+                );
+                return true;
+            }
+        }
+        for (i, &num_cells) in trace_cells.iter().enumerate() {
+            if num_cells > self.max_cells_per_chip_in_segment {
+                tracing::info!(
+                    "Should segment because chip {} (name: {}) has {} cells",
+                    i,
+                    air_names[i],
+                    num_cells
+                );
+                return true;
+            }
+        }
+        false
+    }
+}
+
 pub struct ExecutionSegment<F, VC>
 where
     F: PrimeField32,
@@ -276,19 +354,15 @@ impl<F: PrimeField32, VC: VmConfig<F>> ExecutionSegment<F, VC> {
             return false;
         }
         self.since_last_segment_check = 0;
-        let heights = self.chip_complex.dynamic_trace_heights();
-        for (i, height) in heights.enumerate() {
-            if height > self.system_config().max_segment_len {
-                tracing::info!(
-                    "Should segment because chip {} has height {}",
-                    self.air_names[i],
-                    height
-                );
-                return true;
-            }
-        }
-
-        false
+        let segmentation_strategy = self.system_config().segmentation_strategy.clone();
+        segmentation_strategy.should_segment(
+            &self.air_names,
+            &self
+                .chip_complex
+                .dynamic_trace_heights()
+                .collect::<Vec<_>>(),
+            &self.chip_complex.current_trace_cells(),
+        )
     }
 
     pub fn current_trace_cells(&self) -> Vec<usize> {