[precompile] separate prover backend for cpu/gpu (#972)

to support e2e integration, build on top of #971

Author: hero78119

Cargo.lock

@@ -3,18 +3,15 @@ use std::{collections::BTreeMap, time::Duration};
 use ceno_zkvm::{
     self,
     instructions::{Instruction, riscv::arith::AddInstruction},
-    scheme::{
-        cpu::{CpuBackend, CpuProver},
-        hal::ProofInput,
-        prover::ZKVMProver,
-    },
+    scheme::{hal::ProofInput, prover::ZKVMProver},
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces},
 };
 mod alloc;
 use criterion::*;
 
 use ceno_zkvm::scheme::constants::MAX_NUM_VARIABLES;
 use ff_ext::GoldilocksExt2;
+use gkr_iop::cpu::{CpuBackend, CpuProver};
 use mpcs::{BasefoldDefault, PolynomialCommitmentScheme, SecurityLevel};
 
 use rand::rngs::OsRng;

ceno_zkvm/benches/riscv_add.rs

@@ -3,7 +3,6 @@ use crate::{
     instructions::riscv::{DummyExtraConfig, MemPadder, MmuConfig, Rv32imConfig},
     scheme::{
         PublicValues, ZKVMProof,
-        cpu::{CpuBackend, CpuProver},
         mock_prover::{LkMultiplicityKey, MockProver},
         prover::ZKVMProver,
         verifier::ZKVMVerifier,
@@ -23,6 +22,7 @@ use clap::ValueEnum;
 use ff_ext::ExtensionField;
 #[cfg(debug_assertions)]
 use ff_ext::{Instrumented, PoseidonField};
+use gkr_iop::cpu::{CpuBackend, CpuProver};
 use itertools::{Itertools, MinMaxResult, chain};
 use mpcs::{PolynomialCommitmentScheme, SecurityLevel};
 use std::{

ceno_zkvm/src/e2e.rs

@@ -3,6 +3,7 @@ use ff_ext::ExtensionField;
 use gkr_iop::{
     ProtocolBuilder, ProtocolWitnessGenerator,
     gkr::{GKRCircuit, GKRCircuitOutput, GKRCircuitWitness},
+    hal::ProverBackend,
 };
 use multilinear_extensions::util::max_usable_threads;
 use rayon::{
@@ -95,7 +96,7 @@ pub trait GKRIOPInstruction<E: ExtensionField>
 where
     Self: Instruction<E>,
 {
-    type Layout<'a>: ProtocolWitnessGenerator<'a, E> + ProtocolBuilder<E>;
+    type Layout: ProtocolWitnessGenerator<E> + ProtocolBuilder<E>;
 
     /// Similar to Instruction::construct_circuit; generally
     /// meant to extend InstructionConfig with GKR-specific
@@ -108,8 +109,8 @@ where
     }
 
     /// Should generate phase1 witness for GKR from step records
-    fn phase1_witness_from_steps<'a>(
-        layout: &Self::Layout<'a>,
+    fn phase1_witness_from_steps(
+        layout: &Self::Layout,
         steps: &[StepRecord],
     ) -> RowMajorMatrix<E::BaseField>;
 
@@ -125,17 +126,17 @@ where
 
     /// Similar to Instruction::assign_instances, but with access to the GKR layout.
     #[allow(clippy::type_complexity)]
-    fn assign_instances_with_gkr_iop<'a>(
+    fn assign_instances_with_gkr_iop<'a, PB: ProverBackend<E = E>>(
         _config: &Self::InstructionConfig,
         _num_witin: usize,
         _steps: Vec<StepRecord>,
         _gkr_circuit: &GKRCircuit<E>,
-        _gkr_layout: &Self::Layout<'a>,
+        _gkr_layout: &Self::Layout,
     ) -> Result<
         (
             RowMajorMatrix<E::BaseField>,
-            GKRCircuitWitness<'a, E>,
-            GKRCircuitOutput<'a, E>,
+            GKRCircuitWitness<'a, PB>,
+            GKRCircuitOutput<'a, PB>,
             LkMultiplicity,
         ),
         ZKVMError,

ceno_zkvm/src/instructions.rs

@@ -124,7 +124,7 @@ impl<E: ExtensionField, S: SyscallSpec> Instruction<E> for LargeEcallDummy<E, S>
 }
 
 impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec> {
-    type Layout<'a> = KeccakLayout<E>;
+    type Layout = KeccakLayout<E>;
 
     fn gkr_info() -> crate::instructions::GKRinfo {
         GKRinfo {
@@ -191,8 +191,8 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
         }
     }
 
-    fn phase1_witness_from_steps<'a>(
-        layout: &Self::Layout<'a>,
+    fn phase1_witness_from_steps(
+        layout: &Self::Layout,
         steps: &[StepRecord],
     ) -> RowMajorMatrix<E::BaseField> {
         let instances = steps

ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs

@@ -1,6 +1,5 @@
 use super::hal::{
-    DeviceTransporter, MainSumcheckProver, MultilinearPolynomial, OpeningProver, ProverBackend,
-    ProverDevice, TowerProver, TraceCommitter,
+    DeviceTransporter, MainSumcheckProver, OpeningProver, ProverDevice, TowerProver, TraceCommitter,
 };
 use crate::{
     circuit_builder::ConstraintSystem,
@@ -17,8 +16,12 @@ use crate::{
 };
 use either::Either;
 use ff_ext::ExtensionField;
+use gkr_iop::{
+    cpu::{CpuBackend, CpuProver},
+    hal::ProverBackend,
+};
 use itertools::{Itertools, chain};
-use mpcs::{Point, PolynomialCommitmentScheme, SecurityLevel};
+use mpcs::{Point, PolynomialCommitmentScheme};
 use multilinear_extensions::{
     Expression, Instance,
     mle::{ArcMultilinearExtension, FieldType, IntoMLE, MultilinearExtension},
@@ -28,10 +31,7 @@ use multilinear_extensions::{
     virtual_poly::build_eq_x_r_vec,
     virtual_polys::VirtualPolynomialsBuilder,
 };
-use p3::{
-    field::{FieldAlgebra, TwoAdicField},
-    matrix::dense::RowMajorMatrix,
-};
+use p3::field::FieldAlgebra;
 use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
 use std::{collections::BTreeMap, sync::Arc};
 use sumcheck::{
@@ -42,63 +42,6 @@ use sumcheck::{
 use transcript::Transcript;
 use witness::next_pow2_instance_padding;
 
-pub struct CpuBackend<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
-    pub param: PCS::Param,
-    _marker: std::marker::PhantomData<E>,
-}
-
-impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> Default for CpuBackend<E, PCS> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> CpuBackend<E, PCS> {
-    pub fn new() -> Self {
-        let param =
-            PCS::setup(E::BaseField::TWO_ADICITY, SecurityLevel::Conjecture100bits).unwrap();
-        Self {
-            param,
-            _marker: std::marker::PhantomData,
-        }
-    }
-}
-
-impl<'a, E: ExtensionField> MultilinearPolynomial<E> for MultilinearExtension<'a, E> {
-    fn num_vars(&self) -> usize {
-        self.num_vars()
-    }
-
-    fn eval(&self, point: Point<E>) -> E {
-        self.evaluate(&point)
-    }
-}
-
-impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ProverBackend for CpuBackend<E, PCS> {
-    type E = E;
-    type Pcs = PCS;
-    type MultilinearPoly<'a> = MultilinearExtension<'a, E>;
-    type Matrix = RowMajorMatrix<E::BaseField>;
-    type PcsData = PCS::CommitmentWithWitness;
-}
-
-/// CPU prover for CPU backend
-pub struct CpuProver<PB: ProverBackend> {
-    backend: PB,
-    pp: Option<<<PB as ProverBackend>::Pcs as PolynomialCommitmentScheme<PB::E>>::ProverParam>,
-    largest_poly_size: Option<usize>,
-}
-
-impl<PB: ProverBackend> CpuProver<PB> {
-    pub fn new(backend: PB) -> Self {
-        Self {
-            backend,
-            pp: None,
-            largest_poly_size: None,
-        }
-    }
-}
-
 pub struct CpuTowerProver;
 
 impl CpuTowerProver {

ceno_zkvm/src/scheme/cpu/mod.rs

@@ -6,31 +6,13 @@ use crate::{
     structs::{TowerProofs, ZKVMProvingKey},
 };
 use ff_ext::ExtensionField;
+use gkr_iop::hal::ProverBackend;
 use mpcs::{Point, PolynomialCommitmentScheme};
 use multilinear_extensions::{mle::MultilinearExtension, util::ceil_log2};
 use sumcheck::structs::IOPProverMessage;
 use transcript::Transcript;
 use witness::next_pow2_instance_padding;
 
-pub trait MultilinearPolynomial<E: ExtensionField> {
-    fn num_vars(&self) -> usize;
-    fn eval(&self, point: Point<E>) -> E;
-}
-
-/// Defines basic types like field, pcs that are common among all devices
-/// and also defines the types that are specific to device.
-pub trait ProverBackend {
-    /// types that are common across all devices
-    type E: ExtensionField;
-    type Pcs: PolynomialCommitmentScheme<Self::E>;
-
-    /// device-specific types
-    // TODO: remove lifetime bound
-    type MultilinearPoly<'a>: Send + Sync + Clone + MultilinearPolynomial<Self::E>;
-    type Matrix: Send + Sync + Clone;
-    type PcsData;
-}
-
 pub trait ProverDevice<PB>:
     TraceCommitter<PB>
     + TowerProver<PB>

ceno_zkvm/src/scheme/hal.rs

@@ -1,11 +1,16 @@
 use ff_ext::ExtensionField;
+use gkr_iop::{
+    cpu::{CpuBackend, CpuProver},
+    hal::ProverBackend,
+};
 use std::{
     collections::{BTreeMap, HashMap},
     marker::PhantomData,
     sync::Arc,
 };
 
-use crate::scheme::hal::{MainSumcheckEvals, MultilinearPolynomial};
+use crate::scheme::hal::MainSumcheckEvals;
+use gkr_iop::hal::MultilinearPolynomial;
 use itertools::Itertools;
 use mpcs::{Point, PolynomialCommitmentScheme};
 use multilinear_extensions::{
@@ -27,11 +32,7 @@ use crate::{
     structs::{ProvingKey, TowerProofs, ZKVMProvingKey, ZKVMWitnesses},
 };
 
-use super::{
-    PublicValues, ZKVMChipProof, ZKVMProof,
-    cpu::{CpuBackend, CpuProver},
-    hal::{ProverBackend, ProverDevice},
-};
+use super::{PublicValues, ZKVMChipProof, ZKVMProof, hal::ProverDevice};
 
 type CreateTableProof<E> = (ZKVMChipProof<E>, HashMap<usize, E>, Point<E>);
 

ceno_zkvm/src/scheme/prover.rs

@@ -8,7 +8,7 @@ use crate::{
         riscv::{arith::AddInstruction, ecall::HaltInstruction},
     },
     scheme::{
-        cpu::{CpuBackend, CpuProver, CpuTowerProver},
+        cpu::CpuTowerProver,
         hal::{ProofInput, TowerProverSpec},
         prover::ZkVMCpuProver,
     },
@@ -25,6 +25,7 @@ use ceno_emul::{
     Platform, Program, StepRecord, VMState, encode_rv32,
 };
 use ff_ext::{ExtensionField, FieldInto, FromUniformBytes, GoldilocksExt2};
+use gkr_iop::cpu::{CpuBackend, CpuProver};
 use multilinear_extensions::{ToExpr, WitIn, mle::MultilinearExtension};
 
 #[cfg(debug_assertions)]

ceno_zkvm/src/scheme/tests.rs

@@ -16,8 +16,10 @@ clap.workspace = true
 either.workspace = true
 ff_ext = { path = "../ff_ext" }
 itertools.workspace = true
+mpcs = { path = "../mpcs" }
 multilinear_extensions = { version = "0.1.0", path = "../multilinear_extensions" }
 ndarray.workspace = true
+p3.workspace = true
 p3-field.workspace = true
 p3-goldilocks.workspace = true
 p3-util.workspace = true