feat: give access to cycle count when using CudaProver

Author: leruaa

crates/sdk/src/cuda/mod.rs

@@ -77,32 +77,26 @@ impl CudaProver {
     pub fn prove<'a>(&'a self, pk: &'a SP1ProvingKey, stdin: &'a SP1Stdin) -> CudaProveBuilder<'a> {
         CudaProveBuilder { prover: self, mode: SP1ProofMode::Core, pk, stdin: stdin.clone() }
     }
-}
 
-impl Prover<CpuProverComponents> for CudaProver {
-    fn setup(&self, elf: &[u8]) -> (SP1ProvingKey, SP1VerifyingKey) {
-        let (pk, vk) = self.cuda_prover.setup(elf).unwrap();
-        (pk, vk)
-    }
-
-    fn inner(&self) -> &SP1Prover<CpuProverComponents> {
-        &self.cpu_prover
-    }
-
-    fn prove(
+    /// Proves the given program on the given input in the given proof mode.
+    ///
+    /// Returns the cycle count in addition to the proof.
+    pub fn prove_with_report(
         &self,
         pk: &SP1ProvingKey,
         stdin: &SP1Stdin,
         kind: SP1ProofMode,
-    ) -> Result<SP1ProofWithPublicValues> {
+    ) -> Result<(SP1ProofWithPublicValues, u64)> {
         // Generate the core proof.
         let proof = self.cuda_prover.prove_core(stdin)?;
+        let cycles = proof.cycles;
         if kind == SP1ProofMode::Core {
-            return Ok(SP1ProofWithPublicValues::new(
+            let proof_with_pv = SP1ProofWithPublicValues::new(
                 SP1Proof::Core(proof.proof.0),
                 proof.public_values,
                 self.version().to_string(),
-            ));
+            );
+            return Ok((proof_with_pv, cycles));
         }
 
         // Generate the compressed proof.
@@ -111,11 +105,12 @@ impl Prover<CpuProverComponents> for CudaProver {
         let public_values = proof.public_values.clone();
         let reduce_proof = self.cuda_prover.compress(&pk.vk, proof, deferred_proofs)?;
         if kind == SP1ProofMode::Compressed {
-            return Ok(SP1ProofWithPublicValues::new(
+            let proof_with_pv = SP1ProofWithPublicValues::new(
                 SP1Proof::Compressed(Box::new(reduce_proof)),
                 public_values,
                 self.version().to_string(),
-            ));
+            );
+            return Ok((proof_with_pv, cycles));
         }
 
         // Generate the shrink proof.
@@ -134,11 +129,12 @@ impl Prover<CpuProverComponents> for CudaProver {
                 try_install_circuit_artifacts("plonk")
             };
             let proof = self.cpu_prover.wrap_plonk_bn254(outer_proof, &plonk_bn254_artifacts);
-            return Ok(SP1ProofWithPublicValues::new(
+            let proof_with_pv = SP1ProofWithPublicValues::new(
                 SP1Proof::Plonk(proof),
                 public_values,
                 self.version().to_string(),
-            ));
+            );
+            return Ok((proof_with_pv, cycles));
         } else if kind == SP1ProofMode::Groth16 {
             let groth16_bn254_artifacts = if sp1_prover::build::sp1_dev_mode() {
                 sp1_prover::build::try_build_groth16_bn254_artifacts_dev(
@@ -150,17 +146,38 @@ impl Prover<CpuProverComponents> for CudaProver {
             };
 
             let proof = self.cpu_prover.wrap_groth16_bn254(outer_proof, &groth16_bn254_artifacts);
-            return Ok(SP1ProofWithPublicValues::new(
+            let proof_with_pv = SP1ProofWithPublicValues::new(
                 SP1Proof::Groth16(proof),
                 public_values,
                 self.version().to_string(),
-            ));
+            );
+            return Ok((proof_with_pv, cycles));
         }
 
         unreachable!()
     }
 }
 
+impl Prover<CpuProverComponents> for CudaProver {
+    fn setup(&self, elf: &[u8]) -> (SP1ProvingKey, SP1VerifyingKey) {
+        let (pk, vk) = self.cuda_prover.setup(elf).unwrap();
+        (pk, vk)
+    }
+
+    fn inner(&self) -> &SP1Prover<CpuProverComponents> {
+        &self.cpu_prover
+    }
+
+    fn prove(
+        &self,
+        pk: &SP1ProvingKey,
+        stdin: &SP1Stdin,
+        kind: SP1ProofMode,
+    ) -> Result<SP1ProofWithPublicValues> {
+        self.prove_with_report(pk, stdin, kind).map(|(p, _)| p)
+    }
+}
+
 impl Default for CudaProver {
     fn default() -> Self {
         Self::new(SP1Prover::new(), None)