Skip to content

Commit 7611e70

Browse files
mcalanceamcalancea
committed
add gkr_iop state to ConstraintSystem/ZKVMWitness (wip, Apr 2)
1 parent 5e18d9a commit 7611e70

File tree

9 files changed

+118
-21
lines changed

9 files changed

+118
-21
lines changed

‎Cargo.lock

+2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎ceno_zkvm/Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ serde_json.workspace = true
1919
base64 = "0.22"
2020
ceno_emul = { path = "../ceno_emul" }
2121
ff_ext = { path = "../ff_ext" }
22+
gkr_iop = { path = "../gkr_iop" }
2223
mpcs = { path = "../mpcs" }
2324
multilinear_extensions = { version = "0", path = "../multilinear_extensions" }
2425
p3 = { path = "../p3" }

‎ceno_zkvm/src/structs.rs

+81-2
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,17 @@ use crate::{
22
circuit_builder::{CircuitBuilder, ConstraintSystem},
33
error::ZKVMError,
44
expression::Expression,
5-
instructions::Instruction,
5+
instructions::{Instruction, riscv::dummy::LargeEcallDummy},
66
state::StateCircuit,
77
tables::{RMMCollections, TableCircuit},
88
witness::LkMultiplicity,
99
};
10-
use ceno_emul::{CENO_PLATFORM, Platform, StepRecord};
10+
use ceno_emul::{CENO_PLATFORM, KeccakSpec, Platform, StepRecord};
1111
use ff_ext::ExtensionField;
12+
use gkr_iop::{
13+
ProtocolWitnessGenerator,
14+
precompiles::{KeccakLayout, KeccakTrace},
15+
};
1216
use itertools::{Itertools, chain};
1317
use mpcs::PolynomialCommitmentScheme;
1418
use multilinear_extensions::{
@@ -130,6 +134,13 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> VerifyingKey<E, PCS>
130134
}
131135
}
132136

