airbender_host/prover/
gpu_prover.rs

1use super::{
2    base_path, receipt_from_real_proof, resolve_app_bin_path, ProveResult, Prover, ProverLevel,
3};
4use crate::error::{HostError, Result};
5use crate::proof::{Proof, RealProof};
6use crate::security::SecurityLevel;
7use execution_utils::unrolled_gpu::UnrolledProver;
8use gpu_prover::execution::prover::ExecutionProverConfiguration;
9use riscv_transpiler::abstractions::non_determinism::QuasiUARTSource;
10use std::any::Any;
11use std::path::{Path, PathBuf};
12use std::sync::atomic::{AtomicBool, Ordering};
13use std::sync::{mpsc, Mutex};
14use std::thread::JoinHandle;
15
16/// Low-level configuration for the GPU prover.
17///
18/// The default configuration is meant to be optimal in the vast majority of
19/// cases, so overriding it is only recommended if you either have specific
20/// constraints or you need to fine-tune the configuration for a specific and
21/// unusual workload. Each field left unset keeps the prover's own default.
22#[derive(Clone, Copy, Debug, Default)]
23pub struct GpuProverConfig {
24    worker_threads: Option<usize>,
25    max_device_memory_bytes: Option<usize>,
26    host_allocators_per_job: Option<usize>,
27    host_allocators_per_device: Option<usize>,
28}
29
30impl GpuProverConfig {
31    /// Number of worker threads for the prover's thread pool and replay workers.
32    pub fn with_worker_threads(mut self, worker_threads: usize) -> Self {
33        self.worker_threads = Some(worker_threads);
34        self
35    }
36
37    /// Like [`Self::with_worker_threads`] but a no-op when `None`.
38    pub fn maybe_worker_threads(mut self, worker_threads: Option<usize>) -> Self {
39        if let Some(worker_threads) = worker_threads {
40            self.worker_threads = Some(worker_threads);
41        }
42        self
43    }
44
45    /// Caps the GPU device allocator at `bytes`. By default the allocator grabs
46    /// all free VRAM, which leaves no room for a co-resident prover in the same
47    /// process (e.g. the SNARK wrapper). Capping the FRI prover leaves headroom
48    /// for it. A cap at or above the free amount is a no-op.
49    pub fn with_max_device_memory_bytes(mut self, bytes: usize) -> Self {
50        self.max_device_memory_bytes = Some(bytes);
51        self
52    }
53
54    /// Pinned host transfer buffers pre-allocated per concurrent job (64 MiB
55    /// each). Lowering this below the prover default reclaims committed RAM at
56    /// the cost of less host<->device pipelining.
57    pub fn with_host_allocators_per_job(mut self, count: usize) -> Self {
58        self.host_allocators_per_job = Some(count);
59        self
60    }
61
62    /// Pinned host transfer buffers pre-allocated per GPU device (64 MiB each).
63    /// See [`Self::with_host_allocators_per_job`].
64    pub fn with_host_allocators_per_device(mut self, count: usize) -> Self {
65        self.host_allocators_per_device = Some(count);
66        self
67    }
68}
69
70/// Builder for creating a configured cached GPU prover.
71pub struct GpuProverBuilder {
72    app_bin_path: PathBuf,
73    security: SecurityLevel,
74    level: ProverLevel,
75    config: GpuProverConfig,
76}
77
78impl GpuProverBuilder {
79    pub fn new(app_bin_path: impl AsRef<Path>) -> Self {
80        Self {
81            app_bin_path: app_bin_path.as_ref().to_path_buf(),
82            security: SecurityLevel::default(),
83            level: ProverLevel::RecursionUnified,
84            config: GpuProverConfig::default(),
85        }
86    }
87
88    pub fn with_level(mut self, level: ProverLevel) -> Self {
89        self.level = level;
90        self
91    }
92
93    pub fn with_security(mut self, security: SecurityLevel) -> Self {
94        self.security = security;
95        self
96    }
97
98    /// Applies low-level configuration for the GPU prover. See
99    /// [`GpuProverConfig`]; the default is optimal for the vast majority of
100    /// workloads, so this is only needed for specific constraints or unusual
101    /// workloads.
102    pub fn with_config(mut self, config: GpuProverConfig) -> Self {
103        self.config = config;
104        self
105    }
106
107    pub fn build(self) -> Result<GpuProver> {
108        GpuProver::new(&self.app_bin_path, self.security, self.level, self.config)
109    }
110}
111
112/// GPU prover wrapper that owns and reuses a single `UnrolledProver` instance.
113///
114/// ## Poisoning
115///
116/// Actual proving happens on a separate thread, and in case the program cannot be
117/// proven, the prover can panic. Prover panics are not unwind safe, so the thread
118/// and the prover will be disposed of, making this prover object poisoned, e.g. not
119/// usable for future proving attempts. Once poisoned, the prover will return an error
120/// on all operations.
121///
122/// After poisioning, you can instantiate a new prover if required.
123pub struct GpuProver {
124    command_tx: mpsc::Sender<WorkerCommand>,
125    worker_handle: Mutex<Option<JoinHandle<()>>>,
126    poisoned: AtomicBool,
127}
128
129enum WorkerCommand {
130    Prove {
131        input_words: Vec<u32>,
132        response_tx: mpsc::Sender<Result<ProveResult>>,
133    },
134    Shutdown,
135}
136
137impl GpuProver {
138    fn new(
139        app_bin_path: &Path,
140        security: SecurityLevel,
141        level: ProverLevel,
142        config: GpuProverConfig,
143    ) -> Result<Self> {
144        if matches!(config.worker_threads, Some(0)) {
145            return Err(HostError::Prover(
146                "worker thread count must be greater than zero".to_string(),
147            ));
148        }
149
150        let app_bin_path = resolve_app_bin_path(app_bin_path)?;
151        let (command_tx, worker_handle) = spawn_worker(app_bin_path, security, level, config)?;
152
153        Ok(Self {
154            command_tx,
155            worker_handle: Mutex::new(Some(worker_handle)),
156            poisoned: AtomicBool::new(false),
157        })
158    }
159
160    pub fn is_poisoned(&self) -> bool {
161        self.poisoned.load(Ordering::SeqCst)
162    }
163
164    fn poisoned_error() -> HostError {
165        HostError::Prover("GPU prover is poisoned due to a previous proving panic".to_string())
166    }
167
168    fn handle_worker_failure(&self, operation: &str) -> HostError {
169        if self.poisoned.swap(true, Ordering::SeqCst) {
170            return Self::poisoned_error();
171        }
172
173        match self.take_worker_panic_message() {
174            Some(message) => HostError::Prover(format!(
175                "GPU prover panicked while {operation}; prover is now poisoned: {message}"
176            )),
177            None => HostError::Prover(format!(
178                "GPU prover worker failed while {operation}; prover is now poisoned"
179            )),
180        }
181    }
182
183    fn take_worker_panic_message(&self) -> Option<String> {
184        let mut handle_slot = match self.worker_handle.lock() {
185            Ok(slot) => slot,
186            Err(poisoned) => poisoned.into_inner(),
187        };
188        let handle = handle_slot.take()?;
189
190        match handle.join() {
191            Ok(()) => None,
192            Err(payload) => Some(panic_payload_to_string(payload)),
193        }
194    }
195}
196
197impl Prover for GpuProver {
198    fn prove(&self, input_words: &[u32]) -> Result<ProveResult> {
199        if self.is_poisoned() {
200            return Err(Self::poisoned_error());
201        }
202
203        let (response_tx, response_rx) = mpsc::channel();
204        self.command_tx
205            .send(WorkerCommand::Prove {
206                input_words: input_words.to_vec(),
207                response_tx,
208            })
209            .map_err(|_| self.handle_worker_failure("submitting a prove request"))?;
210
211        response_rx
212            .recv()
213            .map_err(|_| self.handle_worker_failure("receiving a prove response"))?
214    }
215}
216
217impl Drop for GpuProver {
218    fn drop(&mut self) {
219        let _ = self.command_tx.send(WorkerCommand::Shutdown);
220
221        let handle_slot = match self.worker_handle.get_mut() {
222            Ok(slot) => slot,
223            Err(poisoned) => poisoned.into_inner(),
224        };
225
226        if let Some(handle) = handle_slot.take() {
227            let _ = handle.join();
228        }
229    }
230}
231
232fn spawn_worker(
233    app_bin_path: PathBuf,
234    security: SecurityLevel,
235    level: ProverLevel,
236    config: GpuProverConfig,
237) -> Result<(mpsc::Sender<WorkerCommand>, JoinHandle<()>)> {
238    let (command_tx, command_rx) = mpsc::channel();
239    let (init_tx, init_rx) = mpsc::channel();
240
241    let worker_handle = std::thread::Builder::new()
242        .name("airbender-gpu-prover".to_string())
243        .spawn(move || gpu_worker_loop(command_rx, init_tx, app_bin_path, security, level, config))
244        .map_err(|err| {
245            HostError::Prover(format!("failed to spawn GPU prover worker thread: {err}"))
246        })?;
247
248    match init_rx.recv() {
249        Ok(Ok(())) => Ok((command_tx, worker_handle)),
250        Ok(Err(err)) => {
251            let _ = worker_handle.join();
252            Err(err)
253        }
254        Err(_) => {
255            let reason = match worker_handle.join() {
256                Ok(()) => "GPU prover worker exited during initialization".to_string(),
257                Err(payload) => format!(
258                    "GPU prover worker panicked during initialization: {}",
259                    panic_payload_to_string(payload)
260                ),
261            };
262            Err(HostError::Prover(reason))
263        }
264    }
265}
266
267fn gpu_worker_loop(
268    command_rx: mpsc::Receiver<WorkerCommand>,
269    init_tx: mpsc::Sender<Result<()>>,
270    app_bin_path: PathBuf,
271    security: SecurityLevel,
272    level: ProverLevel,
273    config: GpuProverConfig,
274) {
275    // Keep all prover state inside this dedicated thread so a panic does not unwind
276    // through host-call boundaries or require `AssertUnwindSafe`.
277    let prover =
278        match create_unrolled_prover(&app_bin_path, security, level.as_unrolled_level(), config) {
279            Ok(prover) => prover,
280            Err(err) => {
281                let _ = init_tx.send(Err(err));
282                return;
283            }
284        };
285
286    if init_tx.send(Ok(())).is_err() {
287        return;
288    }
289
290    let mut next_batch_id_base: u64 = 0;
291
292    while let Ok(command) = command_rx.recv() {
293        match command {
294            WorkerCommand::Prove {
295                input_words,
296                response_tx,
297            } => {
298                let oracle = QuasiUARTSource::new_with_reads(input_words);
299                let batch_id_base = next_batch_id_base;
300                next_batch_id_base += 1;
301                let (inner_proof, cycles) = prover.prove(batch_id_base, oracle);
302                let receipt = receipt_from_real_proof(&inner_proof);
303                let proof = Proof::Real(RealProof::new(security, level, inner_proof));
304                let result = Ok(ProveResult {
305                    proof,
306                    cycles,
307                    receipt,
308                });
309                let _ = response_tx.send(result);
310            }
311            WorkerCommand::Shutdown => break,
312        }
313    }
314}
315
316fn panic_payload_to_string(payload: Box<dyn Any + Send + 'static>) -> String {
317    if let Some(message) = payload.downcast_ref::<String>() {
318        return message.clone();
319    }
320    if let Some(message) = payload.downcast_ref::<&str>() {
321        return (*message).to_string();
322    }
323
324    "unknown panic payload".to_string()
325}
326
327fn create_unrolled_prover(
328    app_bin_path: &Path,
329    security: SecurityLevel,
330    level: execution_utils::unrolled_gpu::UnrolledProverLevel,
331    config: GpuProverConfig,
332) -> Result<UnrolledProver> {
333    let base_path = base_path(app_bin_path)?;
334    // Map the host-level GpuProverConfig onto the prover's own configuration,
335    // leaving any field the caller did not set at the prover default.
336    let mut configuration = ExecutionProverConfiguration::default();
337    if let Some(threads) = config.worker_threads {
338        configuration.max_thread_pool_threads = Some(threads);
339        configuration.replay_worker_threads_count = threads;
340    }
341    if let Some(count) = config.host_allocators_per_job {
342        configuration.host_allocators_per_job_count = count;
343    }
344    if let Some(count) = config.host_allocators_per_device {
345        configuration.host_allocators_per_device_count = count;
346    }
347    if let Some(bytes) = config.max_device_memory_bytes {
348        // The device allocator works in fixed-size blocks; translate the byte cap
349        // into a block count. A cap below one block is rejected rather than
350        // silently rounded to zero (which would mean "use all free memory").
351        let block_log = configuration.prover_context_config.allocator_block_log_size;
352        let blocks = bytes >> block_log;
353        if blocks == 0 {
354            return Err(HostError::Prover(format!(
355                "max device memory cap of {bytes} bytes is smaller than one allocator block ({} bytes)",
356                1usize << block_log,
357            )));
358        }
359        configuration
360            .prover_context_config
361            .max_device_allocation_blocks_count = Some(blocks);
362    }
363    Ok(UnrolledProver::new(
364        security.into(),
365        &base_path,
366        configuration,
367        level,
368    ))
369}
airbender_host/prover/gpu_prover.rs

airbender_host/prover/
gpu_prover.rs
