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_native_recursion::fri::MAX_TWO_ADICITY;
10use openvm_sdk::{
11    config::SdkVmConfig,
12    prover::vm::{new_local_prover, types::VmProvingKey},
13    Sdk, StdIn, F, SC,
14};
15use openvm_stark_backend::{config::StarkGenericConfig, p3_field::BasedVectorSpace};
16use openvm_stark_sdk::bench::run_with_metric_collection;
17
18fn verify_native_max_trace_heights(
19    sdk: &Sdk,
20    app_exe: Arc<VmExe<F>>,
21    leaf_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
22    num_children_leaf: usize,
23) -> Result<()> {
24    let app_proof = sdk.app_prover(app_exe)?.prove(StdIn::default())?;
25    let leaf_inputs =
26        LeafVmVerifierInput::chunk_continuation_vm_proof(&app_proof, num_children_leaf);
27    let mut leaf_prover = new_local_prover::<Poseidon2Engine, _>(
28        NativeBuilder::default(),
29        &leaf_vm_pk,
30        sdk.app_pk().leaf_committed_exe.exe.clone(),
31    )?;
32
33    for leaf_input in leaf_inputs {
34        let exe = leaf_prover.exe().clone();
35        let vm = &mut leaf_prover.vm;
36        let metered_ctx = vm
37            .build_metered_ctx(&exe)
38            .with_max_cells(usize::MAX) // no segmentation
39            .with_max_trace_height(1 << MAX_TWO_ADICITY);
40        let (segments, _) = vm
41            .metered_interpreter(&exe)?
42            .execute_metered(leaf_input.write_to_stream(), metered_ctx)?;
43        assert_eq!(segments.len(), 1);
44        let estimated_trace_heights = &segments[0].trace_heights;
45        println!("estimated_trace_heights: {estimated_trace_heights:?}");
46
47        // Tracegen without proving since leaf proofs take a while
48        let state = vm.create_initial_state(&exe, leaf_input.write_to_stream());
49        vm.transport_init_memory_to_device(&state.memory);
50        let mut interpreter = vm.preflight_interpreter(&exe)?;
51        let out = vm.execute_preflight(&mut interpreter, state, None, estimated_trace_heights)?;
52        let actual_trace_heights = vm
53            .generate_proving_ctx(out.system_records, out.record_arenas)?
54            .per_air
55            .into_iter()
56            .map(|(_, air_ctx)| air_ctx.main_trace_height())
57            .collect::<Vec<usize>>();
58        println!("actual_trace_heights: {actual_trace_heights:?}");
59
60        actual_trace_heights
61            .iter()
62            .zip(NATIVE_MAX_TRACE_HEIGHTS)
63            .for_each(|(&actual, &expected)| {
64                assert!(
65                    actual <= (expected as usize),
66                    "Actual trace height {actual} exceeds expected height {expected}"
67                );
68            });
69    }
70    Ok(())
71}
72
73fn main() -> Result<()> {
74    let args = BenchmarkCli::parse();
75
76    let vm_config =
77        SdkVmConfig::from_toml(include_str!("../../../guest/kitchen-sink/openvm.toml"))?
78            .app_vm_config;
79    let app_config = args.app_config(vm_config.clone());
80    let elf = args.build_bench_program("kitchen-sink", &vm_config, None)?;
81    let sdk = Sdk::new(app_config)?;
82    let exe = sdk.convert_to_exe(elf)?;
83
84    let (_, app_vk) = sdk.app_keygen();
85    let ext_degree = <<SC as StarkGenericConfig>::Challenge as BasedVectorSpace<F>>::DIMENSION;
86    println!(
87        "app_total_width = {}",
88        app_vk
89            .vk
90            .inner
91            .per_air
92            .iter()
93            .map(|vk| vk.params.width.total_width(ext_degree))
94            .sum::<usize>()
95    );
96    let agg_pk = sdk.agg_pk();
97    println!(
98        "leaf_total_width = {}",
99        agg_pk
100            .leaf_vm_pk
101            .vm_pk
102            .get_vk()
103            .inner
104            .per_air
105            .iter()
106            .map(|vk| vk.params.width.total_width(ext_degree))
107            .sum::<usize>()
108    );
109    // Verify that NATIVE_MAX_TRACE_HEIGHTS remains valid
110    verify_native_max_trace_heights(
111        &sdk,
112        exe.clone(),
113        agg_pk.leaf_vm_pk.clone(),
114        args.agg_tree_config.num_children_leaf,
115    )?;
116
117    run_with_metric_collection("OUTPUT_PATH", || -> eyre::Result<()> {
118        let stdin = StdIn::default();
119        #[cfg(not(feature = "evm"))]
120        {
121            let mut prover = sdk.prover(exe)?.with_program_name("kitchen_sink");
122            let app_commit = prover.app_commit();
123            let proof = prover.prove(stdin)?;
124            Sdk::verify_proof(&agg_pk.get_agg_vk(), app_commit, &proof)?;
125        }
126        #[cfg(feature = "evm")]
127        let _proof = sdk
128            .evm_prover(exe)?
129            .with_program_name("kitchen_sink")
130            .prove_evm(stdin)?;
131        Ok(())
132    })
133}