137+
#[derive(Clone)]
138+
pub struct KeccakGKRIOP<E> {
139+
pub chip: gkr_iop::chip::Chip,
140+
pub layout: KeccakLayout<E>,
141+
pub circuit: gkr_iop::gkr::GKRCircuit,
142+
}
143+
133144
#[derive(Clone, Debug)]
134145
pub struct ProgramParams {
135146
pub platform: Platform,
@@ -154,6 +165,7 @@ pub struct ZKVMConstraintSystem<E: ExtensionField> {
154165
pub(crate) circuit_css: BTreeMap<String, ConstraintSystem<E>>,
155166
pub(crate) initial_global_state_expr: Expression<E>,
156167
pub(crate) finalize_global_state_expr: Expression<E>,
168+
pub keccak_gkr_iop: Option<KeccakGKRIOP<E>>,
157169
pub params: ProgramParams,
158170
}
159171

@@ -164,6 +176,7 @@ impl<E: ExtensionField> Default for ZKVMConstraintSystem<E> {
164176
initial_global_state_expr: Expression::ZERO,
165177
finalize_global_state_expr: Expression::ZERO,
166178
params: ProgramParams::default(),
179+
keccak_gkr_iop: None,
167180
}
168181
}
169182
}
@@ -175,6 +188,25 @@ impl<E: ExtensionField> ZKVMConstraintSystem<E> {
175188
..Default::default()
176189
}
177190
}
191+
192+
pub fn register_keccakf_circuit(
193+
&mut self,
194+
) -> <LargeEcallDummy<E, KeccakSpec> as Instruction<E>>::InstructionConfig {
195+
// Add GKR-IOP instance
196+
let params = gkr_iop::precompiles::KeccakParams {};
197+
let (layout, chip) = <KeccakLayout<E> as gkr_iop::ProtocolBuilder>::build(params);
198+
let circuit = chip.gkr_circuit();
199+
200+
assert!(self.keccak_gkr_iop.is_none());
201+
self.keccak_gkr_iop = Some(KeccakGKRIOP {
202+
layout,
203+
chip,
204+
circuit,
205+
});
206+
207+
self.register_opcode_circuit::<LargeEcallDummy<E, KeccakSpec>>()
208+
}
209+
178210
pub fn register_opcode_circuit<OC: Instruction<E>>(&mut self) -> OC::InstructionConfig {
179211
let mut cs = ConstraintSystem::new(|| format!("riscv_opcode/{}", OC::name()));
180212
let mut circuit_builder =
@@ -220,6 +252,14 @@ pub struct ZKVMFixedTraces<E: ExtensionField> {
220252
}
221253

222254
impl<E: ExtensionField> ZKVMFixedTraces<E> {
255+
pub fn register_keccakf_circuit(&mut self, _cs: &ZKVMConstraintSystem<E>) {
256+
assert!(
257+
self.circuit_fixed_traces
258+
.insert(LargeEcallDummy::<E, KeccakSpec>::name(), None)
259+
.is_none()
260+
);
261+
}
262+
223263
pub fn register_opcode_circuit<OC: Instruction<E>>(&mut self, _cs: &ZKVMConstraintSystem<E>) {
224264
assert!(self.circuit_fixed_traces.insert(OC::name(), None).is_none());
225265
}
@@ -244,6 +284,7 @@ impl<E: ExtensionField> ZKVMFixedTraces<E> {
244284

245285
#[derive(Default, Clone)]
246286
pub struct ZKVMWitnesses<E: ExtensionField> {
287+
keccak_trace: <KeccakLayout<E> as gkr_iop::ProtocolWitnessGenerator<E>>::Trace,
247288
witnesses_opcodes: BTreeMap<String, RowMajorMatrix<E::BaseField>>,
248289
witnesses_tables: BTreeMap<String, RMMCollections<E::BaseField>>,
249290
lk_mlts: BTreeMap<String, LkMultiplicity>,
@@ -263,6 +304,44 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
263304
self.lk_mlts.get(name)
264305
}
265306

307+
pub fn assign_keccakf_circuit(
308+
&mut self,
309+
css: &mut ZKVMConstraintSystem<E>,
310+
config: &<LargeEcallDummy<E, KeccakSpec> as Instruction<E>>::InstructionConfig,
311+
records: Vec<StepRecord>,
312+
) -> Result<(), ZKVMError> {
313+
// Ugly copy paste from assign_opcode_circuit, but we need to use the row major matrix
314+
let cs = css
315+
.get_cs(&LargeEcallDummy::<E, KeccakSpec>::name())
316+
.unwrap();
317+
let (witness, logup_multiplicity) = LargeEcallDummy::<E, KeccakSpec>::assign_instances(
318+
config,
319+
cs.num_witin as usize,
320+
records,
321+
)?;
322+
323+
// GKR-IOP-specific trace from row major witness
324+
self.keccak_trace = KeccakTrace::from(witness.clone());
325+
326+
assert!(
327+
self.witnesses_opcodes
328+
.insert(LargeEcallDummy::<E, KeccakSpec>::name(), witness)
329+
.is_none()
330+
);
331+
assert!(
332+
!self
333+
.witnesses_tables
334+
.contains_key(&LargeEcallDummy::<E, KeccakSpec>::name())
335+
);
336+
assert!(
337+
self.lk_mlts
338+
.insert(LargeEcallDummy::<E, KeccakSpec>::name(), logup_multiplicity)
339+
.is_none()
340+
);
341+
342+
Ok(())
343+
}
344+
266345
pub fn assign_opcode_circuit<OC: Instruction<E>>(
267346
&mut self,
268347
cs: &ZKVMConstraintSystem<E>,

‎gkr_iop/Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ version.workspace = true
1212
[dependencies]
1313
ark-std.workspace = true
1414
ff_ext = { path = "../ff_ext" }
15+
witness = { path = "../witness" }
1516
itertools.workspace = true
1617
multilinear_extensions = { version = "0.1.0", path = "../multilinear_extensions" }
1718
ndarray.workspace = true

‎gkr_iop/src/chip/protocol.rs

+4-4
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@ use super::Chip;
44

55
impl Chip {
66
/// Extract information from Chip that required in the GKR phase.
7-
pub fn gkr_circuit(&'_ self) -> GKRCircuit<'_> {
7+
pub fn gkr_circuit(&self) -> GKRCircuit {
88
GKRCircuit {
9-
layers: &self.layers,
9+
layers: self.layers.clone(),
1010
n_challenges: self.n_challenges,
1111
n_evaluations: self.n_evaluations,
12-
base_openings: &self.base_openings,
13-
ext_openings: &self.ext_openings,
12+
base_openings: self.base_openings.clone(),
13+
ext_openings: self.ext_openings.clone(),
1414
}
1515
}
1616
}

‎gkr_iop/src/gkr.rs

+9-9
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use ff_ext::ExtensionField;
2-
use itertools::{Itertools, chain, izip};
2+
use itertools::{chain, izip, Itertools};
33
use layer::{Layer, LayerWitness};
44
use subprotocols::{expression::Point, sumcheck::SumcheckProof};
55
use transcript::Transcript;
@@ -13,13 +13,13 @@ pub mod layer;
1313
pub mod mock;
1414

1515
#[derive(Clone, Debug)]
16-
pub struct GKRCircuit<'a> {
17-
pub layers: &'a [Layer],
16+
pub struct GKRCircuit {
17+
pub layers: Vec<Layer>,
1818

1919
pub n_challenges: usize,
2020
pub n_evaluations: usize,
21-
pub base_openings: &'a [(usize, EvalExpression)],
22-
pub ext_openings: &'a [(usize, EvalExpression)],
21+
pub base_openings: Vec<(usize, EvalExpression)>,
22+
pub ext_openings: Vec<(usize, EvalExpression)>,
2323
}
2424

2525
#[derive(Clone, Debug)]
@@ -42,7 +42,7 @@ pub struct Evaluation<E: ExtensionField> {
4242

4343
pub struct GKRClaims<Evaluation>(pub Vec<Evaluation>);
4444

45-
impl GKRCircuit<'_> {
45+
impl GKRCircuit {
4646
pub fn prove<E>(
4747
&self,
4848
circuit_wit: GKRCircuitWitness<E>,
@@ -56,7 +56,7 @@ impl GKRCircuit<'_> {
5656
let mut evaluations = out_evals.to_vec();
5757
evaluations.resize(self.n_evaluations, PointAndEval::default());
5858
let mut challenges = challenges.to_vec();
59-
let sumcheck_proofs = izip!(self.layers, circuit_wit.layers)
59+
let sumcheck_proofs = izip!(&self.layers, circuit_wit.layers)
6060
.map(|(layer, layer_wit)| {
6161
layer.prove(layer_wit, &mut evaluations, &mut challenges, transcript)
6262
})
@@ -85,7 +85,7 @@ impl GKRCircuit<'_> {
8585
let mut challenges = challenges.to_vec();
8686
let mut evaluations = out_evals.to_vec();
8787
evaluations.resize(self.n_evaluations, PointAndEval::default());
88-
for (layer, layer_proof) in izip!(self.layers, sumcheck_proofs) {
88+
for (layer, layer_proof) in izip!(self.layers.clone(), sumcheck_proofs) {
8989
layer.verify(layer_proof, &mut evaluations, &mut challenges, transcript)?;
9090
}
9191

@@ -99,7 +99,7 @@ impl GKRCircuit<'_> {
9999
evaluations: &[PointAndEval<E>],
100100
challenges: &[E],
101101
) -> Vec<Evaluation<E>> {
102-
chain!(self.base_openings, self.ext_openings)
102+
chain!(&self.base_openings, &self.ext_openings)
103103
.map(|(poly, eval)| {
104104
let poly = *poly;
105105
let PointAndEval { point, eval: value } = eval.evaluate(evaluations, challenges);

‎gkr_iop/src/gkr/mock.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use std::marker::PhantomData;
22

33
use ff_ext::ExtensionField;
4-
use itertools::{Itertools, izip};
4+
use itertools::{izip, Itertools};
55
use rand::rngs::OsRng;
66
use subprotocols::{
77
expression::{Expression, VectorType},
@@ -12,7 +12,7 @@ use thiserror::Error;
1212

1313
use crate::{evaluation::EvalExpression, utils::SliceIterator};
1414

15-
use super::{GKRCircuit, GKRCircuitWitness, layer::LayerType};
15+
use super::{layer::LayerType, GKRCircuit, GKRCircuitWitness};
1616

1717
pub struct MockProver<E: ExtensionField>(PhantomData<E>);
1818

@@ -35,7 +35,7 @@ pub enum MockProverError<F: ExtensionField> {
3535

3636
impl<E: ExtensionField> MockProver<E> {
3737
pub fn check(
38-
circuit: GKRCircuit<'_>,
38+
circuit: GKRCircuit,
3939
circuit_wit: &GKRCircuitWitness<E>,
4040
mut evaluations: Vec<VectorType<E>>,
4141
mut challenges: Vec<E>,

‎gkr_iop/src/precompiles/faster_keccak.rs

+13-2
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ use crate::{
1818
};
1919
use ndarray::{range, s, ArrayView, Ix2, Ix3};
2020

21+
use witness::RowMajorMatrix;
22+
2123
use super::utils::{u64s_to_felts, zero_eval, CenoLookup};
2224
use p3_field::{extension::BinomialExtensionField, PrimeCharacteristicRing};
2325

@@ -31,10 +33,10 @@ use transcript::BasicTranscript;
3133
type E = BinomialExtensionField<Goldilocks, 2>;
3234

3335
#[derive(Clone, Debug, Default)]
34-
struct KeccakParams {}
36+
pub struct KeccakParams {}
3537

3638
#[derive(Clone, Debug, Default)]
37-
struct KeccakLayout<E> {
39+
pub struct KeccakLayout<E> {
3840
params: KeccakParams,
3941

4042
committed_bits_id: usize,
@@ -727,10 +729,19 @@ impl<E: ExtensionField> ProtocolBuilder for KeccakLayout<E> {
727729
}
728730
}
729731

732+
#[derive(Clone, Default)]
730733
pub struct KeccakTrace {
731734
pub instances: Vec<[u32; KECCAK_INPUT_SIZE]>,
732735
}
733736

737+
use p3_field::Field;
738+
739+
impl<F: Field> From<RowMajorMatrix<F>> for KeccakTrace {
740+
fn from(value: RowMajorMatrix<F>) -> Self {
741+
unimplemented!();
742+
}
743+
}
744+
734745
impl<E> ProtocolWitnessGenerator<E> for KeccakLayout<E>
735746
where
736747
E: ExtensionField,

‎gkr_iop/src/precompiles/mod.rs

+4-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
mod faster_keccak;
22
mod keccak_f;
33
mod utils;
4-
pub use faster_keccak::run_faster_keccakf;
54
pub use keccak_f::run_keccakf;
5+
pub use {
6+
faster_keccak::run_faster_keccakf, faster_keccak::KeccakLayout, faster_keccak::KeccakParams,
7+
faster_keccak::KeccakTrace,
8+
};

0 commit comments

Comments
 (0)