openvm_circuit/arch/execution_mode/metered/
segment_ctx.rs

1use bytesize::ByteSize;
2use itertools::izip;
3use openvm_stark_backend::memory_metering::{ProvingMemoryConfig, ProvingMemoryCounts};
4use serde::{Deserialize, Serialize};
5
6use crate::utils::{add_one_or_zero, next_power_of_two_or_zero};
7
8pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
9
10pub const DEFAULT_MAX_MEMORY: usize = 15 << 30; // 15GiB
11
12#[derive(derive_new::new, Clone, Debug, Serialize, Deserialize)]
13pub struct Segment {
14    pub instret_start: u64,
15    pub num_insns: u64,
16    pub trace_heights: Vec<u32>,
17}
18
19#[derive(Clone, Copy, Debug)]
20pub struct SegmentationLimits {
21    pub max_trace_height_bits: u8,
22    pub max_memory: usize,
23    pub max_interactions: u32,
24}
25
26#[derive(Clone, Debug)]
27struct SegmentationParams {
28    air_names: Vec<String>,
29    widths: Vec<usize>,
30    interactions: Vec<usize>,
31    need_rot: Vec<bool>,
32    constraint_eval_buffers: Vec<usize>,
33    max_trace_height: u32,
34    max_memory: usize,
35    max_interactions: u32,
36    memory_config: ProvingMemoryConfig,
37    segment_check_insns: u64,
38}
39
40impl SegmentationParams {
41    fn new(
42        air_names: Vec<String>,
43        widths: Vec<usize>,
44        interactions: Vec<usize>,
45        need_rot: Vec<bool>,
46        constraint_eval_buffers: Vec<usize>,
47        limits: SegmentationLimits,
48        memory_config: ProvingMemoryConfig,
49    ) -> Self {
50        assert_eq!(air_names.len(), widths.len());
51        assert_eq!(air_names.len(), interactions.len());
52        assert_eq!(air_names.len(), need_rot.len());
53        assert_eq!(air_names.len(), constraint_eval_buffers.len());
54        assert!(
55            limits.max_trace_height_bits < u32::BITS as u8,
56            "max_trace_height_bits must be less than {}",
57            u32::BITS
58        );
59
60        let max_trace_height = 1u32
61            .checked_shl(u32::from(limits.max_trace_height_bits))
62            .expect("max_trace_height_bits must fit in u32 trace height");
63        assert!(
64            u64::from(max_trace_height) >= 2 * DEFAULT_SEGMENT_CHECK_INSNS,
65            "max_trace_height must be at least twice DEFAULT_SEGMENT_CHECK_INSNS"
66        );
67
68        Self {
69            air_names,
70            widths,
71            interactions,
72            need_rot,
73            constraint_eval_buffers,
74            max_trace_height,
75            max_memory: limits.max_memory,
76            max_interactions: limits.max_interactions,
77            memory_config,
78            segment_check_insns: DEFAULT_SEGMENT_CHECK_INSNS,
79        }
80    }
81}
82
83#[derive(Clone, Debug)]
84pub struct SegmentationCtx {
85    pub segments: Vec<Segment>,
86    params: SegmentationParams,
87    pub instret: u64,
88    pub instrets_until_check: u64,
89    /// Checkpoint of trace heights at last known state where all thresholds satisfied
90    pub(crate) checkpoint_trace_heights: Vec<u32>,
91    /// Instruction count at the checkpoint
92    checkpoint_instret: u64,
93}
94
95#[derive(Clone, Copy, Debug)]
96enum SegmentationTrigger {
97    Height {
98        #[cfg(feature = "metrics")]
99        air_id: usize,
100    },
101    Memory,
102    Interactions,
103}
104
105#[cfg(feature = "metrics")]
106impl SegmentationTrigger {
107    fn reason(self) -> &'static str {
108        match self {
109            SegmentationTrigger::Height { .. } => "height",
110            SegmentationTrigger::Memory => "memory",
111            SegmentationTrigger::Interactions => "interactions",
112        }
113    }
114}
115
116#[derive(Default)]
117struct MeteredCounts {
118    /// Rows before power-of-two padding.
119    unpadded_rows: usize,
120    /// Rows added by power-of-two padding.
121    padding_rows: usize,
122    /// Main trace cells for AIRs that open next-row rotations, before padding.
123    main_unpadded_with_rot: usize,
124    /// Main trace cells for AIRs that open next-row rotations, from padding rows.
125    main_padding_with_rot: usize,
126    /// Main trace cells for AIRs without next-row rotations, before padding.
127    main_unpadded_no_rot: usize,
128    /// Main trace cells for AIRs without next-row rotations, from padding rows.
129    main_padding_no_rot: usize,
130    /// Metered row-interaction slots before padding.
131    interaction_cells_unpadded: usize,
132    /// Metered row-interaction slots from padding rows.
133    interaction_cells_padding: usize,
134    /// Constraint eval buffer size without padding.
135    constraint_eval_buffers_unpadded: usize,
136    /// Constraint eval buffer size from padding.
137    constraint_eval_buffers_padding: usize,
138}
139
140struct MeteredMemoryBreakdown {
141    /// Total selected segment memory estimate.
142    total: usize,
143    /// Unpadded-row contribution to the selected memory estimate.
144    unpadded: usize,
145}
146
147impl SegmentationCtx {
148    pub fn new(
149        air_names: Vec<String>,
150        widths: Vec<usize>,
151        interactions: Vec<usize>,
152        need_rot: Vec<bool>,
153        constraint_eval_buffers: Vec<usize>,
154        limits: SegmentationLimits,
155        memory_config: ProvingMemoryConfig,
156    ) -> Self {
157        let num_airs = air_names.len();
158        let params = SegmentationParams::new(
159            air_names,
160            widths,
161            interactions,
162            need_rot,
163            constraint_eval_buffers,
164            limits,
165            memory_config,
166        );
167        Self {
168            segments: Vec::new(),
169            instrets_until_check: params.segment_check_insns,
170            params,
171            instret: 0,
172            checkpoint_trace_heights: vec![0; num_airs],
173            checkpoint_instret: 0,
174        }
175    }
176
177    #[inline(always)]
178    pub(crate) fn air_names(&self) -> &[String] {
179        &self.params.air_names
180    }
181
182    #[inline(always)]
183    pub(crate) fn widths(&self) -> &[usize] {
184        &self.params.widths
185    }
186
187    #[inline(always)]
188    pub(super) fn segment_check_insns(&self) -> u64 {
189        self.params.segment_check_insns
190    }
191
192    pub fn set_max_memory(&mut self, max_memory: usize) {
193        self.params.max_memory = max_memory;
194    }
195
196    /// Calculate the maximum trace height and corresponding air name
197    #[inline(always)]
198    fn calculate_max_trace_height_with_name(&self, trace_heights: &[u32]) -> (u32, &str) {
199        trace_heights
200            .iter()
201            .enumerate()
202            .map(|(i, &height)| (next_power_of_two_or_zero(height as usize) as u32, i))
203            .max_by_key(|(height, _)| *height)
204            .map(|(height, idx)| (height, self.params.air_names[idx].as_str()))
205            .unwrap_or((0, "unknown"))
206    }
207
208    /// Convert main trace cells and interaction cells to memory bytes.
209    #[inline(always)]
210    fn counts_to_memory(
211        &self,
212        main_cnt_with_rot: usize,
213        main_cnt_no_rot: usize,
214        interaction_cells: usize,
215        constraint_eval_cells: usize,
216    ) -> (
217        usize, /* total */
218        usize, /* main */
219        usize, /* secondary */
220    ) {
221        let estimate = self.params.memory_config.estimate(ProvingMemoryCounts::new(
222            main_cnt_with_rot,
223            main_cnt_no_rot,
224            interaction_cells,
225            constraint_eval_cells,
226        ));
227        (estimate.total, estimate.main, estimate.secondary_peak)
228    }
229
230    /// Sum padded main trace cells and interaction cells across all chips, splitting main
231    /// cells by per-AIR `need_rot`.
232    #[inline(always)]
233    fn calculate_count_breakdown(&self, trace_heights: &[u32]) -> MeteredCounts {
234        debug_assert_eq!(trace_heights.len(), self.params.widths.len());
235        debug_assert_eq!(trace_heights.len(), self.params.interactions.len());
236        debug_assert_eq!(trace_heights.len(), self.params.need_rot.len());
237        debug_assert_eq!(
238            trace_heights.len(),
239            self.params.constraint_eval_buffers.len()
240        );
241
242        let mut counts = MeteredCounts::default();
243        for (&height, &width, &interactions, &need_rot, &constraint_eval_buffer) in izip!(
244            trace_heights,
245            &self.params.widths,
246            &self.params.interactions,
247            &self.params.need_rot,
248            &self.params.constraint_eval_buffers
249        ) {
250            let padded_height = next_power_of_two_or_zero(height as usize);
251            let unpadded_height = height as usize;
252            let padding_height = padded_height - unpadded_height;
253            counts.unpadded_rows += unpadded_height;
254            counts.padding_rows += padding_height;
255            let main_unpadded_cells = unpadded_height * width;
256            let main_padding_cells = padding_height * width;
257            if need_rot {
258                counts.main_unpadded_with_rot += main_unpadded_cells;
259                counts.main_padding_with_rot += main_padding_cells;
260            } else {
261                counts.main_unpadded_no_rot += main_unpadded_cells;
262                counts.main_padding_no_rot += main_padding_cells;
263            }
264            counts.interaction_cells_unpadded += unpadded_height * interactions;
265            counts.interaction_cells_padding += padding_height * interactions;
266            counts.constraint_eval_buffers_unpadded += unpadded_height * constraint_eval_buffer;
267            counts.constraint_eval_buffers_padding += padding_height * constraint_eval_buffer;
268        }
269        counts
270    }
271
272    /// Sum padded main trace cells and interaction cells across all chips, splitting main
273    /// cells by per-AIR `need_rot`.
274    #[inline(always)]
275    fn calculate_cell_counts(&self, trace_heights: &[u32]) -> (usize, usize, usize, usize) {
276        debug_assert_eq!(trace_heights.len(), self.params.widths.len());
277        debug_assert_eq!(trace_heights.len(), self.params.interactions.len());
278        debug_assert_eq!(trace_heights.len(), self.params.need_rot.len());
279        debug_assert_eq!(
280            trace_heights.len(),
281            self.params.constraint_eval_buffers.len()
282        );
283
284        let mut main_cnt_with_rot = 0;
285        let mut main_cnt_no_rot = 0;
286        let mut interaction_cells = 0;
287        let mut constraint_eval_cells = 0;
288        for (&height, &width, &interactions, &need_rot, &constraint_eval_buffer) in izip!(
289            trace_heights,
290            &self.params.widths,
291            &self.params.interactions,
292            &self.params.need_rot,
293            &self.params.constraint_eval_buffers
294        ) {
295            let padded_height = next_power_of_two_or_zero(height as usize);
296            let main_cells = padded_height * width;
297            if need_rot {
298                main_cnt_with_rot += main_cells;
299            } else {
300                main_cnt_no_rot += main_cells;
301            }
302            interaction_cells += padded_height * interactions;
303            constraint_eval_cells += padded_height * constraint_eval_buffer;
304        }
305        (
306            main_cnt_with_rot,
307            main_cnt_no_rot,
308            interaction_cells,
309            constraint_eval_cells,
310        )
311    }
312
313    /// Calculate total memory in bytes based on trace heights and widths.
314    #[inline(always)]
315    fn calculate_total_memory(
316        &self,
317        trace_heights: &[u32],
318    ) -> (
319        usize, /* total */
320        usize, /* main */
321        usize, /* secondary */
322    ) {
323        let (main_cnt_with_rot, main_cnt_no_rot, interaction_cells, constraint_eval_cells) =
324            self.calculate_cell_counts(trace_heights);
325        self.counts_to_memory(
326            main_cnt_with_rot,
327            main_cnt_no_rot,
328            interaction_cells,
329            constraint_eval_cells,
330        )
331    }
332
333    #[inline(always)]
334    fn calculate_memory_breakdown(&self, counts: &MeteredCounts) -> MeteredMemoryBreakdown {
335        let unpadded = self.params.memory_config.estimate(ProvingMemoryCounts::new(
336            counts.main_unpadded_with_rot,
337            counts.main_unpadded_no_rot,
338            counts.interaction_cells_unpadded,
339            counts.constraint_eval_buffers_unpadded,
340        ));
341        let total = self.params.memory_config.estimate(ProvingMemoryCounts::new(
342            counts.main_unpadded_with_rot + counts.main_padding_with_rot,
343            counts.main_unpadded_no_rot + counts.main_padding_no_rot,
344            counts.interaction_cells_unpadded + counts.interaction_cells_padding,
345            counts.constraint_eval_buffers_unpadded + counts.constraint_eval_buffers_padding,
346        ));
347
348        MeteredMemoryBreakdown {
349            total: total.total,
350            unpadded: unpadded.total,
351        }
352    }
353
354    /// Calculate the total interactions based on trace heights
355    /// All padding rows contribute a single message to the interactions (+1) since
356    /// we assume chips don't send/receive with nonzero multiplicity on padding rows.
357    #[inline(always)]
358    fn calculate_total_interactions(&self, trace_heights: &[u32]) -> u64 {
359        debug_assert_eq!(trace_heights.len(), self.params.interactions.len());
360
361        trace_heights
362            .iter()
363            .zip(self.params.interactions.iter())
364            .map(|(&height, &interactions)| add_one_or_zero(height) as u64 * interactions as u64)
365            .sum()
366    }
367
368    #[inline(always)]
369    pub(crate) fn should_segment(
370        &self,
371        instret: u64,
372        trace_heights: &[u32],
373        is_trace_height_constant: &[bool],
374    ) -> bool {
375        self.segmentation_trigger(instret, trace_heights, is_trace_height_constant)
376            .is_some()
377    }
378
379    #[inline(always)]
380    fn segmentation_trigger(
381        &self,
382        instret: u64,
383        trace_heights: &[u32],
384        is_trace_height_constant: &[bool],
385    ) -> Option<SegmentationTrigger> {
386        debug_assert_eq!(trace_heights.len(), is_trace_height_constant.len());
387        debug_assert_eq!(trace_heights.len(), self.params.air_names.len());
388        debug_assert_eq!(trace_heights.len(), self.params.widths.len());
389        debug_assert_eq!(trace_heights.len(), self.params.interactions.len());
390        debug_assert_eq!(trace_heights.len(), self.params.need_rot.len());
391
392        let instret_start = self
393            .segments
394            .last()
395            .map_or(0, |s| s.instret_start + s.num_insns);
396        let num_insns = instret - instret_start;
397
398        // Segment should contain at least one cycle
399        if num_insns == 0 {
400            return None;
401        }
402
403        let mut main_cnt_with_rot = 0usize;
404        let mut main_cnt_no_rot = 0usize;
405        let mut interaction_cells = 0usize;
406        let mut constraint_eval_cells = 0usize;
407        let padded_heights = trace_heights
408            .iter()
409            .map(|&height| next_power_of_two_or_zero(height as usize) as u32);
410        for (i, row) in izip!(
411            padded_heights,
412            &self.params.widths,
413            &self.params.interactions,
414            is_trace_height_constant,
415            &self.params.need_rot,
416            &self.params.constraint_eval_buffers
417        )
418        .enumerate()
419        {
420            let (padded_height, &width, &interactions, &is_constant, &need_rot, &constraint_eval) =
421                row;
422            // Only segment if the height is not constant and exceeds the maximum height after
423            // padding
424            if !is_constant && padded_height > self.params.max_trace_height {
425                let air_name = unsafe { self.params.air_names.get_unchecked(i) };
426                tracing::info!(
427                    "overshoot: instret {:10} | height ({:8}) > max ({:8}) | chip {:3} ({}) ",
428                    instret,
429                    padded_height,
430                    self.params.max_trace_height,
431                    i,
432                    air_name,
433                );
434                return Some(SegmentationTrigger::Height {
435                    #[cfg(feature = "metrics")]
436                    air_id: i,
437                });
438            }
439            let main_cells = padded_height as usize * width;
440            if need_rot {
441                main_cnt_with_rot += main_cells;
442            } else {
443                main_cnt_no_rot += main_cells;
444            }
445            interaction_cells += padded_height as usize * interactions;
446            constraint_eval_cells += padded_height as usize * constraint_eval;
447        }
448
449        let (total_memory, main_memory, interaction_memory) = self.counts_to_memory(
450            main_cnt_with_rot,
451            main_cnt_no_rot,
452            interaction_cells,
453            constraint_eval_cells,
454        );
455        if total_memory > self.params.max_memory {
456            tracing::info!(
457                "overshoot: instret {:10} | total memory ({:5}) > max ({:5}) | main ({:5}) | interaction ({:5})",
458                instret,
459                ByteSize::b(total_memory as u64),
460                ByteSize::b(self.params.max_memory as u64),
461                ByteSize::b(main_memory as u64),
462                ByteSize::b(interaction_memory as u64),
463            );
464            return Some(SegmentationTrigger::Memory);
465        }
466
467        let total_interactions = self.calculate_total_interactions(trace_heights);
468        if total_interactions > u64::from(self.params.max_interactions) {
469            tracing::info!(
470                "overshoot: instret {:10} | total interactions ({:10}) > max ({:10})",
471                instret,
472                total_interactions,
473                self.params.max_interactions
474            );
475            return Some(SegmentationTrigger::Interactions);
476        }
477
478        None
479    }
480
481    #[inline(always)]
482    pub fn check_and_segment(
483        &mut self,
484        instret: u64,
485        trace_heights: &mut [u32],
486        is_trace_height_constant: &[bool],
487    ) -> bool {
488        let trigger = self.segmentation_trigger(instret, trace_heights, is_trace_height_constant);
489        let should_segment = trigger.is_some();
490
491        #[cfg(feature = "metrics")]
492        if let Some(trigger) = trigger {
493            self.emit_segmentation_trigger_metric(trigger);
494        }
495
496        if should_segment {
497            self.create_segment_from_checkpoint(instret, trace_heights);
498            true
499        } else {
500            false
501        }
502    }
503
504    #[inline(always)]
505    fn create_segment_from_checkpoint(&mut self, instret: u64, trace_heights: &mut [u32]) {
506        let instret_start = self
507            .segments
508            .last()
509            .map_or(0, |s| s.instret_start + s.num_insns);
510
511        let (segment_instret, segment_heights) = if self.checkpoint_instret > instret_start {
512            (
513                self.checkpoint_instret,
514                self.checkpoint_trace_heights.clone(),
515            )
516        } else {
517            let trace_heights_str = trace_heights
518                .iter()
519                .zip(self.params.air_names.iter())
520                .filter(|(&height, _)| height > 0)
521                .map(|(&height, name)| format!("  {name} = {height}"))
522                .collect::<Vec<_>>()
523                .join("\n");
524            tracing::warn!(
525                "No valid checkpoint, creating segment using instret={instret}\ntrace_heights=[\n{trace_heights_str}\n]"
526            );
527            // No valid checkpoint, use current values
528            (instret, trace_heights.to_vec())
529        };
530
531        let num_insns = segment_instret - instret_start;
532        self.create_segment::<false>(instret_start, num_insns, segment_heights);
533    }
534
535    /// Initialize state for a new segment
536    #[inline(always)]
537    pub(crate) fn initialize_segment(
538        &mut self,
539        trace_heights: &mut [u32],
540        is_trace_height_constant: &[bool],
541    ) {
542        // Reset trace heights by subtracting the last segment's heights
543        let last_segment = self.segments.last().unwrap();
544        self.reset_trace_heights(
545            trace_heights,
546            &last_segment.trace_heights,
547            is_trace_height_constant,
548        );
549    }
550
551    /// Resets trace heights by subtracting segment heights
552    #[inline(always)]
553    fn reset_trace_heights(
554        &self,
555        trace_heights: &mut [u32],
556        segment_heights: &[u32],
557        is_trace_height_constant: &[bool],
558    ) {
559        for ((trace_height, &segment_height), &is_trace_height_constant) in trace_heights
560            .iter_mut()
561            .zip(segment_heights.iter())
562            .zip(is_trace_height_constant.iter())
563        {
564            if !is_trace_height_constant {
565                *trace_height = trace_height.checked_sub(segment_height).unwrap();
566            }
567        }
568    }
569
570    /// Updates the checkpoint with current safe state
571    #[inline(always)]
572    pub(crate) fn update_checkpoint(&mut self, instret: u64, trace_heights: &[u32]) {
573        self.checkpoint_trace_heights.copy_from_slice(trace_heights);
574        self.checkpoint_instret = instret;
575    }
576
577    /// Try segment if there is at least one instruction
578    #[inline(always)]
579    pub fn create_final_segment(&mut self, trace_heights: &[u32]) {
580        self.instret += self.params.segment_check_insns - self.instrets_until_check;
581        self.instrets_until_check = self.params.segment_check_insns;
582        let instret_start = self
583            .segments
584            .last()
585            .map_or(0, |s| s.instret_start + s.num_insns);
586
587        let num_insns = self.instret - instret_start;
588        self.create_segment::<true>(instret_start, num_insns, trace_heights.to_vec());
589    }
590
591    /// Push a new segment with logging
592    #[inline(always)]
593    fn create_segment<const IS_FINAL: bool>(
594        &mut self,
595        instret_start: u64,
596        num_insns: u64,
597        trace_heights: Vec<u32>,
598    ) {
599        debug_assert!(
600            num_insns > 0,
601            "Segment should contain at least one instruction"
602        );
603
604        self.log_segment_info::<IS_FINAL>(instret_start, num_insns, &trace_heights);
605        #[cfg(feature = "metrics")]
606        {
607            let segment = self.segments.len().to_string();
608            self.emit_metered_segment_metrics(&segment, &trace_heights);
609            self.emit_metered_air_metrics(&segment, &trace_heights);
610        }
611        self.segments.push(Segment {
612            instret_start,
613            num_insns,
614            trace_heights,
615        });
616    }
617
618    /// Calculate memory utilization: ratio of unpadded memory estimate to padded memory estimate.
619    /// This measures how much of the selected proving memory estimate is useful work vs
620    /// power-of-two trace padding. Note: this inherits memory-related trace-height overestimates.
621    #[inline(always)]
622    fn calculate_memory_utilization(&self, trace_heights: &[u32]) -> f64 {
623        let counts = self.calculate_count_breakdown(trace_heights);
624        let memory = self.calculate_memory_breakdown(&counts);
625        if memory.total == 0 {
626            0.0
627        } else {
628            100.0 * memory.unpadded as f64 / memory.total as f64
629        }
630    }
631
632    /// Log segment information
633    #[inline(always)]
634    fn log_segment_info<const IS_FINAL: bool>(
635        &self,
636        instret_start: u64,
637        num_insns: u64,
638        trace_heights: &[u32],
639    ) {
640        let (max_trace_height, air_name) = self.calculate_max_trace_height_with_name(trace_heights);
641        let (total_memory, main_memory, interaction_memory) =
642            self.calculate_total_memory(trace_heights);
643        let total_interactions = self.calculate_total_interactions(trace_heights);
644        let utilization = self.calculate_memory_utilization(trace_heights);
645
646        let final_marker = if IS_FINAL { " [TERMINATED]" } else { "" };
647
648        tracing::info!(
649            "Segment {:3} | instret {:10} | {:8} instructions | {:5} memory ({:5}, {:5}) | {:10} interactions | {:8} max height ({}) | {:.2}% memory util{}",
650            self.segments.len(),
651            instret_start,
652            num_insns,
653            ByteSize::b(total_memory as u64),
654            ByteSize::b(main_memory as u64),
655            ByteSize::b(interaction_memory as u64),
656            total_interactions,
657            max_trace_height,
658            air_name,
659            utilization,
660            final_marker
661        );
662    }
663}
664
665#[cfg(feature = "metrics")]
666impl SegmentationCtx {
667    fn emit_segmentation_trigger_metric(&self, trigger: SegmentationTrigger) {
668        let segment = self.segments.len().to_string();
669        let reason = trigger.reason();
670        match trigger {
671            SegmentationTrigger::Height { air_id } => {
672                let labels = [
673                    ("segment", segment),
674                    ("reason", reason.to_string()),
675                    ("air_id", air_id.to_string()),
676                    ("air_name", self.params.air_names[air_id].clone()),
677                ];
678                metrics::counter!("segmentation_trigger", &labels).absolute(1);
679            }
680            SegmentationTrigger::Memory | SegmentationTrigger::Interactions => {
681                let labels = [("segment", segment), ("reason", reason.to_string())];
682                metrics::counter!("segmentation_trigger", &labels).absolute(1);
683            }
684        }
685    }
686
687    fn emit_metered_segment_metrics(&self, segment: &str, trace_heights: &[u32]) {
688        let counts = self.calculate_count_breakdown(trace_heights);
689        let memory = self.calculate_memory_breakdown(&counts);
690        let padding = memory.total - memory.unpadded;
691        let estimate = self.params.memory_config.estimate(ProvingMemoryCounts::new(
692            counts.main_unpadded_with_rot + counts.main_padding_with_rot,
693            counts.main_unpadded_no_rot + counts.main_padding_no_rot,
694            counts.interaction_cells_unpadded + counts.interaction_cells_padding,
695            counts.constraint_eval_buffers_unpadded + counts.constraint_eval_buffers_padding,
696        ));
697        let labels = [("segment", segment.to_string())];
698        metrics::counter!("metered_memory_bytes", &labels).absolute(memory.total as u64);
699        metrics::counter!("metered_memory_unpadded_bytes", &labels)
700            .absolute(memory.unpadded as u64);
701        metrics::counter!("metered_memory_padding_bytes", &labels).absolute(padding as u64);
702        metrics::counter!("metered_stacked_matrix_memory_bytes", &labels)
703            .absolute(estimate.stacked_matrix as u64);
704        metrics::counter!("metered_rs_code_matrix_memory_bytes", &labels)
705            .absolute(estimate.rs_code_matrix as u64);
706        metrics::counter!("metered_batch_constraint_memory_bytes", &labels)
707            .absolute(estimate.batch_constraint as u64);
708        metrics::counter!("metered_gkr_memory_bytes", &labels).absolute(estimate.gkr as u64);
709        metrics::counter!("metered_whir_memory_bytes", &labels).absolute(estimate.whir as u64);
710        metrics::counter!("metered_secondary_peak_memory_bytes", &labels)
711            .absolute(estimate.secondary_peak as u64);
712    }
713
714    fn emit_metered_air_metrics(&self, segment: &str, trace_heights: &[u32]) {
715        let memory_config = self.params.memory_config;
716
717        for (air_id, row) in izip!(
718            trace_heights,
719            &self.params.widths,
720            &self.params.interactions,
721            &self.params.constraint_eval_buffers,
722            &self.params.air_names
723        )
724        .enumerate()
725        {
726            let (&height, &width, &interactions, &constraint_eval_buffer, air_name) = row;
727            let padded_height = next_power_of_two_or_zero(height as usize);
728            let unpadded_height = height as usize;
729            let padding_height = padded_height - unpadded_height;
730            if padded_height == 0 {
731                continue;
732            }
733            let labels = [
734                ("air_name", air_name.clone()),
735                ("air_id", air_id.to_string()),
736                ("segment", segment.to_string()),
737            ];
738            let unpadded_cells = unpadded_height * width;
739            let padding_cells = padding_height * width;
740            // One interaction cell is one metered row-interaction slot.
741            let interaction_cells_unpadded = unpadded_height * interactions;
742            let interaction_cells_padding = padding_height * interactions;
743            // One constraint eval cell is one zerocheck round0 intermediate slot.
744            let constraint_eval_cells_unpadded = unpadded_height * constraint_eval_buffer;
745            let constraint_eval_cells_padding = padding_height * constraint_eval_buffer;
746
747            metrics::counter!("metered_rows_unpadded", &labels).absolute(height as u64);
748            metrics::counter!("metered_rows_padding", &labels).absolute(padding_height as u64);
749            metrics::counter!("metered_main_cells_unpadded", &labels)
750                .absolute(unpadded_cells as u64);
751            metrics::counter!("metered_main_cells_padding", &labels).absolute(padding_cells as u64);
752            metrics::counter!("metered_interaction_cells_unpadded", &labels)
753                .absolute(interaction_cells_unpadded as u64);
754            metrics::counter!("metered_interaction_cells_padding", &labels)
755                .absolute(interaction_cells_padding as u64);
756            metrics::counter!("metered_constraint_eval_cells_unpadded", &labels)
757                .absolute(constraint_eval_cells_unpadded as u64);
758            metrics::counter!("metered_constraint_eval_cells_padding", &labels)
759                .absolute(constraint_eval_cells_padding as u64);
760            metrics::counter!("metered_main_memory_unpadded_bytes", &labels)
761                .absolute(memory_config.main_memory_bytes(unpadded_cells) as u64);
762            metrics::counter!("metered_main_memory_padding_bytes", &labels)
763                .absolute(memory_config.main_memory_bytes(padding_cells) as u64);
764        }
765    }
766}
openvm_circuit/arch/execution_mode/metered/segment_ctx.rs

openvm_circuit/arch/execution_mode/metered/
segment_ctx.rs
