Merge pull request #893 from Inversed-Tech/mihai/keccakf_integration

`keccakf` precompile + GKR-IOP integration

Author: spherel

Cargo.lock

@@ -36,6 +36,7 @@ cfg-if = "1.0"
 criterion = { version = "0.5", features = ["html_reports"] }
 crossbeam-channel = "0.5"
 itertools = "0.13"
+ndarray = "*"
 num-bigint = { version = "0.4.6" }
 num-derive = "0.4"
 num-traits = "0.2"
@@ -60,7 +61,7 @@ rand_core = "0.6"
 rand_xorshift = "0.3"
 rayon = "1.10"
 secp = "0.4.1"
-serde = { version = "1.0", features = ["derive"] }
+serde = { version = "1.0", features = ["derive", "rc"] }
 serde_json = "1.0"
 strum = "0.26"
 strum_macros = "0.26"

Cargo.toml

@@ -20,6 +20,10 @@ pub trait SyscallSpec {
     const REG_OPS_COUNT: usize;
     const MEM_OPS_COUNT: usize;
     const CODE: u32;
+
+    const HAS_LOOKUPS: bool = false;
+
+    const GKR_OUTPUTS: usize = 0;
 }
 
 /// Trace the inputs and effects of a syscall.

ceno_emul/src/syscalls.rs

@@ -16,6 +16,7 @@ impl SyscallSpec for KeccakSpec {
     const REG_OPS_COUNT: usize = 2;
     const MEM_OPS_COUNT: usize = KECCAK_WORDS;
     const CODE: u32 = ceno_rt::syscalls::KECCAK_PERMUTE;
+    const HAS_LOOKUPS: bool = true;
 }
 
 /// Wrapper type for the keccak_permute argument that implements conversions

ceno_emul/src/syscalls/keccak_permute.rs

