openvm_circuit/arch/execution_mode/metered/
ctx.rs

1use getset::{Getters, Setters, WithSetters};
2use itertools::Itertools;
3use openvm_instructions::riscv::{RV32_IMM_AS, RV32_REGISTER_AS};
4use openvm_stark_backend::memory_metering::ProvingMemoryConfig;
5
6use super::{
7    memory_ctx::MemoryCtx,
8    segment_ctx::{Segment, SegmentationCtx, SegmentationLimits},
9};
10use crate::{
11    arch::{
12        execution_mode::{ExecutionCtxTrait, MeteredExecutionCtxTrait},
13        SystemConfig, VmExecState, BOUNDARY_AIR_ID, MERKLE_AIR_ID,
14    },
15    system::memory::online::GuestMemory,
16};
17
18pub const DEFAULT_PAGE_BITS: usize = 6;
19
20#[derive(Clone, Debug, Getters, Setters, WithSetters)]
21pub struct MeteredCtx<const PAGE_BITS: usize = DEFAULT_PAGE_BITS> {
22    pub trace_heights: Vec<u32>,
23    pub is_trace_height_constant: Vec<bool>,
24    pub memory_ctx: MemoryCtx<PAGE_BITS>,
25    pub segmentation_ctx: SegmentationCtx,
26    #[getset(get = "pub", set = "pub", set_with = "pub")]
27    suspend_on_segment: bool,
28}
29
30pub struct MeteredCtxInputs<'a> {
31    pub constant_trace_heights: &'a [Option<usize>],
32    pub air_names: &'a [String],
33    pub widths: &'a [usize],
34    pub interactions: &'a [usize],
35    pub need_rot: &'a [bool],
36    pub constraint_eval_buffers: &'a [usize],
37    pub segmentation_limits: SegmentationLimits,
38}
39
40impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
41    // Note: prefer to use `build_metered_ctx` in `VmExecutor` or `VirtualMachine`.
42    pub fn new(
43        inputs: MeteredCtxInputs<'_>,
44        config: &SystemConfig,
45        memory_config: ProvingMemoryConfig,
46    ) -> Self {
47        let (trace_heights, is_trace_height_constant): (Vec<u32>, Vec<bool>) = inputs
48            .constant_trace_heights
49            .iter()
50            .map(|&constant_height| {
51                if let Some(height) = constant_height {
52                    (height as u32, true)
53                } else {
54                    (0, false)
55                }
56            })
57            .unzip();
58
59        let segmentation_ctx = SegmentationCtx::new(
60            inputs.air_names.to_vec(),
61            inputs.widths.to_vec(),
62            inputs.interactions.to_vec(),
63            inputs.need_rot.to_vec(),
64            inputs.constraint_eval_buffers.to_vec(),
65            inputs.segmentation_limits,
66            memory_config,
67        );
68        let memory_ctx = MemoryCtx::new(config, segmentation_ctx.segment_check_insns());
69
70        // Assert that the indices are correct
71        let air_names = segmentation_ctx.air_names();
72        debug_assert!(
73            air_names[BOUNDARY_AIR_ID].contains("Boundary"),
74            "air_name={}",
75            air_names[BOUNDARY_AIR_ID]
76        );
77        debug_assert!(
78            air_names[MERKLE_AIR_ID].contains("Merkle"),
79            "air_name={}",
80            air_names[MERKLE_AIR_ID]
81        );
82        let mut ctx = Self {
83            trace_heights,
84            is_trace_height_constant,
85            memory_ctx,
86            segmentation_ctx,
87            suspend_on_segment: false,
88        };
89
90        // Add merkle height contributions for all registers
91        ctx.memory_ctx.add_register_merkle_heights();
92        ctx.memory_ctx
93            .lazy_update_boundary_heights(&mut ctx.trace_heights);
94
95        ctx
96    }
97
98    pub fn with_max_memory(mut self, max_memory: usize) -> Self {
99        self.segmentation_ctx.set_max_memory(max_memory);
100        self
101    }
102
103    pub fn segments(&self) -> &[Segment] {
104        &self.segmentation_ctx.segments
105    }
106
107    pub fn into_segments(self) -> Vec<Segment> {
108        self.segmentation_ctx.segments
109    }
110
111    #[inline(always)]
112    pub fn check_and_segment(&mut self) -> bool {
113        // We track the segmentation check by instrets_until_check instead of instret in order to
114        // save a register in AOT mode.
115        if self.segmentation_ctx.instrets_until_check > 0 {
116            return false;
117        }
118        let segment_check_insns = self.segmentation_ctx.segment_check_insns();
119        self.segmentation_ctx.instrets_until_check = segment_check_insns;
120        self.segmentation_ctx.instret += segment_check_insns;
121
122        self.memory_ctx
123            .lazy_update_boundary_heights(&mut self.trace_heights);
124        let did_segment = self.segmentation_ctx.check_and_segment(
125            self.segmentation_ctx.instret,
126            &mut self.trace_heights,
127            &self.is_trace_height_constant,
128        );
129
130        if did_segment {
131            // Initialize contexts for new segment
132            self.segmentation_ctx
133                .initialize_segment(&mut self.trace_heights, &self.is_trace_height_constant);
134            self.memory_ctx.initialize_segment(&mut self.trace_heights);
135
136            // Check if the new segment is within limits
137            if self.segmentation_ctx.should_segment(
138                self.segmentation_ctx.instret,
139                &self.trace_heights,
140                &self.is_trace_height_constant,
141            ) {
142                let trace_heights_str = self
143                    .trace_heights
144                    .iter()
145                    .zip(self.segmentation_ctx.air_names().iter())
146                    .filter(|(&height, _)| height > 0)
147                    .map(|(&height, name)| format!("  {name} = {height}"))
148                    .collect::<Vec<_>>()
149                    .join("\n");
150                tracing::warn!(
151                    "Segment initialized with heights that exceed limits\n\
152                     instret={}\n\
153                     trace_heights=[\n{}\n]",
154                    self.segmentation_ctx.instret,
155                    trace_heights_str
156                );
157            }
158        }
159
160        // Update checkpoints
161        self.segmentation_ctx
162            .update_checkpoint(self.segmentation_ctx.instret, &self.trace_heights);
163        self.memory_ctx.update_checkpoint();
164
165        did_segment
166    }
167
168    #[allow(dead_code)]
169    pub fn print_segment(&self) {
170        println!("{}", "-".repeat(80));
171        println!("Segment {}", self.segmentation_ctx.segments.len() - 1);
172        println!("{}", "-".repeat(80));
173        println!("{:>10} {:>10} {:<30}", "Width", "Height", "Air Name");
174        println!("{}", "-".repeat(80));
175        for ((&width, &height), air_name) in self
176            .segmentation_ctx
177            .widths()
178            .iter()
179            .zip_eq(self.trace_heights.iter())
180            .zip_eq(self.segmentation_ctx.air_names().iter())
181        {
182            println!("{:>10} {:>10} {:<30}", width, height, air_name.as_str());
183        }
184    }
185}
186
187impl<const PAGE_BITS: usize> ExecutionCtxTrait for MeteredCtx<PAGE_BITS> {
188    #[inline(always)]
189    fn on_memory_operation(&mut self, address_space: u32, ptr: u32, size: u32) {
190        debug_assert!(
191            address_space != RV32_IMM_AS,
192            "address space must not be immediate"
193        );
194        debug_assert!(size > 0, "size must be greater than 0, got {size}");
195        debug_assert!(
196            size.is_power_of_two(),
197            "size must be a power of 2, got {size}"
198        );
199
200        // Handle merkle tree updates
201        if address_space != RV32_REGISTER_AS {
202            self.memory_ctx
203                .update_boundary_merkle_heights(address_space, ptr, size);
204        }
205    }
206
207    #[inline(always)]
208    fn should_suspend<F>(exec_state: &mut VmExecState<F, GuestMemory, Self>) -> bool {
209        // ATTENTION: Please make sure to update the corresponding logic in the
210        // `asm_bridge` crate and `aot.rs`` when you change this function.
211        // If `segment_suspend` is set, suspend when a segment is determined (but the VM state might
212        // be after the segment boundary because the segment happens in the previous checkpoint).
213        // Otherwise, execute until termination.
214        if exec_state.ctx.check_and_segment() && exec_state.ctx.suspend_on_segment {
215            true
216        } else {
217            exec_state.ctx.segmentation_ctx.instrets_until_check -= 1;
218            false
219        }
220    }
221
222    #[inline(always)]
223    fn on_terminate<F>(exec_state: &mut VmExecState<F, GuestMemory, Self>) {
224        exec_state
225            .ctx
226            .memory_ctx
227            .lazy_update_boundary_heights(&mut exec_state.ctx.trace_heights);
228        exec_state
229            .ctx
230            .segmentation_ctx
231            .create_final_segment(&exec_state.ctx.trace_heights);
232    }
233}
234
235impl<const PAGE_BITS: usize> MeteredExecutionCtxTrait for MeteredCtx<PAGE_BITS> {
236    #[inline(always)]
237    fn on_height_change(&mut self, chip_idx: usize, height_delta: u32) {
238        debug_assert!(
239            chip_idx < self.trace_heights.len(),
240            "chip_idx out of bounds"
241        );
242        // SAFETY: chip_idx is created in executor_idx_to_air_idx and is always within bounds
243        unsafe {
244            *self.trace_heights.get_unchecked_mut(chip_idx) = self
245                .trace_heights
246                .get_unchecked(chip_idx)
247                .wrapping_add(height_delta);
248        }
249    }
250}
openvm_circuit/arch/execution_mode/metered/ctx.rs

openvm_circuit/arch/execution_mode/metered/
ctx.rs
