openvm_circuit/arch/execution_mode/
metered_cost.rs

1use std::num::NonZero;
2
3use getset::WithSetters;
4use openvm_instructions::riscv::RV32_IMM_AS;
5
6use crate::{
7    arch::{
8        execution_mode::metered::segment_ctx::DEFAULT_MAX_CELLS as DEFAULT_SEGMENT_MAX_CELLS,
9        ExecutionCtxTrait, MeteredExecutionCtxTrait, SystemConfig, VmExecState,
10    },
11    system::memory::online::GuestMemory,
12};
13
14const DEFAULT_MAX_SEGMENTS: u64 = 100;
15pub const DEFAULT_MAX_COST: u64 = DEFAULT_MAX_SEGMENTS * DEFAULT_SEGMENT_MAX_CELLS as u64;
16
17#[derive(Clone, Debug)]
18pub struct AccessAdapterCtx {
19    min_block_size_bits: Vec<u8>,
20    idx_offset: usize,
21}
22
23impl AccessAdapterCtx {
24    pub fn new(config: &SystemConfig) -> Self {
25        Self {
26            min_block_size_bits: config.memory_config.min_block_size_bits(),
27            idx_offset: config.access_adapter_air_id_offset(),
28        }
29    }
30
31    #[inline(always)]
32    pub fn update_cells(
33        &self,
34        cost: &mut u64,
35        address_space: u32,
36        size_bits: u32,
37        widths: &[usize],
38    ) {
39        debug_assert!((address_space as usize) < self.min_block_size_bits.len());
40
41        // SAFETY: address_space passed is usually a hardcoded constant or derived from an
42        // Instruction where it is bounds checked before passing
43        let align_bits = unsafe {
44            *self
45                .min_block_size_bits
46                .get_unchecked(address_space as usize)
47        };
48        debug_assert!(
49            align_bits as u32 <= size_bits,
50            "align_bits ({align_bits}) must be <= size_bits ({size_bits})"
51        );
52
53        for adapter_bits in (align_bits as u32 + 1..=size_bits).rev() {
54            let adapter_idx = self.idx_offset + adapter_bits as usize - 1;
55            debug_assert!(adapter_idx < widths.len());
56            // SAFETY: widths is initialized taking access adapters into account
57            let width = unsafe { *widths.get_unchecked(adapter_idx) };
58            let height_delta = 1 << (size_bits - adapter_bits + 1);
59            *cost += (height_delta as u64) * (width as u64);
60        }
61    }
62}
63
64#[derive(Clone, Debug, WithSetters)]
65pub struct MeteredCostCtx {
66    pub widths: Vec<usize>,
67    pub access_adapter_ctx: AccessAdapterCtx,
68    #[getset(set_with = "pub")]
69    pub max_execution_cost: u64,
70    // Cost is number of trace cells (height * width)
71    pub cost: u64,
72    /// To measure instructions/s
73    pub instret: u64,
74}
75
76impl MeteredCostCtx {
77    pub fn new(widths: Vec<usize>, config: &SystemConfig) -> Self {
78        let access_adapter_ctx = AccessAdapterCtx::new(config);
79        Self {
80            widths,
81            access_adapter_ctx,
82            max_execution_cost: DEFAULT_MAX_COST,
83            cost: 0,
84            instret: 0,
85        }
86    }
87
88    #[cold]
89    fn panic_cost_exceeded(&self) -> ! {
90        panic!(
91            "Execution cost {} exceeded maximum allowed cost of {}",
92            self.cost,
93            2 * DEFAULT_MAX_COST
94        );
95    }
96}
97
98impl ExecutionCtxTrait for MeteredCostCtx {
99    #[inline(always)]
100    fn on_memory_operation(&mut self, address_space: u32, _ptr: u32, size: u32) {
101        debug_assert!(
102            address_space != RV32_IMM_AS,
103            "address space must not be immediate"
104        );
105        debug_assert!(size > 0, "size must be greater than 0, got {size}");
106        debug_assert!(
107            size.is_power_of_two(),
108            "size must be a power of 2, got {size}"
109        );
110        // Prevent unbounded memory accesses per instruction
111        if self.cost > 2 * std::cmp::max(self.max_execution_cost, DEFAULT_MAX_COST) {
112            self.panic_cost_exceeded();
113        }
114
115        // Handle access adapter updates
116        // SAFETY: size passed is always a non-zero power of 2
117        let size_bits = unsafe { NonZero::new_unchecked(size).ilog2() };
118        self.access_adapter_ctx.update_cells(
119            &mut self.cost,
120            address_space,
121            size_bits,
122            &self.widths,
123        );
124    }
125
126    #[inline(always)]
127    fn should_suspend<F>(exec_state: &mut VmExecState<F, GuestMemory, Self>) -> bool {
128        if exec_state.ctx.cost > exec_state.ctx.max_execution_cost {
129            true
130        } else {
131            exec_state.ctx.instret += 1;
132            false
133        }
134    }
135}
136
137impl MeteredExecutionCtxTrait for MeteredCostCtx {
138    #[inline(always)]
139    fn on_height_change(&mut self, chip_idx: usize, height_delta: u32) {
140        debug_assert!(chip_idx < self.widths.len(), "chip_idx out of bounds");
141        // SAFETY: chip_idx is created in executor_idx_to_air_idx and is always within bounds
142        let width = unsafe { *self.widths.get_unchecked(chip_idx) };
143        self.cost += (height_delta as u64) * (width as u64);
144    }
145}