@@ -231,7 +231,7 @@ fn test_ceno_rt_keccak() -> Result<()> {
     let steps = run(&mut state)?;
 
     // Expect the program to have written successive states between Keccak permutations.
-    const ITERATIONS: usize = 3;
+    const ITERATIONS: usize = 4;
     let keccak_outs = sample_keccak_f(ITERATIONS);
 
     let all_messages = read_all_messages(&state);

ceno_host/tests/test_elf.rs

@@ -19,9 +19,11 @@ serde_json.workspace = true
 base64 = "0.22"
 ceno_emul = { path = "../ceno_emul" }
 ff_ext = { path = "../ff_ext" }
+gkr_iop = { path = "../gkr_iop" }
 mpcs = { path = "../mpcs" }
 multilinear_extensions = { version = "0", path = "../multilinear_extensions" }
 p3 = { path = "../p3" }
+subprotocols = { path = "../subprotocols" }
 sumcheck = { version = "0", path = "../sumcheck" }
 transcript = { path = "../transcript" }
 witness = { path = "../witness" }

ceno_zkvm/Cargo.toml

@@ -95,6 +95,7 @@ fn bench_add(c: &mut Criterion) {
                                 "ADD",
                                 &prover.pk.pp,
                                 &circuit_pk,
+                                None,
                                 polys.into_iter().map(|mle| mle.into()).collect_vec(),
                                 commit,
                                 &[],

ceno_zkvm/benches/riscv_add.rs

@@ -1,5 +1,7 @@
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
+use gkr_iop::{ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuitWitness};
+use itertools::Itertools;
 use multilinear_extensions::util::max_usable_threads;
 use rayon::{
     iter::{IndexedParallelIterator, ParallelIterator},
@@ -67,3 +69,170 @@ pub trait Instruction<E: ExtensionField> {
         Ok((raw_witin, lk_multiplicity))
     }
 }
+
+pub struct GKRinfo {
+    pub and_lookups: usize,
+    pub xor_lookups: usize,
+    pub range_lookups: usize,
+    pub aux_wits: usize,
+}
+
+impl GKRinfo {
+    fn lookup_total(&self) -> usize {
+        self.and_lookups + self.xor_lookups + self.range_lookups
+    }
+}
+
+// Trait that should be implemented by opcodes which use the
+// GKR-IOP prover. Presently, such opcodes should also implement
+// the Instruction trait and in general methods of GKRIOPInstruction
+// will call corresponding methods of Instruction and do something extra
+// with respect to syncing state between GKR-IOP and old-style circuits/witnesses
+pub trait GKRIOPInstruction<E: ExtensionField>
+where
+    Self: Instruction<E>,
+{
+    type Layout: ProtocolWitnessGenerator<E> + ProtocolBuilder;
+
+    /// Similar to Instruction::construct_circuit; generally
+    /// meant to extend InstructionConfig with GKR-specific
+    /// fields
+    #[allow(unused_variables)]
+    fn construct_circuit_with_gkr_iop(
+        cb: &mut CircuitBuilder<E>,
+    ) -> Result<Self::InstructionConfig, ZKVMError> {
+        unimplemented!();
+    }
+
+    /// Should generate phase1 witness for GKR from step records
+    fn phase1_witness_from_steps(
+        layout: &Self::Layout,
+        steps: &[StepRecord],
+    ) -> Vec<Vec<E::BaseField>>;
+
+    /// Similar to Instruction::assign_instance, but with access to GKR lookups and wits
+    fn assign_instance_with_gkr_iop(
+        config: &Self::InstructionConfig,
+        instance: &mut [E::BaseField],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+        lookups: &[E::BaseField],
+        aux_wits: &[E::BaseField],
+    ) -> Result<(), ZKVMError>;
+
+    /// Similar to Instruction::assign_instances, but with access to the GKR layout.
+    #[allow(clippy::type_complexity)]
+    fn assign_instances_with_gkr_iop(
+        config: &Self::InstructionConfig,
+        num_witin: usize,
+        steps: Vec<StepRecord>,
+        gkr_layout: &Self::Layout,
+    ) -> Result<
+        (
+            RowMajorMatrix<E::BaseField>,
+            GKRCircuitWitness<E>,
+            LkMultiplicity,
+        ),
+        ZKVMError,
+    > {
+        let nthreads = max_usable_threads();
+        let num_instance_per_batch = if steps.len() > 256 {
+            steps.len().div_ceil(nthreads)
+        } else {
+            steps.len()
+        }
+        .max(1);
+        let lk_multiplicity = LkMultiplicity::default();
+        let mut raw_witin =
+            RowMajorMatrix::<E::BaseField>::new(steps.len(), num_witin, Self::padding_strategy());
+        let raw_witin_iter = raw_witin.par_batch_iter_mut(num_instance_per_batch);
+
+        let gkr_witness =
+            gkr_layout.gkr_witness(&Self::phase1_witness_from_steps(gkr_layout, &steps), &[]);
+
+        let (lookups, aux_wits) = {
+            // Extract lookups and auxiliary witnesses from GKR protocol
+            // Here we assume that the gkr_witness's last layer is a convenience layer which holds
+            // the output records for all instances; further, we assume that the last ```Self::lookup_count()```
+            // elements of this layer are the lookup arguments.
+            let mut lookups = vec![vec![]; steps.len()];
+            let mut aux_wits: Vec<Vec<E::BaseField>> = vec![vec![]; steps.len()];
+            let last_layer = gkr_witness.layers.last().unwrap().bases.clone();
+            let len = last_layer.len();
+            let lookup_total = Self::gkr_info().lookup_total();
+
+            let non_lookups = if len == 0 {
+                0
+            } else {
+                assert!(len >= lookup_total);
+                len - lookup_total
+            };
+
+            for witness in last_layer.iter().skip(non_lookups) {
+                for i in 0..witness.len() {
+                    lookups[i].push(witness[i]);
+                }
+            }
+
+            let n_layers = gkr_witness.layers.len();
+
+            for i in 0..steps.len() {
+                // Omit last layer, which stores outputs and not real witnesses
+                for layer in gkr_witness.layers[..n_layers - 1].iter() {
+                    for base in layer.bases.iter() {
+                        aux_wits[i].push(base[i]);
+                    }
+                }
+            }
+
+            (lookups, aux_wits)
+        };
+
+        raw_witin_iter
+            .zip(
+                steps
+                    .iter()
+                    .enumerate()
+                    .collect_vec()
+                    .par_chunks(num_instance_per_batch),
+            )
+            .flat_map(|(instances, steps)| {
+                let mut lk_multiplicity = lk_multiplicity.clone();
+                instances
+                    .chunks_mut(num_witin)
+                    .zip(steps)
+                    .map(|(instance, (i, step))| {
+                        Self::assign_instance_with_gkr_iop(
+                            config,
+                            instance,
+                            &mut lk_multiplicity,
+                            step,
+                            &lookups[*i],
+                            &aux_wits[*i],
+                        )
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Result<(), ZKVMError>>()?;
+
+        raw_witin.padding_by_strategy();
+        Ok((raw_witin, gkr_witness, lk_multiplicity))
+    }
+
+    /// Lookup and witness counts used by GKR proof
+    fn gkr_info() -> GKRinfo;
+
+    /// Returns corresponding column in RMM for the i-th
+    /// output evaluation of the GKR proof
+    #[allow(unused_variables)]
+    fn output_evals_map(i: usize) -> usize {
+        unimplemented!();
+    }
+
+    /// Returns corresponding column in RMM for the i-th
+    /// witness of the GKR proof
+    #[allow(unused_variables)]
+    fn witness_map(i: usize) -> usize {
+        unimplemented!();
+    }
+}