continuation e2e in gpu (#1121)

reopen from #1118

Author: hero78119

Cargo.toml

@@ -97,7 +97,7 @@ opt-level = 3
 lto = "thin"
 
 # [patch."ssh://[email protected]/scroll-tech/ceno-gpu.git"]
-# ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal" }
+# ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal", default-features = false, features=["bb31"] }
 
 # [patch."https://github.com/scroll-tech/gkr-backend"]
 # ff_ext = { path = "../gkr-backend/crates/ff_ext", package = "ff_ext" }

ceno_zkvm/src/scheme/gpu/mod.rs

@@ -1,5 +1,6 @@
 use super::hal::{
-    DeviceTransporter, MainSumcheckProver, OpeningProver, ProverDevice, TowerProver, TraceCommitter,
+    DeviceTransporter, EccQuarkProver, MainSumcheckProver, OpeningProver, ProverDevice,
+    TowerProver, TraceCommitter,
 };
 use crate::{
     error::ZKVMError,
@@ -9,25 +10,15 @@ use crate::{
     },
     structs::{ComposedConstrainSystem, PointAndEval, TowerProofs},
 };
-use ceno_gpu::bb31::GpuPolynomialExt;
 use ff_ext::{ExtensionField, GoldilocksExt2};
 use gkr_iop::{
-    gkr::{
-        self, Evaluation, GKRProof, GKRProverOutput,
-        layer::{LayerWitness, gpu::utils::extract_mle_relationships_from_monomial_terms},
-    },
+    gkr::{self, Evaluation, GKRProof, GKRProverOutput, layer::LayerWitness},
     gpu::{GpuBackend, GpuProver},
-    hal::ProverBackend,
+    hal::{MultilinearPolynomial, ProverBackend},
 };
 use itertools::{Itertools, chain};
 use mpcs::{Point, PolynomialCommitmentScheme};
-use multilinear_extensions::{
-    Instance, WitnessId,
-    mle::{FieldType, MultilinearExtension},
-    monomialize_expr_to_wit_terms,
-    util::ceil_log2,
-};
-use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use multilinear_extensions::{mle::MultilinearExtension, util::ceil_log2};
 use std::{collections::BTreeMap, sync::Arc};
 use sumcheck::{
     macros::{entered_span, exit_span},
@@ -37,16 +28,20 @@ use sumcheck::{
 use transcript::{BasicTranscript, Transcript};
 use witness::next_pow2_instance_padding;
 
-use crate::circuit_builder::ConstraintSystem;
-use gkr_iop::hal::MultilinearPolynomial;
-
 #[cfg(feature = "gpu")]
 use gkr_iop::gpu::gpu_prover::*;
 
 pub struct GpuTowerProver;
 
-use crate::{e2e::ShardContext, scheme::constants::NUM_FANIN};
-use gkr_iop::gpu::{ArcMultilinearExtensionGpu, MultilinearExtensionGpu};
+use crate::{
+    e2e::ShardContext,
+    scheme::{constants::NUM_FANIN, cpu::CpuEccProver},
+    structs::EccQuarkProof,
+};
+use gkr_iop::{
+    gpu::{ArcMultilinearExtensionGpu, MultilinearExtensionGpu},
+    selector::SelectorContext,
+};
 
 // Extract out_evals from GPU-built tower witnesses
 #[allow(clippy::type_complexity)]
@@ -102,7 +97,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> TraceCommitter<GpuBa
     for GpuProver<GpuBackend<E, PCS>>
 {
     fn commit_traces<'a>(
-        &mut self,
+        &self,
         traces: BTreeMap<usize, witness::RowMajorMatrix<E::BaseField>>,
     ) -> (
         Vec<MultilinearExtensionGpu<'a, E>>,
@@ -534,13 +529,48 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<G
             gkr_circuit,
         } = composed_cs;
 
+        let num_instances = input.num_instances();
         let log2_num_instances = input.log2_num_instances();
         let num_threads = optimal_sumcheck_threads(log2_num_instances);
         let num_var_with_rotation = log2_num_instances + composed_cs.rotation_vars().unwrap_or(0);
 
         let Some(gkr_circuit) = gkr_circuit else {
             panic!("empty gkr circuit")
         };
+        let selector_ctxs = if cs.ec_final_sum.is_empty() {
+            // it's not global chip
+            vec![
+                SelectorContext {
+                    offset: 0,
+                    num_instances,
+                    num_vars: num_var_with_rotation,
+                };
+                gkr_circuit
+                    .layers
+                    .first()
+                    .map(|layer| layer.out_sel_and_eval_exprs.len())
+                    .unwrap_or(0)
+            ]
+        } else {
+            // it's global chip
+            vec![
+                SelectorContext {
+                    offset: 0,
+                    num_instances: input.num_instances[0],
+                    num_vars: num_var_with_rotation,
+                },
+                SelectorContext {
+                    offset: input.num_instances[0],
+                    num_instances: input.num_instances[1],
+                    num_vars: num_var_with_rotation,
+                },
+                SelectorContext {
+                    offset: 0,
+                    num_instances,
+                    num_vars: num_var_with_rotation,
+                },
+            ]
+        };
         let pub_io_mles = cs
             .instance_openings
             .iter()
@@ -574,7 +604,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<G
                 .collect_vec(),
             challenges,
             transcript,
-            num_instances,
+            &selector_ctxs,
         )?;
         assert_eq!(rt.len(), 1, "TODO support multi-layer gkr iop");
         Ok((
@@ -600,6 +630,34 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<G
     }
 }
 
+impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> EccQuarkProver<GpuBackend<E, PCS>>
+    for GpuProver<GpuBackend<E, PCS>>
+{
+    fn prove_ec_sum_quark<'a>(
+        &self,
+        num_instances: usize,
+        xs: Vec<Arc<MultilinearExtensionGpu<'a, E>>>,
+        ys: Vec<Arc<MultilinearExtensionGpu<'a, E>>>,
+        invs: Vec<Arc<MultilinearExtensionGpu<'a, E>>>,
+        transcript: &mut impl Transcript<E>,
+    ) -> Result<EccQuarkProof<E>, ZKVMError> {
+        // TODO implement GPU version of `create_ecc_proof`
+        let xs = xs.iter().map(|mle| mle.inner_to_mle().into()).collect_vec();
+        let ys = ys.iter().map(|mle| mle.inner_to_mle().into()).collect_vec();
+        let invs = invs
+            .iter()
+            .map(|mle| mle.inner_to_mle().into())
+            .collect_vec();
+        Ok(CpuEccProver::create_ecc_proof(
+            num_instances,
+            xs,
+            ys,
+            invs,
+            transcript,
+        ))
+    }
+}
+
 impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> OpeningProver<GpuBackend<E, PCS>>
     for GpuProver<GpuBackend<E, PCS>>
 {
@@ -743,8 +801,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> DeviceTransporter<Gp
         std::mem::forget(pcs_data_basefold);
         let pcs_data = Arc::new(pcs_data);
 
-        let fixed_mles =
-            PCS::get_arc_mle_witness_from_commitment(pk.fixed_commit_wd.as_ref().unwrap());
+        let fixed_mles = PCS::get_arc_mle_witness_from_commitment(pcs_data_original.as_ref());
         let fixed_mles = fixed_mles
             .iter()
             .map(|mle| Arc::new(MultilinearExtensionGpu::from_ceno(&cuda_hal, mle)))

ceno_zkvm/src/scheme/utils.rs

@@ -326,23 +326,6 @@ pub fn build_main_witness<
     );
     assert_eq!(input.fixed.len(), cs.num_fixed);
 
-    // check all witness size are power of 2
-    assert!(
-        input
-            .witness
-            .iter()
-            .all(|v| { v.evaluations_len() == 1 << num_var_with_rotation })
-    );
-
-    if !input.structural_witness.is_empty() {
-        assert!(
-            input
-                .structural_witness
-                .iter()
-                .all(|v| { v.evaluations_len() == 1 << num_var_with_rotation })
-        );
-    }
-
     let Some(gkr_circuit) = gkr_circuit else {
         panic!("empty gkr-iop")
     };
@@ -365,6 +348,17 @@ pub fn build_main_witness<
         .map(|instance| input.public_input[instance.0].clone())
         .collect_vec();
 
+    // check all witness size are power of 2
+    assert!(
+        input
+            .witness
+            .iter()
+            .chain(&input.structural_witness)
+            .chain(&input.fixed)
+            .chain(&pub_io_mles)
+            .all(|v| { v.evaluations_len() == 1 << num_var_with_rotation })
+    );
+
     let (_, gkr_circuit_out) = gkr_witness::<E, PCS, PB, PD>(
         gkr_circuit,
         &input.witness,

gkr_iop/src/gkr/layer/gpu/mod.rs

@@ -43,6 +43,7 @@ use crate::{
 use crate::gpu::{MultilinearExtensionGpu, gpu_prover::*};
 
 pub mod utils;
+use crate::selector::SelectorContext;
 use utils::*;
 
 impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> LinearLayerProver<GpuBackend<E, PCS>>
@@ -112,7 +113,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZerocheckLayerProver
         pub_io_evals: &[<GpuBackend<E, PCS> as ProverBackend>::E],
         challenges: &[<GpuBackend<E, PCS> as ProverBackend>::E],
         transcript: &mut impl Transcript<<GpuBackend<E, PCS> as ProverBackend>::E>,
-        num_instances: usize,
+        selector_ctxs: &[SelectorContext],
     ) -> (
         LayerProof<<GpuBackend<E, PCS> as ProverBackend>::E>,
         Point<<GpuBackend<E, PCS> as ProverBackend>::E>,
@@ -175,8 +176,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZerocheckLayerProver
             .out_sel_and_eval_exprs
             .iter()
             .zip(out_points.iter())
-            .map(|((sel_type, _), point)| {
-                build_eq_x_r_with_sel_gpu(&cuda_hal, point, num_instances, sel_type)
+            .zip(selector_ctxs.iter())
+            .map(|(((sel_type, _), point), selector_ctx)| {
+                build_eq_x_r_with_sel_gpu(&cuda_hal, point, selector_ctx, sel_type)
             })
             // for rotation left point
             .chain(

gkr_iop/src/gkr/layer/gpu/utils.rs

@@ -10,7 +10,7 @@ use multilinear_extensions::{
     Expression, mle::Point, monomial::Term, utils::eval_by_expr_constant,
 };
 
-use crate::selector::SelectorType;
+use crate::selector::{SelectorContext, SelectorType};
 
 use crate::gpu::{MultilinearExtensionGpu, gpu_prover::*};
 
@@ -63,7 +63,7 @@ pub fn extract_mle_relationships_from_monomial_terms<'a, E: ExtensionField>(
 pub fn build_eq_x_r_with_sel_gpu<E: ExtensionField>(
     hal: &CudaHalBB31,
     point: &Point<E>,
-    num_instances: usize,
+    selector_ctx: &SelectorContext,
     selector: &SelectorType<E>,
 ) -> MultilinearExtensionGpu<'static, E> {
     if std::any::TypeId::of::<E::BaseField>() != std::any::TypeId::of::<BB31Base>() {
@@ -74,12 +74,16 @@ pub fn build_eq_x_r_with_sel_gpu<E: ExtensionField>(
     let (num_instances, is_sp32, indices) = match selector {
         SelectorType::None => panic!("SelectorType::None"),
         SelectorType::Whole(_expr) => (eq_len, false, vec![]),
-        SelectorType::Prefix(_, _expr) => (num_instances, false, vec![]),
-        SelectorType::OrderedSparse32 { indices, .. } => (num_instances, true, indices.clone()),
+        SelectorType::Prefix(_expr) => (selector_ctx.num_instances, false, vec![]),
+        SelectorType::OrderedSparse32 { indices, .. } => {
+            (selector_ctx.num_instances, true, indices.clone())
+        }
+        SelectorType::QuarkBinaryTreeLessThan(..) => unimplemented!(),
     };
 
     // type eq
     let eq_mle = if is_sp32 {
+        assert_eq!(selector_ctx.offset, 0);
         let eq = build_eq_x_r_gpu(hal, point);
         let mut eq_buf = match eq.mle {
             GpuFieldType::Base(_) => panic!("should be ext field"),
@@ -103,6 +107,7 @@ pub fn build_eq_x_r_with_sel_gpu<E: ExtensionField>(
             &hal.inner,
             &gpu_points,
             &mut gpu_output,
+            selector_ctx.offset,
             num_instances,
         )
         .unwrap();
@@ -135,6 +140,7 @@ pub fn build_eq_x_r_gpu<E: ExtensionField>(
             &hal.inner,
             &gpu_points,
             &mut gpu_output,
+            0,
             eq_len,
         )
         .unwrap();