Skip to content

Commit f347310

Browse files
committed
with mem bus chip build pass
1 parent d32c71f commit f347310

File tree

6 files changed

+172
-86
lines changed

6 files changed

+172
-86
lines changed

ceno_zkvm/src/e2e.rs

Lines changed: 27 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ use serde::Serialize;
3434
use std::{
3535
borrow::Cow,
3636
collections::{BTreeMap, BTreeSet, HashMap, HashSet},
37-
mem,
3837
sync::Arc,
3938
};
4039
use transcript::BasicTranscript as Transcript;
@@ -102,6 +101,7 @@ pub struct EmulationResult<'a> {
102101
}
103102

104103
pub struct RAMRecord {
104+
pub ram_type: RAMType,
105105
pub id: u64,
106106
pub addr: WordAddr,
107107
pub prev_cycle: Cycle,
@@ -115,14 +115,10 @@ pub struct ShardContext<'a> {
115115
max_num_shards: usize,
116116
max_cycle: Cycle,
117117
addr_future_accesses: Cow<'a, HashMap<(WordAddr, Cycle), Cycle>>,
118-
read_thread_based_record_storage: Either<
119-
Vec<[BTreeMap<WordAddr, RAMRecord>; mem::variant_count::<RAMType>()]>,
120-
&'a mut [BTreeMap<WordAddr, RAMRecord>; mem::variant_count::<RAMType>()],
121-
>,
122-
write_thread_based_record_storage: Either<
123-
Vec<[BTreeMap<WordAddr, RAMRecord>; mem::variant_count::<RAMType>()]>,
124-
&'a mut [BTreeMap<WordAddr, RAMRecord>; mem::variant_count::<RAMType>()],
125-
>,
118+
read_thread_based_record_storage:
119+
Either<Vec<BTreeMap<WordAddr, RAMRecord>>, &'a mut BTreeMap<WordAddr, RAMRecord>>,
120+
write_thread_based_record_storage:
121+
Either<Vec<BTreeMap<WordAddr, RAMRecord>>, &'a mut BTreeMap<WordAddr, RAMRecord>>,
126122
pub cur_shard_cycle_range: std::ops::Range<usize>,
127123
}
128124

