feat: add debug to StarkEngineV2 (#211)

Slightly roundabout because we fix `SC` in some places, but for openvm
compatibility we need it to be generic.

Author: jonathanpwang

crates/stark-backend-v2/src/debug.rs

@@ -0,0 +1,135 @@
+// Merge into debug/mod.rs in v1:
+// - debug_constraints
+// - debug_constraints_and_interactions
+use std::sync::Arc;
+
+use itertools::{izip, Itertools};
+use openvm_stark_backend::{
+    air_builders::{
+        debug::{check_constraints, check_logup, USE_DEBUG_BUILDER},
+        symbolic::SymbolicConstraints,
+    },
+    prover::types::AirProofRawInput,
+    AirRef,
+};
+
+use crate::{
+    keygen::{types::StarkProvingKeyV2, MultiStarkKeygenBuilderV2},
+    prover::{
+        ColMajorMatrix, DeviceDataTransporterV2, ProverBackendV2, ProvingContextV2,
+        StridedColMajorMatrixView,
+    },
+    SystemParams, SC,
+};
+
+// TODO[jpw]: move into StarkEngineV2::debug default implementation after `SC` is made generic.
+/// `airs` should be the full list of all AIRs, not just used AIRs.
+pub fn debug_impl<PB: ProverBackendV2<Val = crate::F>, PD: DeviceDataTransporterV2<PB>>(
+    config: SystemParams,
+    device: &PD,
+    airs: &[AirRef<SC>],
+    ctx: &ProvingContextV2<PB>,
+) {
+    let mut keygen_builder = MultiStarkKeygenBuilderV2::new(config);
+    for air in airs {
+        keygen_builder.add_air(air.clone());
+    }
+    let pk = keygen_builder.generate_pk().unwrap();
+
+    let transpose = |mat: ColMajorMatrix<crate::F>| {
+        let row_major = StridedColMajorMatrixView::from(mat.as_view()).to_row_major_matrix();
+        Arc::new(row_major)
+    };
+    let (inputs, used_airs, used_pks): (Vec<_>, Vec<_>, Vec<_>) = ctx
+        .per_trace
+        .iter()
+        .map(|(air_id, air_ctx)| {
+            // Transfer from device **back** to host so the debugger can read the data.
+            let common_main = device.transport_matrix_from_device_to_host(&air_ctx.common_main);
+            let cached_mains = air_ctx
+                .cached_mains
+                .iter()
+                .map(|cd| transpose(device.transport_matrix_from_device_to_host(&cd.trace)))
+                .collect_vec();
+            let common_main = Some(transpose(common_main));
+            let public_values = air_ctx.public_values.clone();
+            (
+                AirProofRawInput {
+                    cached_mains,
+                    common_main,
+                    public_values,
+                },
+                airs[*air_id].clone(),
+                &pk.per_air[*air_id],
+            )
+        })
+        .multiunzip();
+
+    debug_constraints_and_interactions(&used_airs, &used_pks, &inputs);
+}
+
+/// The debugging will check the main AIR constraints and then separately check LogUp constraints by
+/// checking the actual multiset equalities. Currently it will not debug check any after challenge
+/// phase constraints for implementation simplicity.
+#[allow(dead_code)]
+#[allow(clippy::too_many_arguments)]
+pub fn debug_constraints_and_interactions(
+    airs: &[AirRef<SC>],
+    pk: &[&StarkProvingKeyV2],
+    inputs: &[AirProofRawInput<crate::F>],
+) {
+    USE_DEBUG_BUILDER.with(|debug| {
+        if *debug.lock().unwrap() {
+            let (main_parts_per_air, pvs_per_air): (Vec<_>, Vec<_>) = inputs
+                .iter()
+                .map(|input| {
+                    let mut main_parts = input
+                        .cached_mains
+                        .iter()
+                        .map(|trace| trace.as_view())
+                        .collect_vec();
+                    if let Some(trace) = input.common_main.as_ref() {
+                        main_parts.push(trace.as_view());
+                    }
+                    (main_parts, input.public_values.clone())
+                })
+                .unzip();
+            let preprocessed = izip!(airs, pk, &main_parts_per_air, &pvs_per_air)
+                .map(|(air, pk, main_parts, pvs)| {
+                    let preprocessed_trace = pk
+                        .preprocessed_data
+                        .as_ref()
+                        .map(|data| data.mat_view(0).to_row_major_matrix());
+                    tracing::debug!("Checking constraints for {}", air.name());
+                    check_constraints(
+                        air.as_ref(),
+                        &air.name(),
+                        &preprocessed_trace.as_ref().map(|t| t.as_view()),
+                        main_parts,
+                        pvs,
+                    );
+                    preprocessed_trace
+                })
+                .collect_vec();
+
+            let (air_names, interactions): (Vec<_>, Vec<_>) = pk
+                .iter()
+                .map(|pk| {
+                    let sym_constraints = SymbolicConstraints::from(&pk.vk.symbolic_constraints);
+                    (pk.air_name.clone(), sym_constraints.interactions)
+                })
+                .unzip();
+            let preprocessed_views = preprocessed
+                .iter()
+                .map(|t| t.as_ref().map(|t| t.as_view()))
+                .collect_vec();
+            check_logup(
+                &air_names,
+                &interactions,
+                &preprocessed_views,
+                &main_parts_per_air,
+                &pvs_per_air,
+            );
+        }
+    });
+}

crates/stark-backend-v2/src/engine.rs

@@ -9,6 +9,7 @@ use openvm_stark_backend::{config::StarkGenericConfig, prover::Prover, AirRef};
 use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
 
 use crate::{
+    debug::debug_impl,
     keygen::{
         types::{MultiStarkProvingKeyV2, MultiStarkVerifyingKeyV2},
         MultiStarkKeygenBuilderV2,
@@ -96,6 +97,11 @@ where
         verify(vk, proof, &mut transcript)
     }
 
+    /// The indexing of AIR ID in `ctx` should be consistent with the order of `airs`. In
+    /// particular, `airs` should correspond to the global proving key with all AIRs, including ones
+    /// not present in the `ctx`.
+    fn debug(&self, airs: &[AirRef<Self::SC>], ctx: &ProvingContextV2<Self::PB>);
+
     /// Runs a single end-to-end test for a given set of chips and traces partitions.
     /// This includes proving/verifying key generation, creating a proof, and verifying the proof.
     fn run_test(
@@ -139,12 +145,17 @@ where
     fn device(&self) -> &Self::PD {
         &self.device
     }
+
     fn prover_from_transcript(
         &self,
         transcript: TS,
     ) -> CoordinatorV2<Self::PB, Self::PD, Self::TS> {
         CoordinatorV2::new(CpuBackendV2, self.device.clone(), transcript)
     }
+
+    fn debug(&self, airs: &[AirRef<Self::SC>], ctx: &ProvingContextV2<Self::PB>) {
+        debug_impl::<Self::PB, Self::PD>(self.config().clone(), self.device(), airs, ctx);
+    }
 }
 
 // TODO[jpw]: move to stark-sdk

crates/stark-backend-v2/src/lib.rs

@@ -8,6 +8,7 @@ use p3_util::log2_ceil_u64;
 mod chip;
 pub mod codec;
 mod config;
+pub mod debug;
 pub mod dft;
 mod engine;
 pub mod keygen;

crates/stark-backend-v2/src/test_utils.rs

@@ -482,21 +482,25 @@ impl SystemParams {
     /// **These parameters should not be used in production!**
     pub fn new_for_testing(log_trace_height: usize) -> Self {
         let l_skip = 4;
-        test_system_params_small(4, log_trace_height - l_skip, 4)
+        let k_whir = 4;
+        let mut params = test_system_params_small(l_skip, log_trace_height - l_skip, k_whir);
+        params.max_constraint_degree = 4;
+        params
     }
 }
 
 /// Trace heights cannot exceed 2^{l_skip + n_stack} when using these system params.
 pub fn test_system_params_small(l_skip: usize, n_stack: usize, k_whir: usize) -> SystemParams {
     let log_final_poly_len = (n_stack + l_skip) % k_whir;
-    test_system_params_small_with_poly_len(l_skip, n_stack, k_whir, log_final_poly_len)
+    test_system_params_small_with_poly_len(l_skip, n_stack, k_whir, log_final_poly_len, 3)
 }
 
 pub fn test_system_params_small_with_poly_len(
     l_skip: usize,
     n_stack: usize,
     k_whir: usize,
     log_final_poly_len: usize,
+    max_constraint_degree: usize,
 ) -> SystemParams {
     assert!(log_final_poly_len < l_skip + n_stack);
     let log_blowup = 1;
@@ -507,7 +511,7 @@ pub fn test_system_params_small_with_poly_len(
         log_blowup,
         whir: test_whir_config_small(log_blowup, l_skip + n_stack, k_whir, log_final_poly_len),
         logup: log_up_security_params_baby_bear_100_bits(),
-        max_constraint_degree: 3,
+        max_constraint_degree,
     }
 }
 

crates/stark-backend/src/air_builders/debug/mod.rs

@@ -17,7 +17,7 @@ use crate::{
 
 mod check_constraints;
 
-use check_constraints::*;
+pub use check_constraints::*;
 
 use crate::interaction::BusIndex;