Repository: openvm-org/openvm
Pull request #1553: feat: make provers generic over engine
Author: login: Golovanov399
Created: 2025-04-01T22:15:55Z | Merged: 2025-04-02T23:35:28Z

This adds the ability to use different engines for proving.


Diff:
diff --git a/benchmarks/src/bin/fib_e2e.rs b/benchmarks/src/bin/fib_e2e.rs
index 63c327ddec..d19e0b56e8 100644
--- a/benchmarks/src/bin/fib_e2e.rs
+++ b/benchmarks/src/bin/fib_e2e.rs
@@ -12,7 +12,9 @@ use openvm_rv32im_transpiler::{
 use openvm_sdk::{
     commit::commit_app_exe, prover::ContinuationProver, DefaultStaticVerifierPvHandler, Sdk, StdIn,
 };
-use openvm_stark_sdk::bench::run_with_metric_collection;
+use openvm_stark_sdk::{
+    bench::run_with_metric_collection, config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
+};
 use openvm_transpiler::{transpiler::Transpiler, FromElf};
 
 const NUM_PUBLIC_VALUES: usize = DEFAULT_MAX_NUM_PUBLIC_VALUES;
@@ -56,8 +58,12 @@ async fn main() -> Result<()> {
     let mut stdin = StdIn::default();
     stdin.write(&n);
     run_with_metric_collection("OUTPUT_PATH", || {
-        let mut e2e_prover =
-            ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk);
+        let mut e2e_prover = ContinuationProver::<_, BabyBearPoseidon2Engine>::new(
+            &halo2_params_reader,
+            app_pk,
+            app_committed_exe,
+            full_agg_pk,
+        );
         e2e_prover.set_program_name("fib_e2e");
         let _proof = e2e_prover.generate_proof_for_evm(stdin);
     });
diff --git a/benchmarks/src/bin/kitchen_sink.rs b/benchmarks/src/bin/kitchen_sink.rs
index c5d9ca7ab2..9ec4a26d71 100644
--- a/benchmarks/src/bin/kitchen_sink.rs
+++ b/benchmarks/src/bin/kitchen_sink.rs
@@ -13,7 +13,9 @@ use openvm_sdk::{
     commit::commit_app_exe, config::SdkVmConfig, prover::ContinuationProver,
     DefaultStaticVerifierPvHandler, Sdk, StdIn,
 };
-use openvm_stark_sdk::bench::run_with_metric_collection;
+use openvm_stark_sdk::{
+    bench::run_with_metric_collection, config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
+};
 use openvm_transpiler::FromElf;
 
 fn main() -> Result<()> {
@@ -78,8 +80,12 @@ fn main() -> Result<()> {
     )?;
 
     run_with_metric_collection("OUTPUT_PATH", || -> Result<()> {
-        let mut prover =
-            ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk);
+        let mut prover = ContinuationProver::<_, BabyBearPoseidon2Engine>::new(
+            &halo2_params_reader,
+            app_pk,
+            app_committed_exe,
+            full_agg_pk,
+        );
         prover.set_program_name("kitchen_sink");
         let stdin = StdIn::default();
         let _proof = prover.generate_proof_for_evm(stdin);
diff --git a/benchmarks/src/bin/verify_fibair.rs b/benchmarks/src/bin/verify_fibair.rs
index 95f7c6abbe..aac1f4ea79 100644
--- a/benchmarks/src/bin/verify_fibair.rs
+++ b/benchmarks/src/bin/verify_fibair.rs
@@ -59,8 +59,8 @@ fn main() -> Result<()> {
         let app_pk = sdk.app_keygen(app_config)?;
         let app_vk = app_pk.get_app_vk();
         let committed_exe = sdk.commit_app_exe(app_fri_params, program.into())?;
-        let prover =
-            AppProver::new(app_pk.app_vm_pk, committed_exe).with_program_name("verify_fibair");
+        let prover = AppProver::<_, BabyBearPoseidon2Engine>::new(app_pk.app_vm_pk, committed_exe)
+            .with_program_name("verify_fibair");
         let proof = prover.generate_app_proof_without_continuations(input_stream.into());
         sdk.verify_app_proof_without_continuations(&app_vk, &proof)?;
         Ok(())
diff --git a/benchmarks/src/utils.rs b/benchmarks/src/utils.rs
index 1a04804fd0..198bf8b763 100644
--- a/benchmarks/src/utils.rs
+++ b/benchmarks/src/utils.rs
@@ -24,6 +24,7 @@ use openvm_stark_sdk::{
         baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
         FriParameters,
     },
+    engine::StarkFriEngine,
     openvm_stark_backend::Chip,
     p3_baby_bear::BabyBear,
 };
@@ -157,7 +158,7 @@ impl BenchmarkCli {
         VC::Periphery: Chip<SC>,
     {
         let app_config = self.app_config(vm_config);
-        bench_from_exe(
+        bench_from_exe::<VC, BabyBearPoseidon2Engine>(
             bench_name,
             app_config,
             exe,
@@ -202,7 +203,7 @@ pub fn build_bench(manifest_dir: PathBuf, profile: impl ToString) -> Result<Elf>
 /// 6. Verify STARK proofs.
 ///
 /// Returns the data necessary for proof aggregation.
-pub fn bench_from_exe<VC>(
+pub fn bench_from_exe<VC, E: StarkFriEngine<SC>>(
     bench_name: impl ToString,
     app_config: AppConfig<VC>,
     exe: impl Into<VmExe<F>>,
@@ -231,7 +232,8 @@ where
     // 4. Generate trace
     // 5. Generate STARK proofs for each segment (segmentation is determined by `config`), with timer.
     let app_vk = app_pk.get_app_vk();
-    let prover = AppProver::new(app_pk.app_vm_pk, committed_exe).with_program_name(bench_name);
+    let prover =
+        AppProver::<VC, E>::new(app_pk.app_vm_pk, committed_exe).with_program_name(bench_name);
     let app_proof = prover.generate_app_proof(input_stream);
     // 6. Verify STARK proofs, including boundary conditions.
     let sdk = Sdk::new();
@@ -239,10 +241,8 @@ where
         .expect("Verification failed");
     if let Some(leaf_vm_config) = leaf_vm_config {
         let leaf_vm_pk = leaf_keygen(app_config.leaf_fri_params.fri_params, leaf_vm_config);
-        let leaf_prover = VmLocalProver::<SC, NativeConfig, BabyBearPoseidon2Engine>::new(
-            leaf_vm_pk,
-            app_pk.leaf_committed_exe,
-        );
+        let leaf_prover =
+            VmLocalProver::<SC, NativeConfig, E>::new(leaf_vm_pk, app_pk.leaf_committed_exe);
         let leaf_controller = LeafProvingController {
             num_children: DEFAULT_NUM_CHILDREN_LEAF,
         };
diff --git a/crates/sdk/examples/sdk_app.rs b/crates/sdk/examples/sdk_app.rs
index 9514fb6485..fb78a01a3a 100644
--- a/crates/sdk/examples/sdk_app.rs
+++ b/crates/sdk/examples/sdk_app.rs
@@ -9,7 +9,7 @@ use openvm_sdk::{
     prover::AppProver,
     Sdk, StdIn,
 };
-use openvm_stark_sdk::config::FriParameters;
+use openvm_stark_sdk::config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, FriParameters};
 use openvm_transpiler::elf::Elf;
 use serde::{Deserialize, Serialize};
 
@@ -90,8 +90,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 9a. Generate a proof
     let proof = sdk.generate_app_proof(app_pk.clone(), app_committed_exe.clone(), stdin.clone())?;
     // 9b. Generate a proof with an AppProver with custom fields
-    let app_prover = AppProver::new(app_pk.app_vm_pk.clone(), app_committed_exe.clone())
-        .with_program_name("test_program");
+    let app_prover = AppProver::<_, BabyBearPoseidon2Engine>::new(
+        app_pk.app_vm_pk.clone(),
+        app_committed_exe.clone(),
+    )
+    .with_program_name("test_program");
     let proof = app_prover.generate_app_proof(stdin.clone());
     // ANCHOR_END: proof_generation
 
diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs
index 7f0640bd2b..9d3f62a32e 100644
--- a/crates/sdk/src/lib.rs
+++ b/crates/sdk/src/lib.rs
@@ -176,7 +176,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let app_prover = AppProver::new(app_pk.app_vm_pk.clone(), app_committed_exe);
+        let app_prover = AppProver::<VC, E>::new(app_pk.app_vm_pk.clone(), app_committed_exe);
         let proof = app_prover.generate_app_proof(inputs);
         Ok(proof)
     }
@@ -241,7 +241,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let stark_prover = StarkProver::new(app_pk, app_exe, agg_stark_pk);
+        let stark_prover = StarkProver::<VC, E>::new(app_pk, app_exe, agg_stark_pk);
         let proof = stark_prover.generate_root_verifier_input(inputs);
         Ok(proof)
     }
@@ -258,7 +258,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let e2e_prover = ContinuationProver::new(reader, app_pk, app_exe, agg_pk);
+        let e2e_prover = ContinuationProver::<VC, E>::new(reader, app_pk, app_exe, agg_pk);
         let proof = e2e_prover.generate_proof_for_evm(inputs);
         Ok(proof)
     }
diff --git a/crates/sdk/src/prover/agg.rs b/crates/sdk/src/prover/agg.rs
index bcf202f60d..3055516424 100644
--- a/crates/sdk/src/prover/agg.rs
+++ b/crates/sdk/src/prover/agg.rs
@@ -7,9 +7,7 @@ use openvm_continuations::verifier::{
 };
 use openvm_native_circuit::NativeConfig;
 use openvm_native_recursion::hints::Hintable;
-use openvm_stark_sdk::{
-    config::baby_bear_poseidon2::BabyBearPoseidon2Engine, openvm_stark_backend::proof::Proof,
-};
+use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::proof::Proof};
 use tracing::info_span;
 
 use crate::{
@@ -25,11 +23,11 @@ pub const DEFAULT_NUM_CHILDREN_LEAF: usize = 1;
 const DEFAULT_NUM_CHILDREN_INTERNAL: usize = 2;
 const DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS: usize = 4;
 
-pub struct AggStarkProver {
-    leaf_prover: VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
+pub struct AggStarkProver<E: StarkFriEngine<SC>> {
+    leaf_prover: VmLocalProver<SC, NativeConfig, E>,
     leaf_controller: LeafProvingController,
 
-    internal_prover: VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
+    internal_prover: VmLocalProver<SC, NativeConfig, E>,
     root_prover: RootVerifierLocalProver,
 
     pub num_children_internal: usize,
@@ -41,19 +39,17 @@ pub struct LeafProvingController {
     pub num_children: usize,
 }
 
-impl AggStarkProver {
+impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
     pub fn new(
         agg_stark_pk: AggStarkProvingKey,
         leaf_committed_exe: Arc<NonRootCommittedExe>,
     ) -> Self {
-        let leaf_prover = VmLocalProver::<SC, NativeConfig, BabyBearPoseidon2Engine>::new(
-            agg_stark_pk.leaf_vm_pk,
-            leaf_committed_exe,
-        );
+        let leaf_prover =
+            VmLocalProver::<SC, NativeConfig, E>::new(agg_stark_pk.leaf_vm_pk, leaf_committed_exe);
         let leaf_controller = LeafProvingController {
             num_children: DEFAULT_NUM_CHILDREN_LEAF,
         };
-        let internal_prover = VmLocalProver::<SC, NativeConfig, BabyBearPoseidon2Engine>::new(
+        let internal_prover = VmLocalProver::<SC, NativeConfig, E>::new(
             agg_stark_pk.internal_vm_pk,
             agg_stark_pk.internal_committed_exe,
         );
@@ -184,9 +180,9 @@ impl LeafProvingController {
         self
     }
 
-    pub fn generate_proof(
+    pub fn generate_proof<E: StarkFriEngine<SC>>(
         &self,
-        prover: &VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
+        prover: &VmLocalProver<SC, NativeConfig, E>,
         app_proofs: &ContinuationVmProof<SC>,
     ) -> Vec<Proof<SC>> {
         info_span!("agg_layer", group = "leaf").in_scope(|| {
diff --git a/crates/sdk/src/prover/app.rs b/crates/sdk/src/prover/app.rs
index 6ed7d528b7..095351677e 100644
--- a/crates/sdk/src/prover/app.rs
+++ b/crates/sdk/src/prover/app.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use getset::Getters;
 use openvm_circuit::arch::{ContinuationVmProof, VmConfig};
 use openvm_stark_backend::{proof::Proof, Chip};
-use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine;
+use openvm_stark_sdk::engine::StarkFriEngine;
 use tracing::info_span;
 
 use super::vm::SingleSegmentVmProver;
@@ -13,13 +13,13 @@ use crate::{
 };
 
 #[derive(Getters)]
-pub struct AppProver<VC> {
+pub struct AppProver<VC, E: StarkFriEngine<SC>> {
     pub program_name: Option<String>,
     #[getset(get = "pub")]
-    app_prover: VmLocalProver<SC, VC, BabyBearPoseidon2Engine>,
+    app_prover: VmLocalProver<SC, VC, E>,
 }
 
-impl<VC> AppProver<VC> {
+impl<VC, E: StarkFriEngine<SC>> AppProver<VC, E> {
     pub fn new(
         app_vm_pk: Arc<VmProvingKey<SC, VC>>,
         app_committed_exe: Arc<NonRootCommittedExe>,
@@ -29,10 +29,7 @@ impl<VC> AppProver<VC> {
     {
         Self {
             program_name: None,
-            app_prover: VmLocalProver::<SC, VC, BabyBearPoseidon2Engine>::new(
-                app_vm_pk,
-                app_committed_exe,
-            ),
+            app_prover: VmLocalProver::<SC, VC, E>::new(app_vm_pk, app_committed_exe),
         }
     }
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
diff --git a/crates/sdk/src/prover/mod.rs b/crates/sdk/src/prover/mod.rs
index 5e4fee43ad..e0e4fedd04 100644
--- a/crates/sdk/src/prover/mod.rs
+++ b/crates/sdk/src/prover/mod.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use openvm_circuit::arch::VmConfig;
-use openvm_stark_sdk::openvm_stark_backend::Chip;
+use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::Chip};
 
 use crate::{keygen::AppProvingKey, stdin::StdIn, NonRootCommittedExe, F, SC};
 
@@ -24,12 +24,12 @@ pub use stark::*;
 
 use crate::{keygen::AggProvingKey, prover::halo2::Halo2Prover, types::EvmProof};
 
-pub struct ContinuationProver<VC> {
-    stark_prover: StarkProver<VC>,
+pub struct ContinuationProver<VC, E: StarkFriEngine<SC>> {
+    stark_prover: StarkProver<VC, E>,
     halo2_prover: Halo2Prover,
 }
 
-impl<VC> ContinuationProver<VC> {
+impl<VC, E: StarkFriEngine<SC>> ContinuationProver<VC, E> {
     pub fn new(
         reader: &impl Halo2ParamsReader,
         app_pk: Arc<AppProvingKey<VC>>,
diff --git a/crates/sdk/src/prover/stark.rs b/crates/sdk/src/prover/stark.rs
index 99d053b1f4..d9d8d17d14 100644
--- a/crates/sdk/src/prover/stark.rs
+++ b/crates/sdk/src/prover/stark.rs
@@ -3,6 +3,7 @@ use std::sync::Arc;
 use openvm_circuit::arch::VmConfig;
 use openvm_continuations::verifier::root::types::RootVmVerifierInput;
 use openvm_stark_backend::{proof::Proof, Chip};
+use openvm_stark_sdk::engine::StarkFriEngine;
 
 use crate::{
     keygen::{AggStarkProvingKey, AppProvingKey},
@@ -10,11 +11,11 @@ use crate::{
     NonRootCommittedExe, RootSC, StdIn, F, SC,
 };
 
-pub struct StarkProver<VC> {
-    app_prover: AppProver<VC>,
-    agg_prover: AggStarkProver,
+pub struct StarkProver<VC, E: StarkFriEngine<SC>> {
+    app_prover: AppProver<VC, E>,
+    agg_prover: AggStarkProver<E>,
 }
-impl<VC> StarkProver<VC> {
+impl<VC, E: StarkFriEngine<SC>> StarkProver<VC, E> {
     pub fn new(
         app_pk: Arc<AppProvingKey<VC>>,
         app_committed_exe: Arc<NonRootCommittedExe>,


Reviews:
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T06:48:17Z
user: login: jonathanpwang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T06:55:34Z
user: login: jonathanpwang
body: `AggStarkProver` needs similar `<E>`
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T06:55:55Z
user: login: jonathanpwang
body: LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T23:35:23Z

Review comments:
diff_hunk: @@ -144,7 +142,7 @@ impl BenchmarkCli {
         build_bench(manifest_dir, profile)
     }
 
-    pub fn bench_from_exe<VC>(
+    pub fn bench_from_exe<VC, E: StarkFriEngine<SC>>(
path: benchmarks/src/utils.rs
user: login: jonathanpwang
body: should/can we keep this function for fixed engine just for convenience? so we don't need to change all the benchmarks?
created_at: 2025-04-02T06:48:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -3,18 +3,19 @@ use std::sync::Arc;
 use openvm_circuit::arch::VmConfig;
 use openvm_continuations::verifier::root::types::RootVmVerifierInput;
 use openvm_stark_backend::{proof::Proof, Chip};
+use openvm_stark_sdk::engine::StarkFriEngine;
 
 use crate::{
     keygen::{AggStarkProvingKey, AppProvingKey},
     prover::{agg::AggStarkProver, app::AppProver},
     NonRootCommittedExe, RootSC, StdIn, F, SC,
 };
 
-pub struct StarkProver<VC> {
-    app_prover: AppProver<VC>,
+pub struct StarkProver<VC, E: StarkFriEngine<SC>> {
+    app_prover: AppProver<VC, E>,
     agg_prover: AggStarkProver,
path: crates/sdk/src/prover/stark.rs
user: login: jonathanpwang
body: It looks like `AggStarkProver` also needs generic `E`
created_at: 2025-04-02T06:55:34Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-01T22:22:16Z
body: <!--Benchmark Results-->
| group | app.proof_time_ms | app.cycles | app.cells_used | leaf.proof_time_ms | leaf.cycles | leaf.cells_used |
| -- | -- | -- | -- | -- | -- | -- |
| [verify_fibair](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/verify_fibair-9bce5df86a9a85c1cbe89449a35a43381334f34e.md) | 1,281 |  334,142 |  17,677,298 |- | - | - |
| [fibonacci_program](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/fibonacci-9bce5df86a9a85c1cbe89449a35a43381334f34e.md) |<span style='color: green'>(-57 [-2.1%])</span> 2,702 |  1,500,277 |  50,589,503 |- | - | - |
| [regex_program](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/regex-9bce5df86a9a85c1cbe89449a35a43381334f34e.md) |<span style='color: green'>(-188 [-2.3%])</span> 8,136 |  4,165,226 |  166,511,152 |- | - | - |
| [ecrecover_program](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/ecrecover-9bce5df86a9a85c1cbe89449a35a43381334f34e.md) |<span style='color: green'>(-52 [-3.5%])</span> 1,426 |  289,447 |  14,470,186 |- | - | - |
| [pairing](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/pairing-9bce5df86a9a85c1cbe89449a35a43381334f34e.md) |<span style='color: green'>(-59 [-1.2%])</span> 4,882 |  1,820,436 |  95,832,407 |- | - | - |


Commit: https://github.com/openvm-org/openvm/commit/9bce5df86a9a85c1cbe89449a35a43381334f34e

[Benchmark Workflow](https://github.com/openvm-org/openvm/actions/runs/14207118634)

author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-02T13:48:48Z
body: <!--Benchmark Results-->
| group | app.proof_time_ms | app.cycles | app.cells_used | leaf.proof_time_ms | leaf.cycles | leaf.cells_used |
| -- | -- | -- | -- | -- | -- | -- |
| [verify_fibair](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/verify_fibair-69be2521ceceeb8cfe0f03f9c3522b18a50dabd2.md) |<span style='color: green'>(-3 [-0.2%])</span> 1,279 |  334,086 |  17,676,626 |- | - | - |
| [fibonacci_program](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/fibonacci-69be2521ceceeb8cfe0f03f9c3522b18a50dabd2.md) |<span style='color: green'>(-73 [-2.6%])</span> 2,686 |  1,500,277 |  50,589,503 |- | - | - |
| [regex_program](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/regex-69be2521ceceeb8cfe0f03f9c3522b18a50dabd2.md) |<span style='color: green'>(-221 [-2.7%])</span> 8,103 |  4,165,226 |  166,511,152 |- | - | - |
| [ecrecover_program](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/ecrecover-69be2521ceceeb8cfe0f03f9c3522b18a50dabd2.md) |<span style='color: green'>(-68 [-4.6%])</span> 1,410 |  289,447 |  14,470,186 |- | - | - |
| [pairing](https://github.com/openvm-org/openvm/blob/benchmark-results/benchmarks-pr/1553/pairing-69be2521ceceeb8cfe0f03f9c3522b18a50dabd2.md) |<span style='color: red'>(+40 [+0.8%])</span> 4,981 |  1,820,436 |  95,832,407 |- | - | - |


Commit: https://github.com/openvm-org/openvm/commit/69be2521ceceeb8cfe0f03f9c3522b18a50dabd2

[Benchmark Workflow](https://github.com/openvm-org/openvm/actions/runs/14221161418)

author_association: NONE