fix: add back num_air_id_lookups

Author: stephenh-axiom-xyz

crates/continuations-v2/src/tests/mod.rs

@@ -248,7 +248,8 @@ fn test_internal_recursive_vk_stabilization(def_hook_commit_set: bool) -> Result
     let (_, app_vk) = engine.keygen(&config.create_airs()?.into_airs().collect_vec());
     let def_hook_commit = def_hook_commit_set.then_some([F::ZERO; DIGEST_SIZE]);
 
-    let leaf_prover = InnerProver::<DEFAULT_MAX_NUM_PROOFS>::new::<Engine>(
+    const MAX_LEAF_NUM_PROOFS: usize = 3;
+    let leaf_prover = InnerProver::<MAX_LEAF_NUM_PROOFS>::new::<Engine>(
         Arc::new(app_vk),
         leaf_system_params(),
         false,

crates/recursion/cuda/include/types.h

@@ -20,6 +20,7 @@ typedef struct {
     size_t cached_idx;
     size_t starting_cidx;
     size_t total_interactions;
+    size_t num_air_id_lookups;
 } TraceMetadata;
 
 typedef struct {

crates/recursion/cuda/src/proof_shape/air.cu

@@ -61,6 +61,8 @@ template <typename T, size_t MAX_CACHED> struct ProofShapeCols {
     T n_max;
     T is_n_max_greater;
 
+    T num_air_id_lookups;
+
     // T idx_flags[IDX_FLAGS];
     T cached_commits[MAX_CACHED][DIGEST_SIZE];
 };
@@ -103,6 +105,12 @@ __device__ __forceinline__ void fill_present_row(
     size_t lifted_height = max(height, (size_t)(1 << l_skip));
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, is_present, Fp::one());
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, height, height);
+    COL_WRITE_VALUE(
+        row,
+        typename Cols<MAX_CACHED>::template Type,
+        num_air_id_lookups,
+        trace_data.num_air_id_lookups
+    );
 
     Decomp lifted_height_decomp, num_interactions_decomp, total_interactions_decomp;
     decompose(lifted_height_decomp, lifted_height);
@@ -191,6 +199,7 @@ __device__ __forceinline__ void fill_non_present_row(
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, starting_cidx, final_cidx);
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, is_present, Fp::zero());
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, height, Fp::zero());
+    COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, num_air_id_lookups, Fp::zero());
     row.fill_zero(
         COL_INDEX(typename Cols<MAX_CACHED>::template Type, lifted_height_limbs), NUM_LIMBS
     );
@@ -232,6 +241,7 @@ __device__ __forceinline__ void fill_summary_row(
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, is_last, Fp::one());
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, sorted_idx, Fp::zero());
     COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, is_present, Fp::zero());
+    COL_WRITE_VALUE(row, typename Cols<MAX_CACHED>::template Type, num_air_id_lookups, Fp::zero());
     row.fill_zero(cached_commits_idx, MAX_CACHED * DIGEST_SIZE);
 
     Decomp interaction_decomp, max_interaction_decomp;

crates/recursion/src/cuda/preflight.rs

