sp1_cuda/
lib.rs

1/// The shared API between the client and server.
2pub mod api;
3
4/// The client that interacts with the CUDA server.
5pub mod client;
6
7/// The proving key type, which is a "remote" reference to a key held by the CUDA server.
8pub mod pk;
9
10/// The server startup logic.
11mod server;
12
13mod error;
14pub use error::CudaClientError;
15
16pub use pk::CudaProvingKey;
17use sp1_core_executor::SP1Context;
18use sp1_core_machine::io::SP1Stdin;
19use sp1_core_machine::riscv::RiscvAir;
20use sp1_hypercube::Machine;
21use sp1_primitives::{Elf, SP1Field};
22use sp1_prover::worker::ProofFromNetwork;
23use sp1_prover_types::network_base_types::ProofMode;
24
25use crate::client::CudaClient;
26
27#[derive(Clone)]
28pub struct CudaProver {
29    client: CudaClient,
30}
31
32impl CudaProver {
33    /// Create a new prover, using the 0th CUDA device.
34    pub async fn new() -> Result<Self, CudaClientError> {
35        Ok(Self { client: CudaClient::connect(0).await? })
36    }
37
38    /// Create a new prover, using the given CUDA device.
39    pub async fn new_with_id(cuda_id: u32) -> Result<Self, CudaClientError> {
40        Ok(Self { client: CudaClient::connect(cuda_id).await? })
41    }
42
43    /// Setup a new proving key.
44    pub async fn setup(&self, elf: Elf) -> Result<CudaProvingKey, CudaClientError> {
45        self.setup_with_machine(elf, RiscvAir::machine()).await
46    }
47
48    /// Same as [`Self::setup`] but with a custom machine.
49    pub async fn setup_with_machine(
50        &self,
51        elf: Elf,
52        machine: Machine<SP1Field, RiscvAir<SP1Field>>,
53    ) -> Result<CudaProvingKey, CudaClientError> {
54        self.client.setup(elf, machine).await
55    }
56
57    pub async fn prove_with_mode(
58        &self,
59        pk: &CudaProvingKey,
60        stdin: SP1Stdin,
61        context: SP1Context<'static>,
62        mode: ProofMode,
63    ) -> Result<ProofFromNetwork, CudaClientError> {
64        self.client.prove_with_mode(pk, stdin, context, mode).await
65    }
66}
sp1_cuda/lib.rs

sp1_cuda/
lib.rs