@@ -137,13 +133,13 @@ impl<'a> Default for ShardContext<'a> {
137133
read_thread_based_record_storage: Either::Left(
138134
(0..max_threads)
139135
.into_par_iter()
140-
.map(|_| std::array::from_fn(|_| BTreeMap::new()))
136+
.map(|_| BTreeMap::new())
141137
.collect::<Vec<_>>(),
142138
),
143139
write_thread_based_record_storage: Either::Left(
144140
(0..max_threads)
145141
.into_par_iter()
146-
.map(|_| std::array::from_fn(|_| BTreeMap::new()))
142+
.map(|_| BTreeMap::new())
147143
.collect::<Vec<_>>(),
148144
),
149145
cur_shard_cycle_range: 0..usize::MAX,
@@ -181,14 +177,14 @@ impl<'a> ShardContext<'a> {
181177
read_thread_based_record_storage: Either::Left(
182178
(0..max_threads)
183179
.into_par_iter()
184-
.map(|_| std::array::from_fn(|_| BTreeMap::new()))
180+
.map(|_| BTreeMap::new())
185181
.collect::<Vec<_>>(),
186182
),
187183
// TODO with_capacity optimisation
188184
write_thread_based_record_storage: Either::Left(
189185
(0..max_threads)
190186
.into_par_iter()
191-
.map(|_| std::array::from_fn(|_| BTreeMap::new()))
187+
.map(|_| BTreeMap::new())
192188
.collect::<Vec<_>>(),
193189
),
194190
cur_shard_cycle_range,
@@ -220,6 +216,20 @@ impl<'a> ShardContext<'a> {
220216
}
221217
}
222218

219+
pub fn read_records(&self) -> &[BTreeMap<WordAddr, RAMRecord>] {
220+
match &self.read_thread_based_record_storage {
221+
Either::Left(m) => m,
222+
Either::Right(_) => panic!("undefined behaviour"),
223+
}
224+
}
225+
226+
pub fn write_records(&self) -> &[BTreeMap<WordAddr, RAMRecord>] {
227+
match &self.write_thread_based_record_storage {
228+
Either::Left(m) => m,
229+
Either::Right(_) => panic!("undefined behaviour"),
230+
}
231+
}
232+
223233
#[inline(always)]
224234
pub fn is_first_shard(&self) -> bool {
225235
self.shard_id == 0
@@ -269,9 +279,10 @@ impl<'a> ShardContext<'a> {
269279
.as_mut()
270280
.right()
271281
.expect("illegal type");
272-
ram_record[ram_type as usize].insert(
282+
ram_record.insert(
273283
addr,
274284
RAMRecord {
285+
ram_type,
275286
id,
276287
addr,
277288
prev_cycle,
@@ -291,9 +302,10 @@ impl<'a> ShardContext<'a> {
291302
.as_mut()
292303
.right()
293304
.expect("illegal type");
294-
ram_record[ram_type as usize].insert(
305+
ram_record.insert(
295306
addr,
296307
RAMRecord {
308+
ram_type,
297309
id,
298310
addr,
299311
prev_cycle,

ceno_zkvm/src/instructions/riscv/rv32im.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -416,7 +416,7 @@ impl<E: ExtensionField> Rv32imConfig<E> {
416416
let mut secp256k1_double_records = Vec::new();
417417
steps
418418
.into_iter()
419-
.filter_map(|mut step| {
419+
.filter_map(|step| {
420420
if shard_ctx.is_current_shard_cycle(step.cycle()) {
421421
Some(step)
422422
} else {

ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ pub struct MmuConfig<'a, E: ExtensionField> {
3131
/// finalized circuit for all MMIO
3232
pub local_final_circuit: <LocalFinalCircuit<'a, E> as TableCircuit<E>>::TableConfig,
3333
/// ram bus to deal with cross shard read/write
34-
pub ram_bus_circuit: <RBCircuit<E> as TableCircuit<E>>::TableConfig,
34+
pub ram_bus_circuit: <RBCircuit<'a, E> as TableCircuit<E>>::TableConfig,
3535
pub params: ProgramParams,
3636
}
3737

@@ -160,7 +160,7 @@ impl<E: ExtensionField> MmuConfig<'_, E> {
160160
&(shard_ctx, all_records.as_slice()),
161161
)?;
162162

163-
witness.assign_table_circuit::<RBCircuit<E>>(cs, &self.ram_bus_circuit, todo!())?;
163+
witness.assign_table_circuit::<RBCircuit<E>>(cs, &self.ram_bus_circuit, shard_ctx)?;
164164

165165
Ok(())
166166
}

ceno_zkvm/src/tables/ram.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,4 +159,4 @@ impl NonVolatileTable for PubIOTable {
159159

160160
pub type PubIOCircuit<E> = PubIORamCircuit<E, PubIOTable>;
161161
pub type LocalFinalCircuit<'a, E> = LocalFinalRamCircuit<'a, UINT_LIMBS, E>;
162-
pub type RBCircuit<E> = RamBusCircuit<UINT_LIMBS, E>;
162+
pub type RBCircuit<'a, E> = RamBusCircuit<'a, UINT_LIMBS, E>;

ceno_zkvm/src/tables/ram/ram_circuit.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use super::ram_impl::{
55
};
66
use crate::{
77
circuit_builder::CircuitBuilder,
8-
e2e::{RAMRecord, ShardContext},
8+
e2e::ShardContext,
99
error::ZKVMError,
1010
structs::{ProgramParams, RAMType},
1111
tables::{RMMCollections, TableCircuit},
@@ -323,12 +323,14 @@ impl<'a, E: ExtensionField, const V_LIMBS: usize> TableCircuit<E>
323323
}
324324

325325
/// This circuit is generalized version to handle all mmio records
326-
pub struct RamBusCircuit<const V_LIMBS: usize, E>(PhantomData<E>);
326+
pub struct RamBusCircuit<'a, const V_LIMBS: usize, E>(PhantomData<(&'a (), E)>);
327327

328-
impl<E: ExtensionField, const V_LIMBS: usize> TableCircuit<E> for RamBusCircuit<V_LIMBS, E> {
328+
impl<'a, E: ExtensionField, const V_LIMBS: usize> TableCircuit<E>
329+
for RamBusCircuit<'a, V_LIMBS, E>
330+
{
329331
type TableConfig = RAMBusConfig<V_LIMBS>;
330332
type FixedInput = ();
331-
type WitnessInput = (&'static [RAMRecord], &'static [RAMRecord]);
333+
type WitnessInput = ShardContext<'a>;
332334

333335
fn name() -> String {
334336
"RamBusCircuit".to_string()
@@ -357,16 +359,14 @@ impl<E: ExtensionField, const V_LIMBS: usize> TableCircuit<E> for RamBusCircuit<
357359
num_witin: usize,
358360
num_structural_witin: usize,
359361
_multiplicity: &[HashMap<u64, usize>],
360-
final_v: &Self::WitnessInput,
362+
shard_ctx: &Self::WitnessInput,
361363
) -> Result<RMMCollections<E::BaseField>, ZKVMError> {
362-
let (global_read_mem, global_write_mem) = *final_v;
363364
// assume returned table is well-formed include padding
364365
Ok(Self::TableConfig::assign_instances(
365366
config,
367+
shard_ctx,
366368
num_witin,
367369
num_structural_witin,
368-
global_read_mem,
369-
global_write_mem,
370370
)?)
371371
}
372372
}

0 commit comments

Comments
 (0)