feat: update segmentation weights for hardware memory (#2341)

segment_ctx.rs:
  - DEFAULT_MAX_CELLS → DEFAULT_MAX_MEMORY = 15gb
  - max_cells → max_memory in SegmentationLimits
  - set_max_cells → set_max_memory

  ctx.rs:
  - with_max_cells → with_max_memory

  metered_cost.rs:
  - Updated import to use DEFAULT_MAX_MEMORY

  cli/src/commands/prove.rs:
  - Updated import to DEFAULT_MAX_MEMORY
  - segment_max_cells → segment_max_memory
  - with_max_cells → with_max_memory

  benchmarks/prove/src/bin/async_regex.rs:
  - segment_max_cells → segment_max_memory
  - set_max_cells → set_max_memory

  benchmarks/prove/src/util.rs:
  - segment_max_cells → segment_max_memory
  - set_max_cells → set_max_memory

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: jonathanpwang

.github/workflows/base-tests.cuda.yml

@@ -74,4 +74,4 @@ jobs:
       # - name: Run async concurrency test
       #   working-directory: benchmarks/prove
       #   run: |
-      #     CUDA_OPT_LEVEL=3 MAX_CONCURRENCY=5 cargo run --bin async_regex --features cuda,async -- --max-segment-length $((1<<20)) --segment-max-cells 150000000
+      #     CUDA_OPT_LEVEL=3 MAX_CONCURRENCY=5 cargo run --bin async_regex --features cuda,async -- --max-segment-length $((1<<20)) --segment-max-memory 16106127360

Cargo.toml

@@ -233,6 +233,7 @@ libloading = "0.8"
 tracing-subscriber = { version = "0.3.20", features = ["std", "env-filter"] }
 tokio = "1" # >=1.0.0 to allow downstream flexibility
 abi_stable = "0.11.3"
+bytesize = "2.0"
 
 # default-features = false for no_std for use in guest programs
 itertools = { version = "0.14.0", default-features = false }

benchmarks/prove/src/bin/async_regex.rs

@@ -23,13 +23,13 @@ async fn main() -> eyre::Result<()> {
             .limits
             .set_max_trace_height(max_height);
     }
-    if let Some(max_cells) = args.segment_max_cells {
+    if let Some(max_memory) = args.segment_max_memory {
         config
             .app_vm_config
             .as_mut()
             .segmentation_config
             .limits
-            .set_max_cells(max_cells);
+            .set_max_memory(max_memory);
     }
 
     let sdk = Sdk::new(config)?;

benchmarks/prove/src/util.rs

@@ -66,9 +66,9 @@ pub struct BenchmarkCli {
     #[arg(long, alias = "max_segment_length")]
     pub max_segment_length: Option<u32>,
 
-    /// Total cells used in all chips in segment for continuations
+    /// Total memory in bytes used in all chips in segment for continuations
     #[arg(long)]
-    pub segment_max_cells: Option<usize>,
+    pub segment_max_memory: Option<usize>,
 
     /// Controls the arity (num_children) of the aggregation tree
     #[command(flatten)]
@@ -96,12 +96,12 @@ impl BenchmarkCli {
                 .limits
                 .set_max_trace_height(max_height);
         }
-        if let Some(max_cells) = self.segment_max_cells {
+        if let Some(max_memory) = self.segment_max_memory {
             app_vm_config
                 .as_mut()
                 .segmentation_config
                 .limits
-                .set_max_cells(max_cells);
+                .set_max_memory(max_memory);
         }
         AppConfig {
             app_fri_params: FriParameters::standard_with_100_bits_conjectured_security(

crates/cli/src/commands/prove.rs

@@ -4,7 +4,7 @@ use clap::Parser;
 use eyre::Result;
 use openvm_circuit::arch::{
     execution_mode::metered::segment_ctx::{
-        SegmentationConfig, SegmentationLimits, DEFAULT_MAX_CELLS, DEFAULT_MAX_TRACE_HEIGHT_BITS,
+        SegmentationConfig, SegmentationLimits, DEFAULT_MAX_MEMORY, DEFAULT_MAX_TRACE_HEIGHT_BITS,
     },
     instructions::exe::VmExe,
 };
@@ -133,14 +133,14 @@ pub struct SegmentationArgs {
         help_heading = "OpenVM Options"
     )]
     pub segment_max_height_bits: u8,
-    /// Total cells used across all chips for triggering segmentation for continuations in the app
-    /// proof. These thresholds are not exceeded except when they are too small.
+    /// Total memory in bytes used across all chips for triggering segmentation for continuations
+    /// in the app proof. These thresholds are not exceeded except when they are too small.
     #[arg(
         long,
-        default_value_t = DEFAULT_MAX_CELLS,
+        default_value_t = DEFAULT_MAX_MEMORY,
         help_heading = "OpenVM Options"
     )]
-    pub segment_max_cells: usize,
+    pub segment_max_memory: usize,
 }
 
 impl ProveCmd {
@@ -318,7 +318,7 @@ impl From<SegmentationArgs> for SegmentationConfig {
                     1u32.checked_shl(args.segment_max_height_bits as u32)
                         .expect("segment_max_height_bits too large"),
                 )
-                .with_max_cells(args.segment_max_cells),
+                .with_max_memory(args.segment_max_memory),
             ..Default::default()
         }
     }

crates/vm/Cargo.toml

@@ -46,6 +46,7 @@ cfg-if.workspace = true
 libloading = { workspace = true, optional = true }
 tempfile = { workspace = true, optional = true }
 abi_stable.workspace = true
+bytesize.workspace = true
 
 [build-dependencies]
 openvm-cuda-builder = { workspace = true, optional = true }

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

@@ -103,8 +103,8 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         self
     }
 
