kitchen_sink/
kitchen_sink.rs

1use std::sync::Arc;
2
3use clap::Parser;
4use eyre::Result;
5use openvm_benchmarks_prove::util::BenchmarkCli;
6use openvm_circuit::{arch::instructions::exe::VmExe, utils::TestStarkEngine as Poseidon2Engine};
7use openvm_continuations::verifier::leaf::types::LeafVmVerifierInput;
8use openvm_native_circuit::{NativeBuilder, NativeConfig, NATIVE_MAX_TRACE_HEIGHTS};
9use openvm_sdk::{
10    config::SdkVmConfig,
11    prover::vm::{new_local_prover, types::VmProvingKey},
12    Sdk, StdIn, F, SC,
13};
14use openvm_stark_sdk::bench::run_with_metric_collection;
15
16fn verify_native_max_trace_heights(
17    sdk: &Sdk,
18    app_exe: Arc<VmExe<F>>,
19    leaf_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
20    num_children_leaf: usize,
21) -> Result<()> {
22    let app_proof = sdk.app_prover(app_exe)?.prove(StdIn::default())?;
23    let leaf_inputs =
24        LeafVmVerifierInput::chunk_continuation_vm_proof(&app_proof, num_children_leaf);
25    let mut leaf_prover = new_local_prover::<Poseidon2Engine, _>(
26        NativeBuilder::default(),
27        &leaf_vm_pk,
28        sdk.app_pk().leaf_committed_exe.exe.clone(),
29    )?;
30
31    for leaf_input in leaf_inputs {
32        let exe = leaf_prover.exe().clone();
33        let vm = &mut leaf_prover.vm;
34        let metered_ctx = vm.build_metered_ctx();
35        let (segments, _) = vm
36            .metered_interpreter(&exe)?
37            .execute_metered(leaf_input.write_to_stream(), metered_ctx)?;
38        assert_eq!(segments.len(), 1);
39        let estimated_trace_heights = &segments[0].trace_heights;
40        println!("estimated_trace_heights: {:?}", estimated_trace_heights);
41
42        // Tracegen without proving since leaf proofs take a while
43        let state = vm.create_initial_state(&exe, leaf_input.write_to_stream());
44        vm.transport_init_memory_to_device(&state.memory);
45        let mut interpreter = vm.preflight_interpreter(&exe)?;
46        let out = vm.execute_preflight(&mut interpreter, state, None, estimated_trace_heights)?;
47        let actual_trace_heights = vm
48            .generate_proving_ctx(out.system_records, out.record_arenas)?
49            .per_air
50            .into_iter()
51            .map(|(_, air_ctx)| air_ctx.main_trace_height())
52            .collect::<Vec<usize>>();
53        println!("actual_trace_heights: {:?}", actual_trace_heights);
54
55        actual_trace_heights
56            .iter()
57            .zip(NATIVE_MAX_TRACE_HEIGHTS)
58            .for_each(|(&actual, &expected)| {
59                assert!(
60                    actual <= (expected as usize),
61                    "Actual trace height {} exceeds expected height {}",
62                    actual,
63                    expected
64                );
65            });
66    }
67    Ok(())
68}
69
70fn main() -> Result<()> {
71    let args = BenchmarkCli::parse();
72
73    let vm_config =
74        SdkVmConfig::from_toml(include_str!("../../../guest/kitchen-sink/openvm.toml"))?
75            .app_vm_config;
76    let app_config = args.app_config(vm_config.clone());
77    let elf = args.build_bench_program("kitchen-sink", &vm_config, None)?;
78    let sdk = Sdk::new(app_config)?;
79    let exe = sdk.convert_to_exe(elf)?;
80
81    let agg_pk = sdk.agg_pk();
82    // Verify that NATIVE_MAX_TRACE_HEIGHTS remains valid
83    verify_native_max_trace_heights(
84        &sdk,
85        exe.clone(),
86        agg_pk.leaf_vm_pk.clone(),
87        args.agg_tree_config.num_children_leaf,
88    )?;
89
90    run_with_metric_collection("OUTPUT_PATH", || -> eyre::Result<()> {
91        let stdin = StdIn::default();
92        #[cfg(not(feature = "evm"))]
93        {
94            let mut prover = sdk.prover(exe)?.with_program_name("kitchen_sink");
95            let app_commit = prover.app_commit();
96            let proof = prover.prove(stdin)?;
97            Sdk::verify_proof(&agg_pk.get_agg_vk(), app_commit, &proof)?;
98        }
99        #[cfg(feature = "evm")]
100        let _proof = sdk
101            .evm_prover(exe)?
102            .with_program_name("kitchen_sink")
103            .prove_evm(stdin)?;
104        Ok(())
105    })
106}