Skip to main content

airbender_host/prover/
cpu_prover.rs

1//! CPU-based prover for Airbender programs.
2//!
3//! # Warning: High Memory Requirement
4//!
5//! CPU proving typically requires **96 GB or more of RAM** to complete successfully.
6//! Running this prover on machines with insufficient memory will cause the process to crash.
7//!
8//! # When to use the CPU prover
9//!
10//! Using CPU proving is **very rarely a good idea**. It is primarily a reference
11//! implementation and is most useful for debugging circuit constraints. In almost all
12//! cases you want either the **dev prover** (for rapid iteration without real proofs)
13//! or the **GPU prover** (for production-grade performance). There should be a very
14//! specific reason to use the CPU prover before choosing it over the alternatives.
15
16use super::{
17    receipt_from_real_proof, resolve_app_bin_path, resolve_text_path, resolve_worker_threads,
18    ProveResult, Prover, DEFAULT_CPU_CYCLE_BOUND, DEFAULT_RAM_BOUND_BYTES,
19};
20use crate::error::{HostError, Result};
21use crate::proof::{Proof, RealProof};
22use crate::runner::{Runner, TranspilerRunnerBuilder};
23use crate::security::SecurityLevel;
24use execution_utils::setups;
25use execution_utils::unrolled;
26use riscv_transpiler::abstractions::non_determinism::QuasiUARTSource;
27use riscv_transpiler::common_constants::rom::ROM_BYTE_SIZE;
28use riscv_transpiler::cycle::IMStandardIsaConfigWithUnsignedMulDiv;
29use std::path::{Path, PathBuf};
30
31/// Minimum system RAM (in GB) required to run CPU proving without crashing.
32const MIN_RAM_GB: u64 = 96;
33
34/// Minimum system RAM in bytes derived from [`MIN_RAM_GB`].
35const MIN_RAM_BYTES: u64 = MIN_RAM_GB * 1024 * 1024 * 1024;
36
37/// Environment variable that, when set to `true`, skips the system RAM check.
38const MEM_OVERRIDE_ENV: &str = "AIRBENDER_PLATFORM_CPU_PROVER_MEM_OVERRIDE";
39
40fn check_system_ram() -> Result<()> {
41    if std::env::var(MEM_OVERRIDE_ENV).as_deref() == Ok("true") {
42        return Ok(());
43    }
44
45    let mut sys = sysinfo::System::new();
46    sys.refresh_memory();
47    let total_ram = sys.total_memory();
48
49    if total_ram < MIN_RAM_BYTES {
50        let detected_gb = total_ram / (1024 * 1024 * 1024);
51        return Err(HostError::Prover(format!(
52            "System is expected to have at least {MIN_RAM_GB} GB of ram, but this system only has \
53             {detected_gb} GB. On machines with not enough RAM, process might crash. If you want \
54             to run it anyway, set `{MEM_OVERRIDE_ENV}=true` and run again"
55        )));
56    }
57
58    Ok(())
59}
60
61/// Builder for creating a configured cached CPU prover.
62pub struct CpuProverBuilder {
63    app_bin_path: PathBuf,
64    worker_threads: Option<usize>,
65    cycles: Option<usize>,
66    ram_bound: Option<usize>,
67    security: SecurityLevel,
68}
69
70impl CpuProverBuilder {
71    pub fn new(app_bin_path: impl AsRef<Path>) -> Self {
72        Self {
73            app_bin_path: app_bin_path.as_ref().to_path_buf(),
74            worker_threads: None,
75            cycles: None,
76            ram_bound: None,
77            security: SecurityLevel::default(),
78        }
79    }
80
81    pub fn with_worker_threads(mut self, worker_threads: usize) -> Self {
82        self.worker_threads = Some(worker_threads);
83        self
84    }
85
86    pub fn maybe_worker_threads(self, worker_threads: Option<usize>) -> Self {
87        match worker_threads {
88            Some(v) => self.with_worker_threads(v),
89            None => self,
90        }
91    }
92
93    pub fn with_cycles(mut self, cycles: usize) -> Self {
94        self.cycles = Some(cycles);
95        self
96    }
97
98    pub fn maybe_cycles(self, cycles: Option<usize>) -> Self {
99        match cycles {
100            Some(v) => self.with_cycles(v),
101            None => self,
102        }
103    }
104
105    pub fn with_ram_bound(mut self, ram_bound: usize) -> Self {
106        self.ram_bound = Some(ram_bound);
107        self
108    }
109
110    pub fn maybe_ram_bound(self, ram_bound: Option<usize>) -> Self {
111        match ram_bound {
112            Some(v) => self.with_ram_bound(v),
113            None => self,
114        }
115    }
116
117    pub fn with_security(mut self, security: SecurityLevel) -> Self {
118        self.security = security;
119        self
120    }
121
122    pub fn build(self) -> Result<CpuProver> {
123        CpuProver::new(
124            &self.app_bin_path,
125            self.worker_threads,
126            self.cycles,
127            self.ram_bound,
128            self.security,
129        )
130    }
131}
132
133/// CPU prover wrapper that caches padded artifacts and worker threads.
134pub struct CpuProver {
135    app_bin_path: PathBuf,
136    app_text_path: PathBuf,
137    binary_u32: Vec<u32>,
138    text_u32: Vec<u32>,
139    cycles: Option<usize>,
140    ram_bound: usize,
141    security: SecurityLevel,
142    worker: execution_utils::prover_examples::prover::worker::Worker,
143}
144
145impl CpuProver {
146    fn new(
147        app_bin_path: &Path,
148        worker_threads: Option<usize>,
149        cycles: Option<usize>,
150        ram_bound: Option<usize>,
151        security: SecurityLevel,
152    ) -> Result<Self> {
153        check_system_ram()?;
154
155        if matches!(worker_threads, Some(0)) {
156            return Err(HostError::Prover(
157                "worker thread count must be greater than zero".to_string(),
158            ));
159        }
160
161        let app_bin_path = resolve_app_bin_path(app_bin_path)?;
162        let app_text_path = resolve_text_path(&app_bin_path)?;
163        let (_, binary_u32) = setups::read_and_pad_binary(&app_bin_path);
164        let (_, text_u32) = setups::read_and_pad_binary(&app_text_path);
165
166        let ram_bound = ram_bound.unwrap_or(DEFAULT_RAM_BOUND_BYTES);
167        if ram_bound < ROM_BYTE_SIZE {
168            return Err(HostError::Prover(format!(
169                "ram bound must be at least {} bytes",
170                ROM_BYTE_SIZE
171            )));
172        }
173
174        let threads = resolve_worker_threads(worker_threads);
175        let worker =
176            execution_utils::prover_examples::prover::worker::Worker::new_with_num_threads(threads);
177
178        Ok(Self {
179            app_bin_path,
180            app_text_path,
181            binary_u32,
182            text_u32,
183            cycles,
184            ram_bound,
185            security,
186            worker,
187        })
188    }
189}
190
191impl Prover for CpuProver {
192    fn prove(&self, input_words: &[u32]) -> Result<ProveResult> {
193        let cycles_bound = match self.cycles {
194            Some(value) => value,
195            None => {
196                let cycle_estimator = TranspilerRunnerBuilder::new(&self.app_bin_path)
197                    .with_cycles(DEFAULT_CPU_CYCLE_BOUND)
198                    .with_text_path(&self.app_text_path)
199                    .build()?;
200                let outcome = cycle_estimator.run(input_words)?;
201                if !outcome.reached_end {
202                    return Err(HostError::Prover(format!(
203                        "automatic cycle estimation did not reach program end after {} cycles; provide explicit cycles to prove a bounded run",
204                        outcome.cycles_executed
205                    )));
206                }
207                outcome.cycles_executed
208            }
209        };
210        if cycles_bound == 0 {
211            return Err(HostError::Prover(
212                "cycles bound must be greater than zero".to_string(),
213            ));
214        }
215
216        let oracle = QuasiUARTSource::new_with_reads(input_words.to_vec());
217        let inner_proof = unrolled::prove_unrolled_for_machine_configuration_into_program_proof::<
218            IMStandardIsaConfigWithUnsignedMulDiv,
219        >(
220            &self.binary_u32,
221            &self.text_u32,
222            cycles_bound,
223            oracle,
224            self.ram_bound,
225            &self.worker,
226            self.security.into(),
227        );
228        let receipt = receipt_from_real_proof(&inner_proof);
229        let proof = Proof::Real(RealProof::new(
230            self.security,
231            super::ProverLevel::Base,
232            inner_proof,
233        ));
234
235        Ok(ProveResult {
236            proof,
237            cycles: cycles_bound as u64,
238            receipt,
239        })
240    }
241}