-    pub fn with_max_cells(mut self, max_cells: usize) -> Self {
-        self.segmentation_ctx.set_max_cells(max_cells);
+    pub fn with_max_memory(mut self, max_memory: usize) -> Self {
+        self.segmentation_ctx.set_max_memory(max_memory);
         self
     }
 
@@ -127,11 +127,21 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         self
     }
 
+    pub fn with_main_cell_weight(mut self, weight: usize) -> Self {
+        self.segmentation_ctx.set_main_cell_weight(weight);
+        self
+    }
+
     pub fn with_interaction_cell_weight(mut self, weight: usize) -> Self {
         self.segmentation_ctx.set_interaction_cell_weight(weight);
         self
     }
 
+    pub fn with_base_field_size(mut self, base_field_size: usize) -> Self {
+        self.segmentation_ctx.set_base_field_size(base_field_size);
+        self
+    }
+
     pub fn segments(&self) -> &[Segment] {
         &self.segmentation_ctx.segments
     }

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

@@ -1,3 +1,6 @@
+use std::mem::size_of;
+
+use bytesize::ByteSize;
 use getset::{Setters, WithSetters};
 use openvm_stark_backend::p3_field::PrimeField32;
 use p3_baby_bear::BabyBear;
@@ -7,8 +10,10 @@ pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
 
 pub const DEFAULT_MAX_TRACE_HEIGHT_BITS: u8 = 22;
 pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << DEFAULT_MAX_TRACE_HEIGHT_BITS;
-pub const DEFAULT_MAX_CELLS: usize = 1_200_000_000; // 1.2B
+pub const DEFAULT_MAX_MEMORY: usize = 15 << 30; // 15GiB
 const DEFAULT_MAX_INTERACTIONS: usize = BabyBear::ORDER_U32 as usize;
