feat: hardware abstraction layer (#23)

* chore: move debug out of prover (#21)

* chore: debug checks don't use perm trace

* chore: move debug out of prove

* feat: hardware coordinator (#24)

* wip: add main prover backend traits

* chore: rename MatrixBuffer to MatrixView to reflect metadata

* feat: finished coordinator

* chore: delete circuit_api

* chore: remove unused commit.rs

* chore: mv VerifierConstraintFolder out of air_builders since it's not a builder anymore

* wip: cpu

* cpu prover done

* chore

* chore

* chore: rename Cpu to CpuBackend

* feat: data transport trait

* fix engine trait

* fix lifetimes

* fix all tests

* chore: rename trait

* fix bad regex

* chore: make device memory types owned (#28)

Even if they are buffers, make them owned types (no Clone) so that when
they are dropped, the device memory is automatically deallocated.

* revert Arc in RawAirProofInput

* chore: update test utils

* chore: add assert

Author: jonathanpwang

Cargo.toml

@@ -86,7 +86,7 @@ tracing = "0.1.40"
 serde_json = "1.0.117"
 lazy_static = "1.5.0"
 once_cell = "1.19.0"
-derive-new = "0.6.0"
+derive-new = "0.7.0"
 derive_more = "1.0.0"
 derivative = "2.2.0"
 strum_macros = "0.26.4"

crates/stark-backend/Cargo.toml

@@ -23,9 +23,10 @@ serde = { workspace = true, default-features = false, features = [
     "alloc",
     "rc",
 ] }
-derivative = { workspace = true }
+derivative.workspace = true
+derive-new.workspace = true
 metrics = { workspace = true, optional = true }
-cfg-if = { workspace = true }
+cfg-if.workspace = true
 thiserror.workspace = true
 async-trait.workspace = true
 rustc-hash.workspace = true

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

@@ -21,11 +21,7 @@ pub fn check_constraints<R, SC>(
     rap_name: &str,
     preprocessed: &Option<RowMajorMatrixView<Val<SC>>>,
     partitioned_main: &[RowMajorMatrixView<Val<SC>>],
-    after_challenge: &[RowMajorMatrixView<SC::Challenge>],
-    challenges: &[Vec<SC::Challenge>],
     public_values: &[Val<SC>],
-    exposed_values_after_challenge: &[Vec<SC::Challenge>],
-    rap_phase_seq_kind: RapPhaseSeqKind,
 ) where
     R: for<'a> Rap<DebugConstraintBuilder<'a, SC>>
         + BaseAir<Val<SC>>
@@ -35,7 +31,6 @@ pub fn check_constraints<R, SC>(
 {
     let height = partitioned_main[0].height();
     assert!(partitioned_main.iter().all(|mat| mat.height() == height));
-    assert!(after_challenge.iter().all(|mat| mat.height() == height));
 
     // Check that constraints are satisfied.
     (0..height).into_par_iter().for_each(|i| {
@@ -65,20 +60,6 @@ pub fn check_constraints<R, SC>(
             })
             .collect::<Vec<_>>();
 
-        let after_challenge_row_pair = after_challenge
-            .iter()
-            .map(|mat| (mat.row_slice(i), mat.row_slice(i_next)))
-            .collect::<Vec<_>>();
-        let after_challenge = after_challenge_row_pair
-            .iter()
-            .map(|(local, next)| {
-                VerticalPair::new(
-                    RowMajorMatrixView::new_row(local),
-                    RowMajorMatrixView::new_row(next),
-                )
-            })
-            .collect::<Vec<_>>();
-
         let mut builder = DebugConstraintBuilder {
             air_name: rap_name,
             row_index: i,
@@ -87,14 +68,14 @@ pub fn check_constraints<R, SC>(
                 RowMajorMatrixView::new_row(preprocessed_next.as_slice()),
             ),
             partitioned_main,
-            after_challenge,
-            challenges,
+            after_challenge: vec![], // unreachable
+            challenges: &[],         // unreachable
             public_values,
-            exposed_values_after_challenge,
+            exposed_values_after_challenge: &[], // unreachable
             is_first_row: Val::<SC>::ZERO,
             is_last_row: Val::<SC>::ZERO,
             is_transition: Val::<SC>::ONE,
-            rap_phase_seq_kind,
+            rap_phase_seq_kind: RapPhaseSeqKind::StarkLogUp, // unused
             has_common_main: rap.common_main_width() > 0,
         };
         if i == 0 {
@@ -111,7 +92,7 @@ pub fn check_constraints<R, SC>(
 
 pub fn check_logup<F: Field>(
     air_names: &[String],
-    interactions: &[&[SymbolicInteraction<F>]],
+    interactions: &[Vec<SymbolicInteraction<F>>],
     preprocessed: &[Option<RowMajorMatrixView<F>>],
     partitioned_main: &[Vec<RowMajorMatrixView<F>>],
     public_values: &[Vec<F>],

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

@@ -1,20 +1,79 @@
+use std::sync::{Arc, Mutex};
+
+use itertools::{izip, Itertools};
 use p3_air::{
     AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder, PermutationAirBuilder,
 };
 use p3_field::FieldAlgebra;
 use p3_matrix::{dense::RowMajorMatrixView, stack::VerticalPair};
 
-use super::{PartitionedAirBuilder, ViewPair};
+use super::{symbolic::SymbolicConstraints, PartitionedAirBuilder, ViewPair};
 use crate::{
     config::{StarkGenericConfig, Val},
     interaction::{
         rap::InteractionPhaseAirBuilder, Interaction, InteractionBuilder, InteractionType,
         RapPhaseSeqKind,
     },
-    rap::PermutationAirBuilderWithExposedValues,
+    keygen::types::StarkProvingKey,
+    rap::{AnyRap, PermutationAirBuilderWithExposedValues},
 };
 
-pub mod check_constraints;
+mod check_constraints;
+
+use check_constraints::*;
+
+thread_local! {
+   pub static USE_DEBUG_BUILDER: Arc<Mutex<bool>> = Arc::new(Mutex::new(true));
+}
+
+/// 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<SC: StarkGenericConfig>(
+    airs: &[Arc<dyn AnyRap<SC>>],
+    pk: &[StarkProvingKey<SC>],
+    main_views_per_air: &[Vec<RowMajorMatrixView<'_, Val<SC>>>],
+    public_values_per_air: &[Vec<Val<SC>>],
+) {
+    USE_DEBUG_BUILDER.with(|debug| {
+        if *debug.lock().unwrap() {
+            let preprocessed = izip!(airs, pk, main_views_per_air, public_values_per_air,)
+                .map(|(rap, pk, main, public_values)| {
+                    let preprocessed_trace = pk
+                        .preprocessed_data
+                        .as_ref()
+                        .map(|data| data.trace.as_view());
+                    tracing::debug!("Checking constraints for {}", rap.name());
+                    check_constraints(
+                        rap.as_ref(),
+                        &rap.name(),
+                        &preprocessed_trace,
+                        main,
+                        public_values,
+                    );
+                    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();
+            check_logup(
+                &air_names,
+                &interactions,
+                &preprocessed,
+                main_views_per_air,
+                public_values_per_air,
+            );
+        }
+    });
+}
 
 /// An `AirBuilder` which asserts that each constraint is zero, allowing any failed constraints to
 /// be detected early.

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

@@ -3,8 +3,9 @@ use p3_matrix::{dense::RowMajorMatrixView, stack::VerticalPair};
 
 pub mod debug;
 pub mod sub;
+/// AIR builder that collects the constraints expressed via the [Air](p3_air::Air) trait into
+/// a directed acyclic graph of symbolic expressions for serialization purposes.
 pub mod symbolic;
-pub mod verifier;
 
 pub type ViewPair<'a, T> = VerticalPair<RowMajorMatrixView<'a, T>, RowMajorMatrixView<'a, T>>;
 

crates/stark-backend/src/air_builders/symbolic/dag.rs

@@ -16,7 +16,7 @@ use crate::{
 /// Basically replace `Arc`s in `SymbolicExpression` with node IDs.
 /// Intended to be serializable and deserializable.
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
-#[serde(bound = "F: Field")]
+#[serde(bound(serialize = "F: Serialize", deserialize = "F: Deserialize<'de>"))]
 #[repr(C)]
 pub enum SymbolicExpressionNode<F> {
     Variable(SymbolicVariable<F>),
@@ -46,7 +46,8 @@ pub enum SymbolicExpressionNode<F> {
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
-#[serde(bound = "F: Field")]
+#[serde(bound(serialize = "F: Serialize", deserialize = "F: Deserialize<'de>"))]
+#[repr(C)]
 pub struct SymbolicExpressionDag<F> {
     /// Nodes in **topological** order.
     pub(crate) nodes: Vec<SymbolicExpressionNode<F>>,
@@ -67,7 +68,8 @@ impl<F> SymbolicExpressionDag<F> {
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
-#[serde(bound = "F: Field")]
+#[serde(bound(serialize = "F: Serialize", deserialize = "F: Deserialize<'de>"))]
+#[repr(C)] // TODO[jpw]: device transfer requires usize-independent serialization
 pub struct SymbolicConstraintsDag<F> {
     /// DAG with all symbolic expressions as nodes.
     /// A subset of the nodes represents all constraints that will be
@@ -115,6 +117,9 @@ pub(crate) fn build_symbolic_constraints_dag<F: Field>(
             }
         })
         .collect();
+    // Note[jpw]: there could be few nodes created after `constraint_idx` is built
+    // from `interactions` even though constraints already contain all interactions.
+    // This should be marginal and is not optimized for now.
     let constraints = SymbolicExpressionDag {
         nodes,
         constraint_idx,
@@ -252,23 +257,23 @@ impl<F: Field> SymbolicExpressionDag<F> {
 }
 
 // TEMPORARY conversions until we switch main interfaces to use SymbolicConstraintsDag
-impl<F: Field> From<SymbolicConstraintsDag<F>> for SymbolicConstraints<F> {
-    fn from(dag: SymbolicConstraintsDag<F>) -> Self {
+impl<'a, F: Field> From<&'a SymbolicConstraintsDag<F>> for SymbolicConstraints<F> {
+    fn from(dag: &'a SymbolicConstraintsDag<F>) -> Self {
         let exprs = dag.constraints.to_symbolic_expressions();
         let constraints = dag
             .constraints
             .constraint_idx
-            .into_iter()
-            .map(|idx| exprs[idx].as_ref().clone())
+            .iter()
+            .map(|&idx| exprs[idx].as_ref().clone())
             .collect::<Vec<_>>();
         let interactions = dag
             .interactions
-            .into_iter()
+            .iter()
             .map(|interaction| {
                 let fields = interaction
                     .fields
-                    .into_iter()
-                    .map(|idx| exprs[idx].as_ref().clone())
+                    .iter()
+                    .map(|&idx| exprs[idx].as_ref().clone())
                     .collect();
                 let count = exprs[interaction.count].as_ref().clone();
                 Interaction {
@@ -286,6 +291,12 @@ impl<F: Field> From<SymbolicConstraintsDag<F>> for SymbolicConstraints<F> {
     }
 }
 
+impl<F: Field> From<SymbolicConstraintsDag<F>> for SymbolicConstraints<F> {
+    fn from(dag: SymbolicConstraintsDag<F>) -> Self {
+        (&dag).into()
+    }
+}
+
 impl<F: Field> From<SymbolicConstraints<F>> for SymbolicConstraintsDag<F> {
     fn from(sc: SymbolicConstraints<F>) -> Self {
         build_symbolic_constraints_dag(&sc.constraints, &sc.interactions)

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

@@ -24,10 +24,12 @@ use crate::{
     rap::{BaseAirWithPublicValues, PermutationAirBuilderWithExposedValues, Rap},
 };
 
-pub mod dag;
+mod dag;
 pub mod symbolic_expression;
 pub mod symbolic_variable;
 
+pub use dag::*;
+
 /// Symbolic constraints for a single AIR with interactions.
 /// The constraints contain the constraints on the logup partial sums.
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]

crates/stark-backend/src/circuit_api.rs

@@ -15,6 +15,8 @@ where
     Com<Self>: Send + Sync,
     PcsProof<Self>: Send + Sync,
     PcsProverData<Self>: Send + Sync,
+    RapPhaseSeqPartialProof<Self>: Send + Sync,
+    RapPhaseSeqProvingKey<Self>: Send + Sync,
 {
     /// The PCS used to commit to trace polynomials.
     type Pcs: Pcs<Self::Challenge, Self::Challenger>;
@@ -115,6 +117,8 @@ where
     Pcs::ProverData: Send + Sync,
     Pcs::Proof: Send + Sync,
     Rps: RapPhaseSeq<<Pcs::Domain as PolynomialSpace>::Val, Challenge, Challenger>,
+    Rps::PartialProof: Send + Sync,
+    Rps::ProvingKey: Send + Sync,
     Challenger: FieldChallenger<<Pcs::Domain as PolynomialSpace>::Val>
         + CanObserve<<Pcs as p3_commit::Pcs<Challenge, Challenger>>::Commitment>
         + CanSample<Challenge>,