feat(cuda): logup zerocheck (round 0) (#156)

Closes INT-5182

- [x] GPU state scaffolding (with CPU fallback)
- [x] Interaction metadata plumbing
- [x] Fractional sumcheck kernel
- [x] Interaction evaluation kernel 
- [x] round0 impl (iDFT > pad > DFT > eval DAG)
- [x] fold_ple
- [ ] sumcheck_polys_eval
- [ ] fold_mle
- [ ] column openings 
- [x] small heights support

---------

Co-authored-by: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>

Author: gaxiom

crates/stark-backend-v2/src/prover/cpu_backend.rs

@@ -71,9 +71,16 @@ impl<TS: FiatShamirTranscript> MultiRapProver<CpuBackendV2, TS> for CpuDeviceV2
         transcript: &mut TS,
         mpk: &DeviceMultiStarkProvingKeyV2<CpuBackendV2>,
         ctx: ProvingContextV2<CpuBackendV2>,
+        common_main_pcs_data: &StackedPcsData<F, Digest>,
     ) -> ((GkrProof, BatchConstraintProof), Vec<EF>) {
         let (gkr_proof, batch_constraint_proof, r) =
-            prove_zerocheck_and_logup::<_, _, TS, LogupZerocheckCpu>(self, transcript, mpk, ctx);
+            prove_zerocheck_and_logup::<_, _, TS, LogupZerocheckCpu>(
+                self,
+                transcript,
+                mpk,
+                ctx,
+                common_main_pcs_data,
+            );
         ((gkr_proof, batch_constraint_proof), r)
     }
 }

crates/stark-backend-v2/src/prover/hal.rs

@@ -83,6 +83,7 @@ pub trait MultiRapProver<PB: ProverBackendV2, TS> {
         transcript: &mut TS,
         mpk: &DeviceMultiStarkProvingKeyV2<PB>,
         ctx: ProvingContextV2<PB>,
+        common_main_pcs_data: &PB::PcsData,
     ) -> (Self::PartialProof, Self::Artifacts);
 }
 

crates/stark-backend-v2/src/prover/logup_zerocheck/cpu.rs

@@ -23,7 +23,7 @@ use crate::{
     poseidon2::sponge::FiatShamirTranscript,
     prover::{
         ColMajorMatrix, CpuBackendV2, CpuDeviceV2, DeviceMultiStarkProvingKeyV2,
-        LogupZerocheckProver, MatrixView, ProvingContextV2,
+        LogupZerocheckProver, MatrixView, ProverBackendV2, ProvingContextV2,
         fractional_sumcheck_gkr::{Frac, FracSumcheckProof, fractional_sumcheck},
         logup_zerocheck::EvalHelper,
         poly::evals_eq_hypercube,
@@ -58,14 +58,14 @@ pub struct LogupZerocheckCpu<'a> {
     pub xi: Vec<EF>,
     lambda_pows: Vec<EF>,
 
-    eq_xi_per_trace: Vec<ColMajorMatrix<EF>>,
-    eq_sharp_per_trace: Vec<ColMajorMatrix<EF>>,
+    pub eq_xi_per_trace: Vec<ColMajorMatrix<EF>>,
+    pub eq_sharp_per_trace: Vec<ColMajorMatrix<EF>>,
     eq_3b_per_trace: Vec<Vec<EF>>,
     // TODO[jpw]: delete these
     sels_per_trace_base: Vec<ColMajorMatrix<F>>,
     // After univariate round 0:
-    mat_evals_per_trace: Vec<Vec<ColMajorMatrix<EF>>>,
-    sels_per_trace: Vec<ColMajorMatrix<EF>>,
+    pub mat_evals_per_trace: Vec<Vec<ColMajorMatrix<EF>>>,
+    pub sels_per_trace: Vec<ColMajorMatrix<EF>>,
     // Stores \hat{f}(\vec r_n) * r_{n+1} .. r_{round-1} for polys f that are "done" in the batch
     // sumcheck
     zerocheck_tilde_evals: Vec<EF>,
@@ -81,6 +81,7 @@ where
         transcript: &mut TS,
         pk: &'a DeviceMultiStarkProvingKeyV2<CpuBackendV2>,
         ctx: &ProvingContextV2<CpuBackendV2>,
+        _common_main_data: &'a <CpuBackendV2 as ProverBackendV2>::PcsData,
         n_logup: usize,
         interactions_layout: StackedLayout,
         alpha_logup: EF,
@@ -374,6 +375,7 @@ where
                 ColMajorMatrix::new(mat, 3)
             })
             .collect_vec();
+
         // PERF[jpw]: see Gruen, Section 3.2 and 4 on some ways to reduce the degree of the
         // univariate polynomial. We know s_0 is supposed to vanish on univariate skip
         // domain `D` of size `2^{l_skip}`. Hence `s_0 = Z_D * s'_0(Z)` where `Z_D =

