optimize iop_verifier_state_verify

Author: hero78119

Cargo.lock

@@ -52,11 +52,12 @@ name = "e2e_aggregate"
 path = "src/bin/e2e_aggregate.rs"
 
 [features]
-bench-metrics = ["openvm-circuit/metrics"]
+bench-metrics = ["openvm-sdk/metrics"]
 default = ["parallel", "nightly-features"]
 gpu = ["ceno_zkvm/gpu", "openvm-circuit/cuda", "openvm-native-circuit/cuda", "dep:openvm-cuda-backend"]
 nightly-features = [
   "ceno_zkvm/nightly-features",
   "openvm-stark-sdk/nightly-features",
 ]
 parallel = ["openvm-stark-backend/parallel"]
+perf-metrics = ["openvm-sdk/perf-metrics"]

ceno_recursion/Cargo.toml

@@ -927,6 +927,52 @@ impl<C: Config> UniPolyExtrapolator<C> {
         res
     }
 
+    pub fn extrapolate_uni_poly_deg1(
+        &self,
+        builder: &mut Builder<C>,
+        p_i_0: Ext<C::F, C::EF>,
+        p_i_1: Ext<C::F, C::EF>,
+        eval_at: Ext<C::F, C::EF>,
+    ) -> Ext<C::F, C::EF> {
+        self.extrapolate_uni_poly_deg_1(builder, p_i_0, p_i_1, eval_at)
+    }
+
+    pub fn extrapolate_uni_poly_deg2(
+        &self,
+        builder: &mut Builder<C>,
+        p_i_0: Ext<C::F, C::EF>,
+        p_i_1: Ext<C::F, C::EF>,
+        p_i_2: Ext<C::F, C::EF>,
+        eval_at: Ext<C::F, C::EF>,
+    ) -> Ext<C::F, C::EF> {
+        self.extrapolate_uni_poly_deg_2(builder, p_i_0, p_i_1, p_i_2, eval_at)
+    }
+
+    pub fn extrapolate_uni_poly_deg3(
+        &self,
+        builder: &mut Builder<C>,
+        p_i_0: Ext<C::F, C::EF>,
+        p_i_1: Ext<C::F, C::EF>,
+        p_i_2: Ext<C::F, C::EF>,
+        p_i_3: Ext<C::F, C::EF>,
+        eval_at: Ext<C::F, C::EF>,
+    ) -> Ext<C::F, C::EF> {
+        self.extrapolate_uni_poly_deg_3(builder, p_i_0, p_i_1, p_i_2, p_i_3, eval_at)
+    }
+
+    pub fn extrapolate_uni_poly_deg4(
+        &self,
+        builder: &mut Builder<C>,
+        p_i_0: Ext<C::F, C::EF>,
+        p_i_1: Ext<C::F, C::EF>,
+        p_i_2: Ext<C::F, C::EF>,
+        p_i_3: Ext<C::F, C::EF>,
+        p_i_4: Ext<C::F, C::EF>,
+        eval_at: Ext<C::F, C::EF>,
+    ) -> Ext<C::F, C::EF> {
+        self.extrapolate_uni_poly_deg_4(builder, p_i_0, p_i_1, p_i_2, p_i_3, p_i_4, eval_at)
+    }
+
     fn extrapolate_uni_poly_deg_1(
         &self,
         builder: &mut Builder<C>,

ceno_recursion/src/arithmetics/mod.rs

@@ -5,7 +5,9 @@ use crate::{
         exts_to_felts,
     },
     tower_verifier::binding::IOPProverMessageVecVariable,
-    transcript::transcript_observe_label,
+    transcript::{
+        transcript_label_as_array, transcript_observe_label, transcript_observe_label_felts,
+    },
     zkvm_verifier::binding::TowerProofInputVariable,
 };
 use openvm_native_compiler::prelude::*;
