Refactor expression tracegen for recursion vk

Author: hero78119

ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/air.rs

@@ -1,12 +1,12 @@
 use std::borrow::Borrow;
 
-use openvm_circuit_primitives::{utils::assert_array_eq, SubAir};
+use openvm_circuit_primitives::{SubAir, utils::assert_array_eq};
 use openvm_stark_backend::{
-    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
+    BaseAirWithPublicValues, PartitionedBaseAir, interaction::InteractionBuilder,
 };
 use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
 use p3_air::{Air, AirBuilder, BaseAir};
-use p3_field::{extension::BinomiallyExtendable, PrimeCharacteristicRing};
+use p3_field::{PrimeCharacteristicRing, extension::BinomiallyExtendable};
 use p3_matrix::Matrix;
 use stark_recursion_circuit_derive::AlignedBorrow;
 

ceno_recursion_v2/src/batch_constraint/expr_eval/constraints_folding/trace.rs

@@ -1,15 +1,13 @@
 use std::borrow::BorrowMut;
 
-use itertools::Itertools;
-use openvm_stark_backend::keygen::types::MultiStarkVerifyingKey0;
-use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, D_EF, EF, F};
+use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F};
 use p3_field::{BasedVectorSpace, PrimeCharacteristicRing};
 use p3_matrix::dense::RowMajorMatrix;
 use p3_maybe_rayon::prelude::*;
 
 use crate::{
     batch_constraint::expr_eval::constraints_folding::air::ConstraintsFoldingCols,
-    system::Preflight,
+    system::{Preflight, RecursionVk},
     tracegen::RowMajorChip,
     utils::{MultiProofVecVec, MultiVecWithBounds},
 };