crates/stark-backend-v2/src/prover/logup_zerocheck/mod.rs

@@ -49,6 +49,7 @@ pub trait LogupZerocheckProver<'a, PB: ProverBackendV2, PD, TS>: Sized {
         transcript: &mut TS,
         pk: &'a DeviceMultiStarkProvingKeyV2<PB>,
         ctx: &ProvingContextV2<PB>,
+        common_main_pcs_data: &'a PB::PcsData,
         n_logup: usize,
         interactions_layout: StackedLayout,
         alpha_logup: PB::Challenge,
@@ -91,6 +92,7 @@ pub fn prove_zerocheck_and_logup<'a, PB, PD, TS, LZP>(
     transcript: &mut TS,
     mpk: &'a DeviceMultiStarkProvingKeyV2<PB>,
     ctx: ProvingContextV2<PB>,
+    common_main_pcs_data: &'a PB::PcsData,
 ) -> (GkrProof, BatchConstraintProof, Vec<PB::Challenge>)
 where
     PB: ProverBackendV2<Val = F, Challenge = EF>,
@@ -139,6 +141,7 @@ where
         transcript,
         mpk,
         &ctx,
+        common_main_pcs_data,
         n_logup,
         interactions_layout,
         alpha_logup,

crates/stark-backend-v2/src/prover/logup_zerocheck/single.rs

@@ -15,7 +15,7 @@ use crate::prover::{
 };
 
 /// For a single AIR
-pub(super) struct EvalHelper<'a, F> {
+pub struct EvalHelper<'a, F> {
     /// AIR constraints
     pub constraints_dag: &'a SymbolicExpressionDag<F>,
     /// Interactions

crates/stark-backend-v2/src/prover/mod.rs

@@ -146,7 +146,9 @@ where
             })
             .collect();
 
