memory_profile/
memory_profile.rs

1//! End-to-end prover memory profile.
2//!
3//! Proves either the canonical benchmark circuit (`kimchi::bench::BenchmarkCtx`)
4//! or a serialised mina circuit fixture (the same `kimchi_inputs_*.ser` files
5//! the `proof_criterion_mina` bench consumes) under a counting global
6//! allocator and reports, for the proof-creation window: total bytes
7//! allocated, allocation count, peak live bytes (and its delta over the bytes
8//! live when the window opened), plus jemalloc's peak resident set sampled by
9//! a background thread. The process exists for this one measurement, so the
10//! counters are exact for a deterministic workload — a single run is the
11//! answer.
12//!
13//! Usage:
14//!
15//! ```text
16//! cargo run --release -p kimchi --bin memory_profile --features diagnostics -- synthetic [--srs-log2 16]
17//! cargo run --release -p kimchi --bin memory_profile --features diagnostics -- fixture <kimchi_inputs_CURVE_SEED.ser>
18//! ```
19//!
20//! `synthetic` proves the benchmark circuit at the given domain/SRS size;
21//! `fixture` proves a mina fixture, whose curve and seed are parsed from the
22//! filename exactly as in `proof_criterion_mina`. Output is a single JSON
23//! object on stdout (byte counts, not MB), meant to be collected across
24//! fixtures and diffed against a baseline run — see
25//! `scripts/memory-profile-mina-circuits.sh` and
26//! `scripts/memory-profile-diff.py`.
27//! One fixture per invocation: jemalloc retains pages across proofs, so
28//! `peak_resident` is only trustworthy for the first proof in a process.
29//! The resident-set sampler interval is `KIMCHI_MEMORY_PROFILE_SAMPLE_MS`
30//! (default 25).
31
32use ark_ff::PrimeField;
33use clap::Parser;
34use groupmap::GroupMap;
35use kimchi::{
36    bench::{
37        bench_arguments_from_file, BaseSpongePallas, BaseSpongeVesta, BenchmarkCtx,
38        ScalarSpongePallas, ScalarSpongeVesta,
39    },
40    curve::KimchiCurve,
41    plonk_sponge::FrSponge,
42    proof::ProverProof,
43};
44use mina_curves::{
45    named::NamedCurve,
46    pasta::{Pallas, Vesta},
47};
48use mina_poseidon::{pasta::FULL_ROUNDS, FqSponge};
49use poly_commitment::ipa::OpeningProof;
50use std::time::Instant;
51
52// A counting wrapper around jemalloc. jemalloc-ctl statistics still observe
53// the real allocator, while the counters capture what jemalloc cannot report
54// over a window: total bytes allocated, allocation count, and peak live
55// bytes. Counting is unconditional: this binary is the measurement.
56mod counting_alloc {
57    use core::sync::atomic::{AtomicUsize, Ordering::Relaxed};
58    use std::alloc::{GlobalAlloc, Layout};
59    use tikv_jemallocator::Jemalloc;
60
61    pub static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
62    pub static COUNT: AtomicUsize = AtomicUsize::new(0);
63    pub static CURRENT: AtomicUsize = AtomicUsize::new(0);
64    pub static PEAK: AtomicUsize = AtomicUsize::new(0);
65
66    fn record_alloc(size: usize) {
67        ALLOCATED.fetch_add(size, Relaxed);
68        COUNT.fetch_add(1, Relaxed);
69        let live = CURRENT.fetch_add(size, Relaxed) + size;
70        PEAK.fetch_max(live, Relaxed);
71    }
72
73    struct CountingAlloc;
74
75    unsafe impl GlobalAlloc for CountingAlloc {
76        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
77            let p = Jemalloc.alloc(layout);
78            if !p.is_null() {
79                record_alloc(layout.size());
80            }
81            p
82        }
83
84        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
85            let p = Jemalloc.alloc_zeroed(layout);
86            if !p.is_null() {
87                record_alloc(layout.size());
88            }
89            p
90        }
91
92        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
93            let p = Jemalloc.realloc(ptr, layout, new_size);
94            if !p.is_null() {
95                CURRENT.fetch_sub(layout.size(), Relaxed);
96                record_alloc(new_size);
97            }
98            p
99        }
100
101        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
102            CURRENT.fetch_sub(layout.size(), Relaxed);
103            Jemalloc.dealloc(ptr, layout)
104        }
105    }
106
107    #[global_allocator]
108    static GLOBAL: CountingAlloc = CountingAlloc;
109}
110
111/// An allocation profile over a window of execution.
112mod mem_profile {
113    use super::counting_alloc::{ALLOCATED, COUNT, CURRENT, PEAK};
114    use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed};
115
116    static STOP: AtomicBool = AtomicBool::new(false);
117    static PEAK_RESIDENT: AtomicUsize = AtomicUsize::new(0);
118
119    pub struct Window {
120        live_before: usize,
121        sampler: std::thread::JoinHandle<()>,
122    }
123
124    pub fn start() -> Window {
125        let sample_ms: u64 = match std::env::var("KIMCHI_MEMORY_PROFILE_SAMPLE_MS") {
126            Err(_) => 25,
127            Ok(s) => s.parse().ok().filter(|ms| *ms >= 1).unwrap_or_else(|| {
128                panic!("KIMCHI_MEMORY_PROFILE_SAMPLE_MS must be a positive integer (ms), got {s:?}")
129            }),
130        };
131        let sampler = std::thread::spawn(move || {
132            use tikv_jemalloc_ctl::{epoch, stats};
133            while !STOP.load(Relaxed) {
134                if epoch::advance().is_ok() {
135                    if let Ok(r) = stats::resident::read() {
136                        PEAK_RESIDENT.fetch_max(r, Relaxed);
137                    }
138                }
139                std::thread::sleep(core::time::Duration::from_millis(sample_ms));
140            }
141        });
142
143        // Reset the window counters after spawning the sampler, so the
144        // spawn's own allocations stay out of the profile.
145        let live_before = CURRENT.load(Relaxed);
146        ALLOCATED.store(0, Relaxed);
147        COUNT.store(0, Relaxed);
148        PEAK.store(live_before, Relaxed);
149        Window {
150            live_before,
151            sampler,
152        }
153    }
154
155    #[derive(serde::Serialize)]
156    pub struct Metrics {
157        #[serde(rename = "allocated_bytes")]
158        pub allocated: usize,
159        pub allocs: usize,
160        #[serde(rename = "peak_live_bytes")]
161        pub peak_live: usize,
162        #[serde(rename = "peak_delta_bytes")]
163        pub peak_delta: usize,
164        #[serde(rename = "peak_resident_bytes")]
165        pub peak_resident: usize,
166    }
167
168    impl Window {
169        pub fn finish(self) -> Metrics {
170            STOP.store(true, Relaxed);
171            let _ = self.sampler.join();
172            let peak_live = PEAK.load(Relaxed);
173            Metrics {
174                allocated: ALLOCATED.load(Relaxed),
175                allocs: COUNT.load(Relaxed),
176                peak_live,
177                peak_delta: peak_live.saturating_sub(self.live_before),
178                peak_resident: PEAK_RESIDENT.load(Relaxed),
179            }
180        }
181    }
182}
183
184/// One profiled proof, as the JSON object written to stdout.
185#[derive(serde::Serialize)]
186struct Report<'a> {
187    workload: &'static str,
188    #[serde(skip_serializing_if = "Option::is_none")]
189    srs_log2: Option<u32>,
190    #[serde(skip_serializing_if = "Option::is_none")]
191    curve: Option<&'a str>,
192    #[serde(skip_serializing_if = "Option::is_none")]
193    seed: Option<&'a str>,
194    #[serde(skip_serializing_if = "Option::is_none")]
195    domain_log2: Option<u32>,
196    setup_ms: u128,
197    prove_ms: u128,
198    #[serde(flatten)]
199    metrics: mem_profile::Metrics,
200}
201
202impl std::fmt::Display for Report<'_> {
203    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
204        f.write_str(&serde_json::to_string(self).map_err(|_| std::fmt::Error)?)
205    }
206}
207
208/// A serialised mina circuit fixture, addressed by the filename convention
209/// `kimchi_inputs_CURVE_SEED.ser` that `proof_criterion_mina` also uses.
210#[derive(Clone)]
211struct Fixture {
212    path: String,
213    curve: String,
214    seed: String,
215}
216
217impl std::str::FromStr for Fixture {
218    type Err = String;
219
220    fn from_str(path: &str) -> Result<Self, Self::Err> {
221        let (curve, seed) = path
222            .split('/')
223            .next_back()
224            .unwrap_or(path)
225            .strip_prefix("kimchi_inputs_")
226            .and_then(|s| s.strip_suffix(".ser"))
227            .and_then(|s| s.split_once('_'))
228            .ok_or_else(|| {
229                format!(
230                    "fixture filename must look like kimchi_inputs_CURVE_SEED.ser, got {path:?}"
231                )
232            })?;
233        Ok(Fixture {
234            path: path.to_string(),
235            curve: curve.to_string(),
236            seed: seed.to_string(),
237        })
238    }
239}
240
241fn profile_synthetic(srs_log2: u32) {
242    let setup_start = Instant::now();
243    let ctx = BenchmarkCtx::new(srs_log2);
244    let setup_ms = setup_start.elapsed().as_millis();
245
246    let window = mem_profile::start();
247    let prove_start = Instant::now();
248    let proof_and_public = ctx.create_proof();
249    let prove_ms = prove_start.elapsed().as_millis();
250    let metrics = window.finish();
251    std::hint::black_box(proof_and_public);
252
253    println!(
254        "{}",
255        Report {
256            workload: "synthetic",
257            srs_log2: Some(srs_log2),
258            curve: None,
259            seed: None,
260            domain_log2: None,
261            setup_ms,
262            prove_ms,
263            metrics,
264        }
265    );
266}
267
268fn profile_fixture_curve<G, BaseSponge, ScalarSponge>(fixture: &Fixture)
269where
270    G: KimchiCurve<FULL_ROUNDS> + NamedCurve,
271    G::BaseField: PrimeField,
272    BaseSponge: Clone + FqSponge<G::BaseField, G, G::ScalarField, FULL_ROUNDS>,
273    ScalarSponge: FrSponge<G::ScalarField>
274        + From<&'static mina_poseidon::poseidon::ArithmeticSpongeParams<G::ScalarField, FULL_ROUNDS>>,
275{
276    let setup_start = Instant::now();
277    let srs = poly_commitment::precomputed_srs::get_srs_test();
278    let (index, witness, runtime_tables, prev) =
279        bench_arguments_from_file::<FULL_ROUNDS, G, BaseSponge>(srs, fixture.path.clone());
280    let group_map = GroupMap::<_>::setup();
281    let domain_log2 = index.cs.domain.d1.size.trailing_zeros();
282    let setup_ms = setup_start.elapsed().as_millis();
283
284    let window = mem_profile::start();
285    let prove_start = Instant::now();
286    let proof = ProverProof::<G, OpeningProof<G, FULL_ROUNDS>, FULL_ROUNDS>::create_recursive::<
287        BaseSponge,
288        ScalarSponge,
289        _,
290    >(
291        &group_map,
292        witness,
293        &runtime_tables,
294        &index,
295        prev,
296        None,
297        &mut rand::rngs::OsRng,
298    )
299    .expect("proof creation failed: the fixture no longer satisfies the constraint system");
300    let prove_ms = prove_start.elapsed().as_millis();
301    let metrics = window.finish();
302    std::hint::black_box(proof);
303
304    println!(
305        "{}",
306        Report {
307            workload: "fixture",
308            srs_log2: None,
309            curve: Some(&fixture.curve),
310            seed: Some(&fixture.seed),
311            domain_log2: Some(domain_log2),
312            setup_ms,
313            prove_ms,
314            metrics,
315        }
316    );
317}
318
319fn profile_mina_fixture(fixture: &Fixture) {
320    if fixture.curve == Vesta::NAME {
321        profile_fixture_curve::<Vesta, BaseSpongeVesta, ScalarSpongeVesta>(fixture);
322    } else if fixture.curve == Pallas::NAME {
323        profile_fixture_curve::<Pallas, BaseSpongePallas, ScalarSpongePallas>(fixture);
324    } else {
325        panic!("Unsupported curve: {}", fixture.curve);
326    }
327}
328
329#[derive(Parser)]
330#[command(about = "End-to-end prover memory profile; emits one JSON object on stdout")]
331enum Cli {
332    /// Prove the synthetic benchmark circuit (kimchi::bench::BenchmarkCtx)
333    Synthetic {
334        /// log2 of the domain/SRS size
335        #[arg(long, default_value_t = 16, value_parser = clap::value_parser!(u32).range(4..=28))]
336        srs_log2: u32,
337    },
338    /// Prove a serialised mina circuit fixture (kimchi_inputs_CURVE_SEED.ser)
339    Fixture { fixture: Fixture },
340}
341
342fn main() {
343    match Cli::parse() {
344        Cli::Synthetic { srs_log2 } => profile_synthetic(srs_log2),
345        Cli::Fixture { fixture } => profile_mina_fixture(&fixture),
346    }
347}
memory_profile/memory_profile.rs

memory_profile/
memory_profile.rs
