Skip to content

Commit

Permalink
program-table (#210)
Browse files Browse the repository at this point in the history 
_Issue #110._

---------

Co-authored-by: Aurélien Nicolas <[email protected]>
  • Loading branch information
@naure
naure and Aurélien Nicolas authored Sep 14, 2024
1 parent c9f6ba1 commit a8174ae
Show file tree
Hide file tree
Showing 12 changed files with 339 additions and 52 deletions.
4 changes: 2 additions & 2 deletions ceno_emul/src/rv32im.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ impl DecodedInstruction {
self.rd
}

pub fn func3(&self) -> u32 {
pub fn funct3(&self) -> u32 {
self.func3
}

Expand All @@ -221,7 +221,7 @@ impl DecodedInstruction {
self.rs2
}

pub fn func7(&self) -> u32 {
pub fn funct7(&self) -> u32 {
self.func7
}

Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/.gitignore
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
tracing.folded
27 changes: 23 additions & 4 deletions ceno_zkvm/examples/riscv_add.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
use std::time::Instant;

use ark_std::test_rng;
use ceno_zkvm::{instructions::riscv::addsub::AddInstruction, scheme::prover::ZKVMProver};
use ceno_zkvm::{
instructions::riscv::addsub::AddInstruction, scheme::prover::ZKVMProver,
tables::ProgramTableCircuit,
};
use clap::Parser;
use const_env::from_env;

Expand Down Expand Up @@ -90,11 +93,20 @@ fn main() {
let mut zkvm_cs = ZKVMConstraintSystem::default();
let add_config = zkvm_cs.register_opcode_circuit::<AddInstruction<E>>();
let range_config = zkvm_cs.register_table_circuit::<RangeTableCircuit<E>>();
let prog_config = zkvm_cs.register_table_circuit::<ProgramTableCircuit<E>>();

let mut zkvm_fixed_traces = ZKVMFixedTraces::default();
zkvm_fixed_traces.register_opcode_circuit::<AddInstruction<E>>(&zkvm_cs);
zkvm_fixed_traces
.register_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, range_config.clone());
zkvm_fixed_traces.register_table_circuit::<RangeTableCircuit<E>>(
&zkvm_cs,
range_config.clone(),
&(),
);
zkvm_fixed_traces.register_table_circuit::<ProgramTableCircuit<E>>(
&zkvm_cs,
prog_config.clone(),
&PROGRAM_ADD_LOOP,
);

let pk = zkvm_cs
.clone()
Expand Down Expand Up @@ -136,7 +148,14 @@ fn main() {
zkvm_witness.finalize_lk_multiplicities();
// assign table circuits
zkvm_witness
.assign_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, &range_config)
.assign_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, &range_config, &())
.unwrap();
zkvm_witness
.assign_table_circuit::<ProgramTableCircuit<E>>(
&zkvm_cs,
&prog_config,
&PROGRAM_ADD_LOOP.len(),
)
.unwrap();

let timer = Instant::now();
Expand Down
12 changes: 12 additions & 0 deletions ceno_zkvm/src/chip_handler/general.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use crate::{
expression::{Expression, Fixed, ToExpr, WitIn},
instructions::riscv::config::ExprLtConfig,
structs::ROMType,
tables::InsnRecord,
};

use super::utils::rlc_chip_record;
Expand Down Expand Up @@ -60,6 +61,17 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
self.cs.lk_table_record(name_fn, rlc_record, multiplicity)
}

/// Fetch an instruction at a given PC from the Program table.
pub fn lk_fetch(&mut self, record: &InsnRecord<Expression<E>>) -> Result<(), ZKVMError> {
let rlc_record = {
let mut fields = vec![E::BaseField::from(ROMType::Instruction as u64).expr()];
fields.extend_from_slice(record.as_slice());
self.rlc_chip_record(fields)
};

self.cs.lk_record(|| "fetch", rlc_record)
}

pub fn read_record<NR, N>(
&mut self,
name_fn: N,
Expand Down
70 changes: 39 additions & 31 deletions ceno_zkvm/src/instructions/riscv/addsub.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,10 @@ use itertools::Itertools;

use super::{
config::ExprLtConfig,
constants::{OPType, OpcodeType, RegUInt, PC_STEP_SIZE},
constants::{
OPType, OpcodeType, RegUInt, FUNCT3_ADD_SUB, FUNCT7_ADD, FUNCT7_SUB, OPCODE_OP,
PC_STEP_SIZE,
},
RIVInstruction,
};
use crate::{
Expand All @@ -16,6 +19,7 @@ use crate::{
expression::{ToExpr, WitIn},
instructions::{riscv::config::ExprLtInput, Instruction},
set_val,
tables::InsnRecord,
uint::UIntValue,
witness::LkMultiplicity,
};
Expand Down Expand Up @@ -93,6 +97,17 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
let rs2_id = circuit_builder.create_witin(|| "rs2_id")?;
let rd_id = circuit_builder.create_witin(|| "rd_id")?;

// Fetch the instruction.
circuit_builder.lk_fetch(&InsnRecord::new(
pc.expr(),
OPCODE_OP.into(),
rd_id.expr(),
FUNCT3_ADD_SUB.into(),
rs1_id.expr(),
rs2_id.expr(),
(if IS_ADD { FUNCT7_ADD } else { FUNCT7_SUB }).into(),
))?;

let prev_rs1_ts = circuit_builder.create_witin(|| "prev_rs1_ts")?;
let prev_rs2_ts = circuit_builder.create_witin(|| "prev_rs2_ts")?;
let prev_rd_ts = circuit_builder.create_witin(|| "prev_rd_ts")?;
Expand Down Expand Up @@ -145,6 +160,7 @@ fn add_sub_assignment<E: ExtensionField, const IS_ADD: bool>(
lk_multiplicity: &mut LkMultiplicity,
step: &StepRecord,
) -> Result<(), ZKVMError> {
lk_multiplicity.fetch(step.pc().before.0);
set_val!(instance, config.pc, step.pc().before.0 as u64);
set_val!(instance, config.ts, step.cycle());
let addend_1 = UIntValue::new_unchecked(step.rs2().unwrap().value);
Expand Down Expand Up @@ -268,15 +284,15 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction<E> {

#[cfg(test)]
mod test {
use ceno_emul::{Change, ReadOp, StepRecord, WriteOp};
use ceno_emul::{Change, ReadOp, StepRecord, WriteOp, CENO_PLATFORM};
use goldilocks::GoldilocksExt2;
use itertools::Itertools;
use multilinear_extensions::mle::IntoMLEs;

use crate::{
circuit_builder::{CircuitBuilder, ConstraintSystem},
instructions::{riscv::constants::PC_STEP_SIZE, Instruction},
scheme::mock_prover::MockProver,
scheme::mock_prover::{MockProver, MOCK_PC_ADD, MOCK_PC_SUB, MOCK_PROGRAM},
};

use super::{AddInstruction, SubInstruction};
Expand All @@ -302,22 +318,20 @@ mod test {
cb.cs.num_witin as usize,
vec![StepRecord {
cycle: 3,
pc: Change {
before: 2u32.into(),
after: (2u32 + PC_STEP_SIZE as u32).into(),
},
pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE),
insn_code: MOCK_PROGRAM[0],
rs1: Some(ReadOp {
addr: 2.into(),
addr: CENO_PLATFORM.register_vma(2).into(),
value: 11u32,
previous_cycle: 0,
}),
rs2: Some(ReadOp {
addr: 3.into(),
addr: CENO_PLATFORM.register_vma(3).into(),
value: 0xfffffffeu32,
previous_cycle: 0,
}),
rd: Some(WriteOp {
addr: 4.into(),
addr: CENO_PLATFORM.register_vma(4).into(),
value: Change {
before: 0u32,
after: 11u32.wrapping_add(0xfffffffeu32),
Expand Down Expand Up @@ -362,22 +376,20 @@ mod test {
cb.cs.num_witin as usize,
vec![StepRecord {
cycle: 3,
pc: Change {
before: 2u32.into(),
after: (2u32 + PC_STEP_SIZE as u32).into(),
},
pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE),
insn_code: MOCK_PROGRAM[0],
rs1: Some(ReadOp {
addr: 2.into(),
addr: CENO_PLATFORM.register_vma(2).into(),
value: u32::MAX - 1,
previous_cycle: 0,
}),
rs2: Some(ReadOp {
addr: 3.into(),
addr: CENO_PLATFORM.register_vma(3).into(),
value: u32::MAX - 1,
previous_cycle: 0,
}),
rd: Some(WriteOp {
addr: 4.into(),
addr: CENO_PLATFORM.register_vma(4).into(),
value: Change {
before: 0u32,
after: (u32::MAX - 1).wrapping_add(u32::MAX - 1),
Expand Down Expand Up @@ -422,22 +434,20 @@ mod test {
cb.cs.num_witin as usize,
vec![StepRecord {
cycle: 3,
pc: Change {
before: 2u32.into(),
after: (2u32 + PC_STEP_SIZE as u32).into(),
},
pc: Change::new(MOCK_PC_SUB, MOCK_PC_SUB + PC_STEP_SIZE),
insn_code: MOCK_PROGRAM[1],
rs1: Some(ReadOp {
addr: 2.into(),
addr: CENO_PLATFORM.register_vma(2).into(),
value: 11u32,
previous_cycle: 0,
}),
rs2: Some(ReadOp {
addr: 3.into(),
addr: CENO_PLATFORM.register_vma(3).into(),
value: 2u32,
previous_cycle: 0,
}),
rd: Some(WriteOp {
addr: 4.into(),
addr: CENO_PLATFORM.register_vma(4).into(),
value: Change {
before: 0u32,
after: 11u32.wrapping_sub(2u32),
Expand Down Expand Up @@ -482,22 +492,20 @@ mod test {
cb.cs.num_witin as usize,
vec![StepRecord {
cycle: 3,
pc: Change {
before: 2u32.into(),
after: (2u32 + PC_STEP_SIZE as u32).into(),
},
pc: Change::new(MOCK_PC_SUB, MOCK_PC_SUB + PC_STEP_SIZE),
insn_code: MOCK_PROGRAM[1],
rs1: Some(ReadOp {
addr: 2.into(),
addr: CENO_PLATFORM.register_vma(2).into(),
value: 3u32,
previous_cycle: 0,
}),
rs2: Some(ReadOp {
addr: 3.into(),
addr: CENO_PLATFORM.register_vma(3).into(),
value: 11u32,
previous_cycle: 0,
}),
rd: Some(WriteOp {
addr: 4.into(),
addr: CENO_PLATFORM.register_vma(4).into(),
value: Change {
before: 0u32,
after: 3u32.wrapping_sub(11u32),
Expand Down
5 changes: 5 additions & 0 deletions ceno_zkvm/src/instructions/riscv/constants.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,11 @@ use crate::uint::UInt;

pub(crate) const PC_STEP_SIZE: usize = 4;

pub const OPCODE_OP: usize = 0x33;
pub const FUNCT3_ADD_SUB: usize = 0;
pub const FUNCT7_ADD: usize = 0;
pub const FUNCT7_SUB: usize = 0x20;

#[allow(clippy::upper_case_acronyms)]
#[derive(Debug, Clone, Copy)]
pub enum OPType {
Expand Down
49 changes: 44 additions & 5 deletions ceno_zkvm/src/scheme/mock_prover.rs
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,32 @@
use super::utils::{eval_by_expr, wit_infer_by_expr};
use crate::{
circuit_builder::CircuitBuilder,
circuit_builder::{CircuitBuilder, ConstraintSystem},
expression::Expression,
scheme::utils::eval_by_expr_with_fixed,
structs::{ROMType, WitnessId},
tables::{ProgramTableCircuit, TableCircuit},
};
use ark_std::test_rng;
use ceno_emul::{ByteAddr, CENO_PLATFORM};
use ff_ext::ExtensionField;
use generic_static::StaticTypeMap;
use goldilocks::SmallField;
use itertools::Itertools;
use multilinear_extensions::virtual_poly_v2::ArcMultilinearExtension;
use std::{collections::HashSet, hash::Hash, marker::PhantomData, ops::Neg, sync::OnceLock};

/// The program baked in the MockProver.
/// TODO: Make this a parameter?
pub const MOCK_PROGRAM: &[u32] = &[
// add x4, x2, x3
0x00 << 25 | 3 << 20 | 2 << 15 | 4 << 7 | 0x33,
// sub x4, x2, x3
0x20 << 25 | 3 << 20 | 2 << 15 | 4 << 7 | 0x33,
];
// Addresses of particular instructions in the mock program.
pub const MOCK_PC_ADD: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start());
pub const MOCK_PC_SUB: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start() + 4);

#[allow(clippy::enum_variant_names)]
#[derive(Debug, PartialEq, Clone)]
pub(crate) enum MockProverError<E: ExtensionField> {
Expand Down Expand Up @@ -297,13 +312,37 @@ fn load_tables<E: ExtensionField>(cb: &CircuitBuilder<E>, challenge: [E; 2]) ->
}
}

fn load_program_table<E: ExtensionField>(
t_vec: &mut Vec<Vec<u8>>,
_cb: &CircuitBuilder<E>,
challenge: [E; 2],
) {
let mut cs = ConstraintSystem::<E>::new(|| "mock_program");
let mut cb = CircuitBuilder::new(&mut cs);
let config = ProgramTableCircuit::construct_circuit(&mut cb).unwrap();
let fixed =
ProgramTableCircuit::<E>::generate_fixed_traces(&config, cs.num_fixed, MOCK_PROGRAM);
for table_expr in &cs.lk_table_expressions {
for row in fixed.iter_rows() {
// TODO: Find a better way to obtain the row content.
let row = row
.iter()
.map(|v| unsafe { v.clone().assume_init() }.into())
.collect::<Vec<_>>();
let rlc_record = eval_by_expr_with_fixed(&row, &[], &challenge, &table_expr.values);
t_vec.push(rlc_record.to_repr().as_ref().to_vec());
}
}
}

let mut table_vec = vec![];
// TODO load more tables here
load_u5_table(&mut table_vec, cb, challenge);
load_u16_table(&mut table_vec, cb, challenge);
load_lt_table(&mut table_vec, cb, challenge);
load_and_table(&mut table_vec, cb, challenge);
load_ltu_table(&mut table_vec, cb, challenge);
load_program_table(&mut table_vec, cb, challenge);
HashSet::from_iter(table_vec)
}

Expand Down Expand Up @@ -336,22 +375,22 @@ fn load_once_tables<E: ExtensionField + 'static + Sync + Send>(

impl<'a, E: ExtensionField + Hash> MockProver<E> {
pub fn run_with_challenge(
cb: &mut CircuitBuilder<E>,
cb: &CircuitBuilder<E>,
wits_in: &[ArcMultilinearExtension<'a, E>],
challenge: [E; 2],
) -> Result<(), Vec<MockProverError<E>>> {
Self::run_maybe_challenge(cb, wits_in, Some(challenge))
}

pub fn run(
cb: &mut CircuitBuilder<E>,
cb: &CircuitBuilder<E>,
wits_in: &[ArcMultilinearExtension<'a, E>],
) -> Result<(), Vec<MockProverError<E>>> {
Self::run_maybe_challenge(cb, wits_in, None)
}

fn run_maybe_challenge(
cb: &mut CircuitBuilder<E>,
cb: &CircuitBuilder<E>,
wits_in: &[ArcMultilinearExtension<'a, E>],
challenge: Option<[E; 2]>,
) -> Result<(), Vec<MockProverError<E>>> {
Expand Down Expand Up @@ -451,7 +490,7 @@ impl<'a, E: ExtensionField + Hash> MockProver<E> {
}

pub fn assert_satisfied(
cb: &mut CircuitBuilder<E>,
cb: &CircuitBuilder<E>,
wits_in: &[ArcMultilinearExtension<'a, E>],
challenge: Option<[E; 2]>,
) {
Expand Down
Loading

0 comments on commit a8174ae

Please sign in to comment.