@@ -10,8 +10,7 @@ use crate::{
     cuda::{
         to_device_or_nullptr,
         types::{TraceHeight, TraceMetadata},
-    },
-    system::Preflight,
+    }, proof_shape::proof_shape::compute_air_shape_lookup_counts, system::Preflight
 };
 
 /*
@@ -106,6 +105,8 @@ impl PreflightGpu {
         let mut total_interactions = 0;
         let l_skip = vk.inner.params.l_skip;
 
+        let bc_air_shape_lookups = compute_air_shape_lookup_counts(vk);
+
         let (sorted_trace_heights, sorted_trace_metadata): (Vec<_>, Vec<_>) = preflight
             .proof_shape
             .sorted_trace_vdata
@@ -119,6 +120,7 @@ impl PreflightGpu {
                     cached_idx: sorted_cached_commits.len(),
                     starting_cidx: cidx,
                     total_interactions,
+                    num_air_id_lookups: bc_air_shape_lookups[*air_idx],
                 };
                 cidx += vdata.cached_commitments.len()
                     + vk.inner.per_air[*air_idx].preprocessed_data.is_some() as usize;

crates/recursion/src/cuda/types.rs

@@ -13,6 +13,7 @@ pub struct TraceMetadata {
     pub cached_idx: usize,
     pub starting_cidx: usize,
     pub total_interactions: usize,
+    pub num_air_id_lookups: usize,
 }
 
 #[repr(C)]

crates/recursion/src/proof_shape/mod.rs

@@ -46,7 +46,6 @@ pub struct AirMetadata {
     need_rot: bool,
     num_public_values: usize,
     num_interactions: usize,
-    num_dag_nodes: usize,
     main_width: usize,
     cached_widths: Vec<usize>,
     preprocessed_width: Option<usize>,
@@ -115,7 +114,6 @@ impl ProofShapeModule {
                 need_rot: avk.params.need_rot,
                 num_public_values: avk.params.num_public_values,
                 num_interactions: avk.num_interactions,
-                num_dag_nodes: avk.num_dag_nodes,
                 main_width: avk.params.width.common_main,
                 cached_widths: avk.params.width.cached_mains.clone(),
                 preprocessed_width: avk.params.width.preprocessed,

crates/recursion/src/proof_shape/proof_shape/air.rs

@@ -90,6 +90,8 @@ pub struct ProofShapeCols<F, const NUM_LIMBS: usize> {
     /// Computed as max(0, n0, n1, ...) where ni = log_height_i - l_skip for each present trace.
     pub n_max: F,
     pub is_n_max_greater: F,
+
+    pub num_air_id_lookups: F,
 }
 
 // Variable-length columns are stored at the end
@@ -261,7 +263,6 @@ where
         let mut main_common_width = AB::Expr::ZERO;
         let mut preprocessed_stacked_width = AB::Expr::ZERO;
         let mut cached_widths = vec![AB::Expr::ZERO; self.max_cached];
-        let mut num_dag_nodes = AB::Expr::ZERO;
 
         // Select values for CommitmentsBus
         let mut preprocessed_commit = [AB::Expr::ZERO; DIGEST_SIZE];
@@ -279,7 +280,6 @@ where
             air_idx += is_current_air.clone() * AB::F::from_usize(i);
             need_rot += is_current_air.clone() * AB::F::from_bool(air_data.need_rot);
             main_common_width += is_current_air.clone() * AB::F::from_usize(air_data.main_width);
-            num_dag_nodes += is_current_air.clone() * AB::F::from_usize(air_data.num_dag_nodes);
 
             if air_data.num_public_values != 0 {
                 has_pvs += is_current_air.clone();
@@ -498,7 +498,7 @@ where
                 property_idx: AirShapeProperty::AirId.to_field(),
                 value: air_idx.clone(),
             },
-            local.is_present * (num_dag_nodes.clone() + AB::Expr::TWO),
+            local.is_present * (local.num_air_id_lookups + AB::Expr::TWO),
         );
 
         self.air_shape_bus.add_key_with_lookups(
@@ -555,7 +555,7 @@ where
                 n_abs: n_abs.clone(),
                 n_sign_bit: local.n_sign_bit.into(),
             },
-            local.is_present * (num_dag_nodes + AB::F::ONE),
+            local.is_present * (local.num_air_id_lookups + AB::F::ONE),
         );
 
         ///////////////////////////////////////////////////////////////////////////////////////////

crates/recursion/src/proof_shape/proof_shape/mod.rs

@@ -2,7 +2,7 @@ mod air;
 mod trace;
 
 pub use air::*;
-pub(in crate::proof_shape) use trace::*;
+pub(crate) use trace::*;
 
 #[cfg(feature = "cuda")]
 pub(crate) mod cuda;

crates/recursion/src/proof_shape/proof_shape/trace.rs

@@ -13,8 +13,29 @@ use crate::{
     },
     system::{Preflight, POW_CHECKER_HEIGHT},
     tracegen::RowMajorChip,
+    utils::interaction_length,
 };
 
+pub(crate) fn compute_air_shape_lookup_counts(
+    child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
+) -> Vec<usize> {
+    child_vk
+        .inner
+        .per_air
+        .iter()
+        .map(|avk| {
+            let dag = &avk.symbolic_constraints;
+            dag.constraints.nodes.len()
+                + avk.unused_variables.len()
+                + dag
+                    .interactions
+                    .iter()
+                    .map(interaction_length)
+                    .sum::<usize>()
+        })
+        .collect::<Vec<_>>()
+}
+
 #[derive(derive_new::new)]
 pub(in crate::proof_shape) struct ProofShapeChip<const NUM_LIMBS: usize, const LIMB_BITS: usize> {
     idx_encoder: Arc<Encoder>,
@@ -69,6 +90,7 @@ impl<const NUM_LIMBS: usize, const LIMB_BITS: usize> RowMajorChip<F>
             let mut total_interactions = 0usize;
             let mut cidx = 1usize;
             let mut num_present = 0usize;
+            let bc_air_shape_lookups = compute_air_shape_lookup_counts(child_vk);
 
             // Present AIRs
             for (idx, vdata) in &preflight.proof_shape.sorted_trace_vdata {
@@ -110,6 +132,7 @@ impl<const NUM_LIMBS: usize, const LIMB_BITS: usize> RowMajorChip<F>
                 total_interactions += num_interactions;
 
                 cols.n_max = F::from_usize(preflight.proof_shape.n_max);
+                cols.num_air_id_lookups = F::from_usize(bc_air_shape_lookups[*idx]);
 
                 let vcols: &mut ProofShapeVarColsMut<'_, F> = &mut borrow_var_cols_mut(
                     &mut chunk[cols_width..],

crates/recursion/src/system/frame.rs

@@ -8,8 +8,6 @@ use openvm_stark_backend::{
 };
 use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, Digest, F};
 
-use crate::utils::interaction_length;
-
 /*
  * Modified versions of the STARK and multi-STARK verifying keys for AirModule
  * implementations. AirModules should use MultiStarkVerifyingKeyFrame instead
@@ -22,7 +20,6 @@ use crate::utils::interaction_length;
 pub struct StarkVkeyFrame {
     pub preprocessed_data: Option<VerifierSinglePreprocessedData<Digest>>,
     pub params: StarkVerifyingParams,
-    pub num_dag_nodes: usize,
     pub num_interactions: usize,
     pub max_constraint_degree: u8,
     pub is_required: bool,
@@ -40,7 +37,6 @@ impl From<&StarkVerifyingKey<F, Digest>> for StarkVkeyFrame {
         Self {
             preprocessed_data: vk.preprocessed_data.clone(),
             params: vk.params.clone(),
-            num_dag_nodes: compute_num_dag_nodes(vk),
             num_interactions: vk.num_interactions(),
             max_constraint_degree: vk.max_constraint_degree,
             is_required: vk.is_required,
@@ -57,14 +53,3 @@ impl From<&MultiStarkVerifyingKey<BabyBearPoseidon2Config>> for MultiStarkVkeyFr
         }
     }
 }
-
-fn compute_num_dag_nodes<F, DIGEST>(vk: &StarkVerifyingKey<F, DIGEST>) -> usize {
-    let dag = &vk.symbolic_constraints;
-    dag.constraints.nodes.len()
-        + vk.unused_variables.len()
-        + dag
-            .interactions
-            .iter()
-            .map(interaction_length)
-            .sum::<usize>()
-}