@@ -33,15 +31,28 @@ pub(crate) struct ConstraintsFoldingBlob {
 
 impl ConstraintsFoldingBlob {
     pub fn new(
-        vk: &MultiStarkVerifyingKey0<BabyBearPoseidon2Config>,
+        child_vk: &RecursionVk,
         expr_evals: &MultiVecWithBounds<EF, 2>,
         preflights: &[&Preflight],
     ) -> Self {
-        let constraints = vk
-            .per_air
-            .iter()
-            .map(|vk| vk.symbolic_constraints.constraints.constraint_idx.clone())
-            .collect_vec();
+        let mut max_air_idx = 0usize;
+        for key in child_vk.circuit_index_to_name.keys().copied() {
+            max_air_idx = max_air_idx.max(key);
+        }
+        let mut constraints = vec![Vec::<usize>::new(); max_air_idx + 1];
+        for (&air_idx, name) in &child_vk.circuit_index_to_name {
+            let expr_len = child_vk
+                .circuit_vks
+                .get(name)
+                .and_then(|vk| vk.cs.gkr_circuit.as_ref())
+                .and_then(|circuit| circuit.layers.get(0))
+                .map(|layer| layer.exprs.len())
+                .unwrap_or_default();
+            if air_idx >= constraints.len() {
+                constraints.resize(air_idx + 1, vec![]);
+            }
+            constraints[air_idx] = (0..expr_len).collect();
+        }
 
         let mut records = MultiProofVecVec::new();
         let mut folded = MultiProofVecVec::new();
@@ -79,8 +90,9 @@ impl ConstraintsFoldingBlob {
                         value,
                     });
                 }
-                let n_lift = v.log_height.saturating_sub(vk.params.l_skip);
-                let n = v.log_height as isize - vk.params.l_skip as isize;
+                let l_skip = preflight.proof_shape.l_skip;
+                let n_lift = v.log_height.saturating_sub(l_skip);
+                let n = v.log_height as isize - l_skip as isize;
                 folded.push((
                     n,
                     folded_claim * preflight.batch_constraint.eq_ns_frontloaded[n_lift],
@@ -186,15 +198,15 @@ impl RowMajorChip<F> for ConstraintsFoldingTraceGenerator {
 #[cfg(feature = "cuda")]
 pub(in crate::batch_constraint) mod cuda {
     use openvm_circuit_primitives::cuda_abi::UInt2;
-    use openvm_cuda_backend::{base::DeviceMatrix, GpuBackend};
+    use openvm_cuda_backend::{GpuBackend, base::DeviceMatrix};
     use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer};
     use openvm_stark_backend::prover::AirProvingContext;
 
     use super::*;
     use crate::{
         batch_constraint::cuda_abi::{
-            constraints_folding_tracegen, constraints_folding_tracegen_temp_bytes, AffineFpExt,
-            FpExtWithTidx,
+            AffineFpExt, FpExtWithTidx, constraints_folding_tracegen,
+            constraints_folding_tracegen_temp_bytes,
         },
         cuda::{preflight::PreflightGpu, vk::VerifyingKeyGpu},
         tracegen::ModuleChip,

ceno_recursion_v2/src/batch_constraint/expr_eval/symbolic_expression/air.rs

@@ -3,12 +3,12 @@ use std::borrow::Borrow;
 
 use openvm_circuit_primitives::{encoder::Encoder, utils::assert_array_eq};
 use openvm_stark_backend::{
-    air_builders::PartitionedAirBuilder, interaction::InteractionBuilder, BaseAirWithPublicValues,
-    PartitionedBaseAir,
+    BaseAirWithPublicValues, PartitionedBaseAir, air_builders::PartitionedAirBuilder,
+    interaction::InteractionBuilder,
 };
 use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
 use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
-use p3_field::{extension::BinomiallyExtendable, Field, PrimeCharacteristicRing};
+use p3_field::{Field, PrimeCharacteristicRing, extension::BinomiallyExtendable};
 use p3_matrix::Matrix;
 use stark_recursion_circuit_derive::AlignedBorrow;
 use strum::{EnumCount, IntoEnumIterator};
@@ -37,25 +37,26 @@ pub const ENCODER_MAX_DEGREE: u32 = 2;
 
 #[derive(Debug, Clone, Copy, EnumIter, EnumCount)]
 pub enum NodeKind {
-    VarPreprocessed = 0,
-    VarMain = 1,
-    VarPublicValue = 2,
-    SelIsFirst = 3,
-    SelIsLast = 4,
-    SelIsTransition = 5,
-    Constant = 6,
-    Add = 7,
-    Sub = 8,
-    Neg = 9,
-    Mul = 10,
-    InteractionMult = 11,
-    InteractionMsgComp = 12,
-    InteractionBusIndex = 13,
+    WitIn = 0,
+    StructuralWitIn = 1,
+    Fixed = 2,
+    Instance = 3,
+    SelIsFirst = 4,
+    SelIsLast = 5,
+    SelIsTransition = 6,
+    Constant = 7,
+    Add = 8,
+    Sub = 9,
+    Neg = 10,
+    Mul = 11,
+    InteractionMult = 12,
+    InteractionMsgComp = 13,
+    InteractionBusIndex = 14,
 }
 
 impl Default for NodeKind {
     fn default() -> Self {
-        NodeKind::VarPreprocessed
+        NodeKind::WitIn
     }
 }
 
@@ -161,12 +162,22 @@ where
                 NodeKind::Neg,
                 NodeKind::InteractionMult,
                 NodeKind::InteractionMsgComp,
+                NodeKind::WitIn,
+                NodeKind::StructuralWitIn,
+                NodeKind::Fixed,
+                NodeKind::Instance,
             ]
             .map(|x| x as usize),
         );
         let is_arg1_node_idx = enc.contains_flag::<AB>(
             &flags,
-            &[NodeKind::Add, NodeKind::Sub, NodeKind::Mul].map(|x| x as usize),
+            &[
+                NodeKind::Add,
+                NodeKind::Sub,
+                NodeKind::Mul,
+                NodeKind::InteractionMsgComp,
+            ]
+            .map(|x| x as usize),
         );
 
         for (proof_idx, (&cols, &next_cols)) in main_cols.iter().zip(&next_main_cols).enumerate() {
@@ -180,9 +191,8 @@ where
             let next_proof_present = next_slot_state.clone()
                 * (AB::Expr::from_u8(3) - next_slot_state)
                 * AB::F::TWO.inverse();
-            let air_present = slot_state.clone()
-                * (slot_state.clone() - AB::Expr::ONE)
-                * AB::F::TWO.inverse();
+            let air_present =
+                slot_state.clone() * (slot_state.clone() - AB::Expr::ONE) * AB::F::TWO.inverse();
 
             let arg_ef0: [AB::Var; D_EF] = cols.args[..D_EF].try_into().unwrap();
             let arg_ef1: [AB::Var; D_EF] = cols.args[D_EF..2 * D_EF].try_into().unwrap();
@@ -207,15 +217,16 @@ where
                     NodeKind::Neg => scalar_subtract_ext_field::<AB::Expr>(AB::Expr::ZERO, arg_ef0),
                     NodeKind::Mul => ext_field_multiply::<AB::Expr>(arg_ef0, arg_ef1),
                     NodeKind::Constant => base_to_ext(cached_cols.attrs[0]),
-                    NodeKind::VarPublicValue => base_to_ext(cols.args[0]),
+                    NodeKind::Instance => base_to_ext(cols.args[0]),
                     NodeKind::SelIsFirst => ext_field_multiply(arg_ef0, arg_ef1),
                     NodeKind::SelIsLast => ext_field_multiply(arg_ef0, arg_ef1),
                     NodeKind::SelIsTransition => scalar_subtract_ext_field(
                         AB::Expr::ONE,
                         ext_field_multiply(arg_ef0, arg_ef1),
                     ),
-                    NodeKind::VarPreprocessed
-                    | NodeKind::VarMain
+                    NodeKind::WitIn
+                    | NodeKind::StructuralWitIn
+                    | NodeKind::Fixed
                     | NodeKind::InteractionMult
                     | NodeKind::InteractionMsgComp => arg_ef0.map(Into::into),
                     NodeKind::InteractionBusIndex => {
@@ -261,7 +272,7 @@ where
 
             let is_var = enc.contains_flag::<AB>(
                 &flags,
-                &[NodeKind::VarMain, NodeKind::VarPreprocessed].map(|x| x as usize),
+                &[NodeKind::WitIn, NodeKind::StructuralWitIn, NodeKind::Fixed].map(|x| x as usize),
             );
             self.column_claims_bus.receive(
                 builder,
@@ -283,8 +294,7 @@ where
                     pv_idx: cached_cols.attrs[0],
                     value: cols.args[0],
                 },
-                enc.get_flag_expr::<AB>(NodeKind::VarPublicValue as usize, &flags)
-                    * air_present.clone(),
+                enc.get_flag_expr::<AB>(NodeKind::Instance as usize, &flags) * air_present.clone(),
             );
             self.air_shape_bus.lookup_key(
                 builder,