openvm_stark_backend/
memory_metering.rs

1//! Memory estimates for proving.
2
3use std::{cmp::max, mem::size_of};
4
5use crate::{StarkProtocolConfig, SystemParams};
6
7/// Fixed interaction scratch not proportional to interaction cells.
8pub const INTERACTION_MEMORY_OVERHEAD: usize = 2 << 20;
9
10/// Cell counts for a proving memory estimate.
11///
12/// `main_cells_*` are `Σ(padded_height * width)` in base-field cells, split by whether a trace
13/// opens next-row rotations. `interaction_cells` is the metered row-interaction slot count after
14/// power-of-two trace padding.
15#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
16pub struct ProvingMemoryCounts {
17    /// Main trace cells for AIRs that open next-row rotations after power-of-two padding.
18    pub main_cells_with_rot: usize,
19    /// Main trace cells for AIRs without next-row rotations after power-of-two padding.
20    pub main_cells_without_rot: usize,
21    /// Metered row-interaction slots after power-of-two padding.
22    pub interaction_cells: usize,
23}
24
25impl ProvingMemoryCounts {
26    pub const fn new(
27        main_cells_with_rot: usize,
28        main_cells_without_rot: usize,
29        interaction_cells: usize,
30    ) -> Self {
31        Self {
32            main_cells_with_rot,
33            main_cells_without_rot,
34            interaction_cells,
35        }
36    }
37
38    #[inline]
39    pub const fn main_cells(&self) -> usize {
40        self.main_cells_with_rot + self.main_cells_without_rot
41    }
42}
43
44/// Estimated memory components, in bytes.
45#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
46pub struct ProvingMemoryEstimate {
47    /// Selected peak estimate.
48    pub total: usize,
49    /// Cached main trace data.
50    pub main: usize,
51    /// Reed-Solomon code matrix for main traces.
52    pub rs_code_matrix: usize,
53    /// Secondary main-trace buffers after PCS opening.
54    pub main_secondary: usize,
55    /// Interaction GKR buffers plus fixed interaction overhead.
56    pub interaction: usize,
57    /// Peak among secondary phases, excluding cached main trace data.
58    pub secondary_peak: usize,
59}
60
61/// Configuration for proving memory estimates.
62#[derive(Clone, Copy, Debug, PartialEq)]
63pub struct ProvingMemoryConfig {
64    /// Size of one base-field element in bytes.
65    pub base_field_size: usize,
66    /// Degree of the extension field over the base field.
67    pub extension_degree: usize,
68    /// `-log_2` of the rate for the initial Reed-Solomon code.
69    pub log_blowup: usize,
70    /// `log_2` of the univariate skip domain.
71    pub l_skip: usize,
72    /// Maximum constraint degree across AIR and interaction constraints.
73    pub max_constraint_degree: usize,
74    /// Whether the prover keeps the Reed-Solomon code matrix cached after `stacked_commit`.
75    pub cache_rs_code_matrix: bool,
76}
77
78impl ProvingMemoryConfig {
79    pub fn from_protocol_config<SC: StarkProtocolConfig>(
80        config: &SC,
81        cache_rs_code_matrix: bool,
82    ) -> Self {
83        Self::from_params::<SC::F>(config.params(), SC::D_EF, cache_rs_code_matrix)
84    }
85
86    fn from_params<F>(
87        params: &SystemParams,
88        extension_degree: usize,
89        cache_rs_code_matrix: bool,
90    ) -> Self {
91        Self {
92            base_field_size: size_of::<F>(),
93            extension_degree,
94            log_blowup: params.log_blowup,
95            l_skip: params.l_skip,
96            max_constraint_degree: params.max_constraint_degree,
97            cache_rs_code_matrix,
98        }
99    }
100
101    #[inline]
102    pub fn main_memory_bytes(&self, main_cells: usize) -> usize {
103        main_cells * self.base_field_size
104    }
105
106    #[inline]
107    pub fn rs_code_matrix_memory_bytes(&self, main_cells: usize) -> usize {
108        main_cells * (1usize << self.log_blowup) * self.base_field_size
109    }
110
111    #[inline]
112    pub fn main_cell_secondary_weight(&self) -> f64 {
113        default_main_cell_secondary_weight(
114            self.extension_degree,
115            self.l_skip,
116            self.max_constraint_degree,
117        )
118    }
119
120    /// Secondary memory for `main_cells = Σ(padded_height * width)`.
121    ///
122    /// ```text
123    /// mat_eval_bytes = (padded_height / 2^l_skip) * ((1 + need_rot) * width) * sizeof(EF)
124    ///                = (1 + need_rot) * main_cells * D_EF / 2^l_skip * sizeof(F)
125    ///
126    /// interp_bytes      = (constraint_degree / 2) * mat_eval_bytes
127    /// secondary_bytes   = (1 + constraint_degree / 2) * mat_eval_bytes
128    /// secondary_weight  = (1 + constraint_degree / 2) * D_EF / 2^l_skip
129    /// ```
130    ///
131    /// AIRs with `need_rot = true` open two PCS cells per column.
132    #[inline]
133    pub fn main_secondary_memory_bytes_for_rot(&self, main_cells: usize, need_rot: bool) -> usize {
134        let main_cell_secondary_weight = self.main_cell_secondary_weight();
135        let weight = if need_rot {
136            2.0 * main_cell_secondary_weight
137        } else {
138            main_cell_secondary_weight
139        };
140        ceil_weighted_bytes(main_cells, self.base_field_size, weight)
141    }
142
143    #[inline]
144    pub fn main_secondary_memory_bytes(&self, counts: ProvingMemoryCounts) -> usize {
145        self.main_secondary_memory_bytes_for_rot(counts.main_cells_with_rot, true)
146            + self.main_secondary_memory_bytes_for_rot(counts.main_cells_without_rot, false)
147    }
148
149    #[inline]
150    pub fn interaction_memory_bytes_without_overhead(&self, interaction_cells: usize) -> usize {
151        ceil_weighted_bytes(
152            interaction_cells,
153            self.base_field_size,
154            default_interaction_cell_weight(self.extension_degree),
155        )
156    }
157
158    #[inline]
159    pub fn interaction_memory_bytes(&self, interaction_cells: usize) -> usize {
160        self.interaction_memory_bytes_without_overhead(interaction_cells)
161            + INTERACTION_MEMORY_OVERHEAD
162    }
163
164    /// Convert main trace cells and interaction cells to memory bytes.
165    ///
166    /// ```text
167    /// main_cells       = main_cells_with_rot + main_cells_without_rot
168    /// main             = main_cells * sizeof(F)
169    /// rs_code_matrix   = main_cells * 2^log_blowup * sizeof(F)
170    /// main_secondary   = sizeof(F) *
171    ///                    (2 * main_cell_secondary_weight() * main_cells_with_rot
172    ///                     + main_cell_secondary_weight() * main_cells_without_rot)
173    /// interaction      = sizeof(F) * interaction_cell_weight * interaction_cells
174    ///                    + INTERACTION_MEMORY_OVERHEAD
175    /// ```
176    ///
177    /// Cached RS code matrix:
178    ///
179    /// ```text
180    /// total = main + rs_code_matrix + max(main_secondary, interaction)
181    /// ```
182    ///
183    /// Dropped RS code matrix:
184    ///
185    /// ```text
186    /// total = main + max(rs_code_matrix, main_secondary, interaction)
187    /// ```
188    #[inline]
189    pub fn estimate(&self, counts: ProvingMemoryCounts) -> ProvingMemoryEstimate {
190        let main_cells = counts.main_cells();
191        let main = self.main_memory_bytes(main_cells);
192        let rs_code_matrix = self.rs_code_matrix_memory_bytes(main_cells);
193        let main_secondary = self.main_secondary_memory_bytes(counts);
194        let interaction = self.interaction_memory_bytes(counts.interaction_cells);
195        let secondary_peak = if self.cache_rs_code_matrix {
196            rs_code_matrix + max(main_secondary, interaction)
197        } else {
198            max(rs_code_matrix, max(main_secondary, interaction))
199        };
200
201        ProvingMemoryEstimate {
202            total: main + secondary_peak,
203            main,
204            rs_code_matrix,
205            main_secondary,
206            interaction,
207            secondary_peak,
208        }
209    }
210}
211
212fn default_main_cell_secondary_weight(
213    extension_degree: usize,
214    l_skip: usize,
215    max_constraint_degree: usize,
216) -> f64 {
217    (1.0 + max_constraint_degree as f64 / 2.0) * extension_degree as f64 / (1usize << l_skip) as f64
218}
219
220/// ```
221/// leaf_weight = 2 * extension_degree
222///
223/// work_buffer    <= logical_len / 16
224/// tmp_block_sums ~= logical_len / 256
225/// logical_len    <= 2 * real_len
226///
227/// interaction_weight = leaf_weight * (1 + 2 * (1 / 16 + 1 / 256))
228/// ```
229fn default_interaction_cell_weight(extension_degree: usize) -> f64 {
230    (2 * extension_degree) as f64 * (1.0 + 2.0 * (1.0 / 16.0 + 1.0 / 256.0))
231}
232
233/// `ceil(cell_count * base_field_size * weight)`
234fn ceil_weighted_bytes(cell_count: usize, base_field_size: usize, weight: f64) -> usize {
235    ((cell_count * base_field_size) as f64 * weight).ceil() as usize
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::test_utils::default_test_params_small;
242
243    fn test_memory_config() -> ProvingMemoryConfig {
244        let params = default_test_params_small();
245        ProvingMemoryConfig::from_params::<u32>(&params, 4, true)
246    }
247
248    #[test]
249    fn dropped_rs_code_matrix_is_phase_disjoint() {
250        let params = default_test_params_small();
251        let config = ProvingMemoryConfig::from_params::<u32>(&params, 4, false);
252        let counts = ProvingMemoryCounts::new(10, 20, 5);
253
254        let estimate = config.estimate(counts);
255
256        assert_eq!(estimate.main, 30 * 4);
257        assert_eq!(estimate.rs_code_matrix, 30 * 2 * 4);
258        assert_eq!(estimate.total, estimate.main + estimate.secondary_peak);
259        assert_eq!(
260            estimate.secondary_peak,
261            max(
262                estimate.rs_code_matrix,
263                max(estimate.main_secondary, estimate.interaction)
264            )
265        );
266    }
267
268    #[test]
269    fn cached_rs_code_matrix_is_additive() {
270        let config = test_memory_config();
271        let counts = ProvingMemoryCounts::new(10, 20, 5);
272
273        let estimate = config.estimate(counts);
274
275        assert_eq!(
276            estimate.secondary_peak,
277            estimate.rs_code_matrix + max(estimate.main_secondary, estimate.interaction)
278        );
279    }
280}
openvm_stark_backend/memory_metering.rs

openvm_stark_backend/
memory_metering.rs