@@ -22,51 +24,78 @@ pub fn iop_verifier_state_verify<C: Config>(
     out_claim: &Ext<C::F, C::EF>,
     prover_messages: &IOPProverMessageVecVariable<C>,
     max_num_variables: Felt<C::F>,
-    max_degree: Felt<C::F>,
+    max_degree: usize,
     unipoly_extrapolator: &UniPolyExtrapolator<C>,
 ) -> (
     Array<C, Ext<<C as Config>::F, <C as Config>::EF>>,
     Ext<<C as Config>::F, <C as Config>::EF>,
 ) {
-    // TODO: either store it in a global cache or pass them as parameters
-    let zero: Ext<C::F, C::EF> = builder.constant(C::EF::ZERO);
     let zero_f: Felt<C::F> = builder.constant(C::F::ZERO);
 
+    let max_degree_felt: Felt<C::F> = builder.constant(C::F::from_canonical_usize(max_degree));
     let max_num_variables_usize: Usize<C::N> =
         Usize::from(builder.cast_felt_to_var(max_num_variables));
     challenger.observe(builder, max_num_variables);
     challenger.observe(builder, zero_f);
-    challenger.observe(builder, max_degree);
+    challenger.observe(builder, max_degree_felt);
     challenger.observe(builder, zero_f);
 
     builder.assert_var_eq(max_num_variables_usize.get_var(), prover_messages.len());
 
     let challenges: Array<C, Ext<C::F, C::EF>> = builder.dyn_array(max_num_variables_usize.clone());
-    let expected: Ext<C::F, C::EF> = builder.eval(*out_claim + zero);
+    let expected: Ext<C::F, C::EF> = builder.eval(*out_claim);
+    let internal_round_label = transcript_label_as_array(builder, b"Internal round");
 
     builder.cycle_tracker_start("IOPVerifierState::verify_round_and_update_state");
+    let curr_offset: Var<C::N> = builder.eval(C::N::ZERO);
     builder
         .range(0, max_num_variables_usize.clone())
         .for_each(|i_vec, builder| {
             let i = i_vec[0];
 
-            // TODO: this takes 7 cycles, can we optimize it?
-            let prover_msg = prover_messages.get(builder, i.variable());
+            let next_offset: Var<C::N> =
+                builder.eval(curr_offset + prover_messages.prover_message_size);
+            let prover_msg = prover_messages
+                .evaluations
+                .slice(builder, curr_offset, next_offset);
+            builder.assign(&curr_offset, next_offset);
 
             unsafe {
                 let prover_msg_felts = exts_to_felts(builder, &prover_msg);
                 challenger_multi_observe(builder, challenger, &prover_msg_felts);
             }
 
-            transcript_observe_label(builder, challenger, b"Internal round");
+            transcript_observe_label_felts(builder, challenger, &internal_round_label);
             let challenge = challenger.sample_ext(builder);
 
             let e1 = builder.get(&prover_msg, 0);
             let e0 = builder.eval(expected - e1);
-            let p_r =
-                unipoly_extrapolator.extrapolate_uni_poly(builder, e0, &prover_msg, challenge);
-
-            builder.assign(&expected, p_r + zero);
+            let p_r = match max_degree {
+                1 => unipoly_extrapolator.extrapolate_uni_poly_deg1(builder, e0, e1, challenge),
+                2 => {
+                    let p1 = e1;
+                    let p2 = builder.get(&prover_msg, 1);
+                    unipoly_extrapolator.extrapolate_uni_poly_deg2(builder, e0, p1, p2, challenge)
+                }
+                3 => {
+                    let p1 = e1;
+                    let p2 = builder.get(&prover_msg, 1);
+                    let p3 = builder.get(&prover_msg, 2);
+                    unipoly_extrapolator
+                        .extrapolate_uni_poly_deg3(builder, e0, p1, p2, p3, challenge)
+                }
+                4 => {
+                    let p1 = e1;
+                    let p2 = builder.get(&prover_msg, 1);
+                    let p3 = builder.get(&prover_msg, 2);
+                    let p4 = builder.get(&prover_msg, 3);
+                    unipoly_extrapolator
+                        .extrapolate_uni_poly_deg4(builder, e0, p1, p2, p3, p4, challenge)
+                }
+                _ => panic!("unsupported max_degree {max_degree}"),
+            };
+
+            builder.assign(&expected, p_r);
             builder.set_value(&challenges, i, challenge);
         });
     builder.cycle_tracker_end("IOPVerifierState::verify_round_and_update_state");
@@ -258,16 +287,14 @@ pub fn verify_tower_proof<C: Config>(
         let max_num_variables: Felt<C::F> = builder.constant(C::F::ONE);
         builder.assign(&max_num_variables, max_num_variables + round);
 
-        let max_degree = builder.constant(C::F::from_canonical_usize(3));
-
         builder.cycle_tracker_start("sumcheck verify");
         let (sub_rt, sub_e) = iop_verifier_state_verify(
             builder,
             challenger,
             out_claim,
             &prover_messages,
             max_num_variables,
-            max_degree,
+            3,
             unipoly_extrapolator,
         );
         builder.cycle_tracker_end("sumcheck verify");

ceno_recursion/src/tower_verifier/program.rs

@@ -5,6 +5,8 @@ use openvm_native_recursion::challenger::{
 };
 use openvm_stark_backend::p3_field::FieldAlgebra;
 
+use crate::arithmetics::challenger_multi_observe;
+
 pub fn transcript_observe_label<C: Config>(
     builder: &mut Builder<C>,
     challenger: &mut DuplexChallengerVariable<C>,
@@ -17,6 +19,27 @@ pub fn transcript_observe_label<C: Config>(
     }
 }
 
+pub fn transcript_label_as_array<C: Config>(
+    builder: &mut Builder<C>,
+    label: &[u8],
+) -> Array<C, Felt<C::F>> {
+    let label_f = <BabyBearExt4 as CenoExtensionField>::BaseField::bytes_to_field_elements(label);
+    let arr: Array<C, Felt<C::F>> = builder.dyn_array(label_f.len());
+    for (idx, n) in label_f.into_iter().enumerate() {
+        let f: Felt<C::F> = builder.constant(C::F::from_canonical_u64(n.to_canonical_u64()));
+        builder.set_value(&arr, idx, f);
+    }
+    arr
+}
+
+pub fn transcript_observe_label_felts<C: Config>(
+    builder: &mut Builder<C>,
+    challenger: &mut DuplexChallengerVariable<C>,
+    label_felts: &Array<C, Felt<C::F>>,
+) {
+    challenger_multi_observe(builder, challenger, label_felts);
+}
+
 pub fn transcript_check_pow_witness<C: Config>(
     builder: &mut Builder<C>,
     challenger: &mut DuplexChallengerVariable<C>,

ceno_recursion/src/transcript/mod.rs

@@ -953,8 +953,6 @@ pub fn verify_gkr_circuit<C: Config>(
             });
 
         // sigma = \sum_b sel(b) * zero_expr(b)
-        let max_degree = builder.constant(C::F::from_canonical_usize(layer.max_expr_degree + 1));
-
         let max_num_variables_f = builder.unsafe_cast_var_to_felt(max_num_variables.get_var());
 
         let (in_point, expected_evaluation) = iop_verifier_state_verify(
@@ -963,7 +961,7 @@ pub fn verify_gkr_circuit<C: Config>(
             &sigma,
             &proof,
             max_num_variables_f,
-            max_degree,
+            layer.max_expr_degree + 1,
             unipoly_extrapolator,
         );
 
@@ -1204,15 +1202,14 @@ pub fn verify_rotation<C: Config>(
     let sigma: Ext<C::F, C::EF> = builder.constant(C::EF::ZERO);
 
     let max_num_variables = builder.unsafe_cast_var_to_felt(max_num_variables.get_var());
-    let max_degree: Felt<C::F> = builder.constant(C::F::TWO);
 
     let (origin_point, expected_evaluation) = iop_verifier_state_verify(
         builder,
         challenger,
         &sigma,
         proof,
         max_num_variables,
-        max_degree,
+        2,
         unipoly_extrapolator,
     );
 
@@ -1756,15 +1753,14 @@ pub fn verify_ecc_proof<C: Config>(
 
     let one_ext: Ext<C::F, C::EF> = builder.constant(C::EF::ONE);
     let zero_ext: Ext<C::F, C::EF> = builder.constant(C::EF::ZERO);
-    let three_f: Felt<C::F> = builder.constant(C::F::from_canonical_u32(3));
     let num_vars_f = builder.unsafe_cast_var_to_felt(num_vars.get_var());
     let (rt, sumcheck_expected_evaluation) = iop_verifier_state_verify(
         builder,
         challenger,
         &zero_ext,
         &proof.zerocheck_proof,
         num_vars_f,
-        three_f,
+        3,
         unipoly_extrapolator,
     );