-        let (constraints_proof, r) = self.device.prove_rap_constraints(transcript, mpk, ctx);
+        let (constraints_proof, r) =
+            self.device
+                .prove_rap_constraints(transcript, mpk, ctx, &common_main_pcs_data);
 
         let opening_proof = self.device.prove_openings(
             transcript,

crates/stark-backend-v2/src/prover/stacked_pcs.rs

@@ -13,10 +13,11 @@ use crate::{
     prover::{ColMajorMatrix, MatrixView, StridedColMajorMatrixView, col_maj_idx, poly::Ple},
 };
 
-#[derive(Clone, Serialize, Deserialize, Debug)]
+#[derive(Clone, Serialize, Deserialize, Debug, CopyGetters)]
 pub struct StackedLayout {
     /// The minimum log2 height of a stacked slice. When stacking columns with smaller height, the
     /// column is expanded to `2^l_skip` by striding.
+    #[getset(get_copy = "pub")]
     l_skip: usize,
     /// The columns of the unstacked matrices in sorted order. Each entry `(matrix index, column
     /// index, coordinate)` contains the pointer `(matrix index, column index)` to a column of the

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

@@ -27,7 +27,8 @@ use crate::{
     proof::Proof,
     prover::{
         AirProvingContextV2, ColMajorMatrix, CommittedTraceDataV2, CpuBackendV2,
-        DeviceDataTransporterV2, ProverBackendV2, ProvingContextV2, stacked_pcs::stacked_commit,
+        DeviceDataTransporterV2, DeviceMultiStarkProvingKeyV2, MultiRapProver, ProverBackendV2,
+        ProvingContextV2, TraceCommitterV2, stacked_pcs::stacked_commit,
     },
 };
 
@@ -66,6 +67,25 @@ where
     ProvingContextV2::new(per_trace)
 }
 
+pub fn prove_up_to_batch_constraints<E: StarkEngineV2>(
+    engine: &E,
+    transcript: &mut E::TS,
+    pk: &DeviceMultiStarkProvingKeyV2<E::PB>,
+    ctx: ProvingContextV2<E::PB>,
+) -> (
+    <E::PD as MultiRapProver<E::PB, E::TS>>::PartialProof,
+    <E::PD as MultiRapProver<E::PB, E::TS>>::Artifacts,
+) {
+    let (_, common_main_pcs_data) = engine.device().commit(
+        &ctx.common_main_traces()
+            .map(|(_, trace)| trace)
+            .collect_vec(),
+    );
+    engine
+        .device()
+        .prove_rap_constraints(transcript, pk, ctx, &common_main_pcs_data)
+}
+
 fn get_fib_number(mut a: u32, mut b: u32, n: usize) -> u32 {
     for _ in 0..n - 1 {
         let c = (a + b) % BabyBear::ORDER_U32;

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

@@ -20,7 +20,7 @@ use crate::{
     },
     test_utils::{
         CachedFixture11, FibFixture, InteractionsFixture11, PreprocessedFibFixture, TestFixture,
-        test_engine_small,
+        prove_up_to_batch_constraints, test_engine_small,
     },
     verifier::{
         batch_constraints::{BatchConstraintError, verify_zerocheck_and_logup},
@@ -159,9 +159,8 @@ fn test_batch_sumcheck_zero_interactions(
 
     let pvs = vec![ctx.per_trace[0].1.public_values.clone()];
     let ((gkr_proof, batch_proof), _) =
-        engine
-            .device()
-            .prove_rap_constraints(&mut prover_sponge, &pk, ctx);
+        prove_up_to_batch_constraints(&engine, &mut prover_sponge, &pk, ctx);
+
     let r = verify_zerocheck_and_logup(
         &mut verifier_sponge,
         &vk.inner,
@@ -211,8 +210,12 @@ fn test_stacked_opening_reduction(log_trace_degree: usize) -> Result<(), Stacked
 
     let device = engine.device();
     // We need batch_proof to obtain the column openings
-    let ((_, batch_proof), r) =
-        device.prove_rap_constraints(&mut DuplexSponge::default(), &pk, ctx);
+    let ((_, batch_proof), r) = device.prove_rap_constraints(
+        &mut DuplexSponge::default(),
+        &pk,
+        ctx,
+        &common_main_pcs_data,
+    );
 
     let (stacking_proof, _) = prove_stacked_opening_reduction::<_, _, _, StackedReductionCpu>(
         device,
@@ -237,7 +240,6 @@ fn test_stacked_opening_reduction(log_trace_degree: usize) -> Result<(), Stacked
     assert_eq!(u_prism.len(), params.n_stack + 1);
     Ok(())
 }
-
 #[test_case(3)]
 #[test_case(2 ; "when fib log_height equals l_skip")]
 #[test_case(1 ; "when fib log_height less than l_skip")]
@@ -299,41 +301,8 @@ fn test_single_fib_and_dummy_trace_stark(log_trace_degree: usize) {
     per_trace.push((per_trace.len(), fib_ctx));
     let combined_ctx = ProvingContextV2::new(per_trace).into_sorted();
 
-    let l_skip = engine.config().l_skip;
-    let mut pvs = vec![vec![]; 3];
-    let (trace_id_to_air_ids, ns): (Vec<_>, Vec<_>) = combined_ctx
-        .per_trace
-        .iter()
-        .map(|(air_idx, air_ctx)| {
-            pvs[*air_idx] = air_ctx.public_values.clone();
-            (
-                *air_idx,
-                log2_strict_usize(air_ctx.common_main.height()) as isize - l_skip as isize,
-            )
-        })
-        .multiunzip();
-    let omega_pows = F::two_adic_generator(l_skip)
-        .powers()
-        .take(1 << l_skip)
-        .collect_vec();
-
-    let mut transcript = DuplexSponge::default();
-    let ((gkr_proof, batch_proof), _) =
-        engine
-            .device()
-            .prove_rap_constraints(&mut transcript, &combined_pk, combined_ctx);
-    let mut transcript = DuplexSponge::default();
-    verify_zerocheck_and_logup(
-        &mut transcript,
-        &combined_pk.get_vk().inner,
-        &pvs,
-        &gkr_proof,
-        &batch_proof,
-        &trace_id_to_air_ids,
-        &ns,
-        &omega_pows,
-    )
-    .unwrap();
+    let proof = engine.prove(&combined_pk, combined_ctx);
+    engine.verify(&combined_pk.get_vk(), &proof).unwrap();
 }
 
 #[test]
@@ -380,9 +349,7 @@ fn test_gkr_verify_zero_interactions() -> eyre::Result<()> {
     let pk = engine.device().transport_pk_to_device(&pk);
     let ctx = fx.generate_proving_ctx().into_sorted();
     let mut transcript = DuplexSponge::default();
-    let ((gkr_proof, _), _) = engine
-        .device()
-        .prove_rap_constraints(&mut transcript, &pk, ctx);
+    let ((gkr_proof, _), _) = prove_up_to_batch_constraints(&engine, &mut transcript, &pk, ctx);
 
     let mut transcript = DuplexSponge::default();
     assert!(transcript.check_witness(params.logup_pow_bits, gkr_proof.logup_pow_witness));
@@ -425,9 +392,7 @@ fn test_batch_constraints_with_interactions() -> eyre::Result<()> {
 
     let mut transcript = DuplexSponge::default();
     let ((gkr_proof, batch_proof), _) =
-        engine
-            .device()
-            .prove_rap_constraints(&mut transcript, &pk, ctx);
+        prove_up_to_batch_constraints(&engine, &mut transcript, &pk, ctx);
     let mut transcript = DuplexSponge::default();
     verify_zerocheck_and_logup(
         &mut transcript,