+const DEFAULT_MAIN_CELL_WEIGHT: usize = 3; // 1 + 2^{log_blowup=1}
+const DEFAULT_INTERACTION_CELL_WEIGHT: usize = 8; // 2 * D_EF
 
 #[derive(derive_new::new, Clone, Debug, Serialize, Deserialize)]
 pub struct Segment {
@@ -17,19 +22,36 @@ pub struct Segment {
     pub trace_heights: Vec<u32>,
 }
 
-#[derive(Clone, Debug, Default, WithSetters)]
+#[derive(Clone, Debug, WithSetters)]
 pub struct SegmentationConfig {
     pub limits: SegmentationLimits,
-    /// Cells per row contributed by each interaction in cell count.
+    /// Weight multiplier for main trace cells in memory calculation.
+    #[getset(set_with = "pub")]
+    pub main_cell_weight: usize,
+    /// Weight multiplier for interaction cells in memory calculation.
     #[getset(set_with = "pub")]
     pub interaction_cell_weight: usize,
+    /// Size of the base field in bytes. Used to convert cell count to memory bytes.
+    #[getset(set_with = "pub")]
+    pub base_field_size: usize,
+}
+
+impl Default for SegmentationConfig {
+    fn default() -> Self {
+        Self {
+            limits: SegmentationLimits::default(),
+            main_cell_weight: DEFAULT_MAIN_CELL_WEIGHT,
+            interaction_cell_weight: DEFAULT_INTERACTION_CELL_WEIGHT,
+            base_field_size: size_of::<u32>(),
+        }
+    }
 }
 
 #[derive(Clone, Debug, WithSetters, Setters)]
 pub struct SegmentationLimits {
     pub max_trace_height: u32,
     #[getset(set = "pub", set_with = "pub")]
-    pub max_cells: usize,
+    pub max_memory: usize,
     #[getset(set_with = "pub")]
     pub max_interactions: usize,
 }
@@ -38,21 +60,21 @@ impl Default for SegmentationLimits {
     fn default() -> Self {
         Self {
             max_trace_height: DEFAULT_MAX_TRACE_HEIGHT,
-            max_cells: DEFAULT_MAX_CELLS,
+            max_memory: DEFAULT_MAX_MEMORY,
             max_interactions: DEFAULT_MAX_INTERACTIONS,
         }
     }
 }
 
 impl SegmentationLimits {
-    pub fn new(max_trace_height: u32, max_cells: usize, max_interactions: usize) -> Self {
+    pub fn new(max_trace_height: u32, max_memory: usize, max_interactions: usize) -> Self {
         debug_assert!(
             max_trace_height.is_power_of_two(),
             "max_trace_height should be a power of two"
         );
         Self {
             max_trace_height,
-            max_cells,
+            max_memory,
             max_interactions,
         }
     }
@@ -120,18 +142,26 @@ impl SegmentationCtx {
         self.config.limits.set_max_trace_height(max_trace_height);
     }
 
-    pub fn set_max_cells(&mut self, max_cells: usize) {
-        self.config.limits.max_cells = max_cells;
+    pub fn set_max_memory(&mut self, max_memory: usize) {
+        self.config.limits.max_memory = max_memory;
     }
 
     pub fn set_max_interactions(&mut self, max_interactions: usize) {
         self.config.limits.max_interactions = max_interactions;
     }
 
+    pub fn set_main_cell_weight(&mut self, weight: usize) {
+        self.config.main_cell_weight = weight;
+    }
+
     pub fn set_interaction_cell_weight(&mut self, weight: usize) {
         self.config.interaction_cell_weight = weight;
     }
 
+    pub fn set_base_field_size(&mut self, base_field_size: usize) {
+        self.config.base_field_size = base_field_size;
+    }
+
     /// Calculate the maximum trace height and corresponding air name
     #[inline(always)]
     fn calculate_max_trace_height_with_name(&self, trace_heights: &[u32]) -> (u32, &str) {
@@ -144,22 +174,44 @@ impl SegmentationCtx {
             .unwrap_or((0, "unknown"))
     }
 
-    /// Calculate the total cells used based on trace heights and widths,
-    /// including weighted contribution from interactions if `interaction_cell_weight > 0`.
+    /// Calculate total memory in bytes based on trace heights and widths.
+    /// Formula: base_field_size * (main_cell_weight * main_cells + interaction_cell_weight *
+    /// interaction_cells)
     #[inline(always)]
-    fn calculate_total_cells(&self, trace_heights: &[u32]) -> usize {
+    fn calculate_total_memory(
+        &self,
+        trace_heights: &[u32],
+    ) -> (
+        usize, /* memory */
+        usize, /* main */
+        usize, /* interaction */
+    ) {
         debug_assert_eq!(trace_heights.len(), self.widths.len());
 
+        let main_weight = self.config.main_cell_weight;
         let interaction_weight = self.config.interaction_cell_weight;
-        trace_heights
+        let base_field_size = self.config.base_field_size;
+
+        let mut main_cnt = 0;
+        let mut interaction_cnt = 0;
+        for ((&height, &width), &interactions) in trace_heights
             .iter()
             .zip(self.widths.iter())
             .zip(self.interactions.iter())
-            .map(|((&height, &width), &interactions)| {
-                let padded_height = height.next_power_of_two() as usize;
-                padded_height * (width + interactions * interaction_weight)
-            })
-            .sum()
+        {
+            let padded_height = height.next_power_of_two() as usize;
+            main_cnt += padded_height * width;
+            interaction_cnt += padded_height * interactions;
+        }
+
+        let main_memory = main_cnt * main_weight * base_field_size;
+        let interaction_memory =
+            (interaction_cnt + 1).next_power_of_two() * interaction_weight * base_field_size;
+        (
+            main_memory + interaction_memory,
+            main_memory,
+            interaction_memory,
+        )
     }
 
     /// Calculate the total interactions based on trace heights
@@ -199,8 +251,11 @@ impl SegmentationCtx {
             return false;
         }
 
+        let main_weight = self.config.main_cell_weight;
         let interaction_weight = self.config.interaction_cell_weight;
-        let mut total_cells = 0;
+        let base_field_size = self.config.base_field_size;
+        let mut main_cnt = 0usize;
+        let mut interaction_cnt = 0usize;
         for (i, (((padded_height, width), interactions), is_constant)) in trace_heights
             .iter()
             .map(|&height| height.next_power_of_two())
@@ -214,7 +269,7 @@ impl SegmentationCtx {
             if !is_constant && padded_height > self.config.limits.max_trace_height {
                 let air_name = unsafe { self.air_names.get_unchecked(i) };
                 tracing::info!(
-                    "instret {:10} | height ({:8}) > max ({:8}) | chip {:3} ({}) ",
+                    "overshoot: instret {:10} | height ({:8}) > max ({:8}) | chip {:3} ({}) ",
                     instret,
                     padded_height,
                     self.config.limits.max_trace_height,
@@ -223,23 +278,30 @@ impl SegmentationCtx {
                 );
                 return true;
             }
-            total_cells += padded_height as usize * (width + interactions * interaction_weight);
+            main_cnt += padded_height as usize * width;
+            interaction_cnt += padded_height as usize * interactions;
         }
+        // interaction rounding to match n_logup calculation
+        let total_memory = (main_cnt * main_weight
+            + (interaction_cnt + 1).next_power_of_two() * interaction_weight)
+            * base_field_size;
 
-        if total_cells > self.config.limits.max_cells {
+        if total_memory > self.config.limits.max_memory {
             tracing::info!(
-                "instret {:10} | total cells ({:10}) > max ({:10})",
+                "overshoot: instret {:10} | total memory ({:10}) > max ({:10}) | main ({:10}) | interaction ({:10})",
                 instret,
-                total_cells,
-                self.config.limits.max_cells
+                total_memory,
+                self.config.limits.max_memory,
+                main_cnt,
+                interaction_cnt
             );
             return true;
         }
 
         let total_interactions = self.calculate_total_interactions(trace_heights);
         if total_interactions > self.config.limits.max_interactions {
             tracing::info!(
-                "instret {:10} | total interactions ({:10}) > max ({:10})",
+                "overshoot: instret {:10} | total interactions ({:10}) > max ({:10})",
                 instret,
                 total_interactions,
                 self.config.limits.max_interactions
@@ -403,18 +465,21 @@ impl SegmentationCtx {
         trace_heights: &[u32],
     ) {
         let (max_trace_height, air_name) = self.calculate_max_trace_height_with_name(trace_heights);
-        let total_cells = self.calculate_total_cells(trace_heights);
+        let (total_memory, main_memory, interaction_memory) =
+            self.calculate_total_memory(trace_heights);
         let total_interactions = self.calculate_total_interactions(trace_heights);
         let utilization = self.calculate_trace_utilization(trace_heights);
 
         let final_marker = if IS_FINAL { " [TERMINATED]" } else { "" };
 
         tracing::info!(
-            "Segment {:3} | instret {:10} | {:8} instructions | {:10} cells | {:10} interactions | {:8} max height ({}) | {:.2}% utilization{}",
+            "Segment {:3} | instret {:10} | {:8} instructions | {:5} memory ({:5}, {:5}) | {:10} interactions | {:8} max height ({}) | {:.2}% utilization{}",
             self.segments.len(),
             instret_start,
             num_insns,
-            total_cells,
+            ByteSize::b(total_memory as u64),
+            ByteSize::b(main_memory as u64),
+            ByteSize::b(interaction_memory as u64),
             total_interactions,
             max_trace_height,
             air_name,