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

openvm_circuit/arch/execution_mode/metered/
ctx.rs
