diff --git a/crates/zkwasm/src/circuits/utils/image_table.rs b/crates/zkwasm/src/circuits/utils/image_table.rs index f071909e8..bb5020205 100644 --- a/crates/zkwasm/src/circuits/utils/image_table.rs +++ b/crates/zkwasm/src/circuits/utils/image_table.rs @@ -1,6 +1,8 @@ use anyhow::Error; use halo2_proofs::arithmetic::FieldExt; use num_bigint::BigUint; +use rayon::iter::IntoParallelIterator; +use rayon::iter::ParallelIterator; use specs::brtable::BrTable; use specs::brtable::ElemTable; use specs::encode::image_table::ImageTableEncoder; @@ -29,18 +31,25 @@ pub(crate) struct InitMemoryLayouter { } impl InitMemoryLayouter { - fn for_each(self, mut f: impl FnMut((LocationType, u32))) { - for offset in 0..STACK_CAPABILITY { - f((LocationType::Stack, offset as u32)) - } + fn len(&self) -> usize { + STACK_CAPABILITY + GLOBAL_CAPABILITY + (self.pages * PAGE_ENTRIES) as usize + } - for offset in 0..GLOBAL_CAPABILITY { - f((LocationType::Global, offset as u32)) + fn memory_location_from_offset(&self, offset: usize) -> (LocationType, u32) { + let mut offset = offset; + + if offset < STACK_CAPABILITY { + return (LocationType::Stack, offset as u32); } - for offset in 0..(self.pages * PAGE_ENTRIES) { - f((LocationType::Heap, offset)) + offset -= STACK_CAPABILITY; + + if offset < GLOBAL_CAPABILITY { + return (LocationType::Global, offset as u32); } + + offset -= GLOBAL_CAPABILITY; + return (LocationType::Heap, offset as u32); } } @@ -275,40 +284,48 @@ pub(crate) fn encode_compilation_table_values( let padding_handler = |start, end| Ok(vec![F::zero(); end - start]); let init_memory_entries_handler = |_| { - let mut cells = vec![]; - - cells.push(bn_to_field( - &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), - )); - let layouter = InitMemoryLayouter { pages: page_capability, }; - layouter.for_each(|(ltype, offset)| { - if let Some(entry) = init_memory_table.try_find(ltype, offset) { - cells.push(bn_to_field::( - &ImageTableEncoder::InitMemory.encode(entry.encode()), - )); - } else if ltype == LocationType::Heap { - let entry = InitMemoryTableEntry { - ltype, - is_mutable: true, - offset, - vtype: VarType::I64, - value: 0, - eid: 0, + // The first entry is a default entry. + let mut cells = Vec::with_capacity(layouter.len() + 1); + cells.push(bn_to_field( + &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), + )); + unsafe { cells.set_len(layouter.len() + 1) }; + + { + let address = &cells; + + (0..layouter.len()).into_par_iter().for_each(|pos| { + let (ltype, offset) = layouter.memory_location_from_offset(pos); + + let entry = if let Some(entry) = init_memory_table.try_find(ltype, offset) { + bn_to_field::(&ImageTableEncoder::InitMemory.encode(entry.encode())) + } else if ltype == LocationType::Heap { + let entry = InitMemoryTableEntry { + ltype, + is_mutable: true, + offset, + vtype: VarType::I64, + value: 0, + eid: 0, + }; + + bn_to_field::(&ImageTableEncoder::InitMemory.encode(entry.encode())) + } else { + bn_to_field::(&ImageTableEncoder::InitMemory.encode(BigUint::from(0u64))) }; - cells.push(bn_to_field::( - &ImageTableEncoder::InitMemory.encode(entry.encode()), - )); - } else { - cells.push(bn_to_field::( - &ImageTableEncoder::InitMemory.encode(BigUint::from(0u64)), - )); - } - }); + let addr = address.as_ptr(); + unsafe { + let addr = addr as *mut F; + + *addr.offset((pos + 1) as isize) = entry; + } + }); + } Ok(cells) }; diff --git a/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs b/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs index 4a7b052db..d46baf4db 100644 --- a/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs +++ b/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs @@ -348,39 +348,40 @@ impl Circuit for ZkWasmCircuit { let _layouter = layouter.clone(); let _assigned_cells = assigned_cells.clone(); s.spawn(move |_| { - let pre_image_table = self.slice.encode_pre_compilation_table_values(config.k); + exec_with_profile!(|| "Assign pre image table chip", { + let pre_image_table = self.slice.encode_pre_compilation_table_values(config.k); - let cells = exec_with_profile!( - || "Assign pre image table chip", - image_chip + let cells = image_chip .assign(_layouter, &image_table_assigner, pre_image_table) - .unwrap() - ); + .unwrap(); - *_assigned_cells.pre_image_table_cells.lock().unwrap() = Some(cells); + *_assigned_cells.pre_image_table_cells.lock().unwrap() = Some(cells); + }); }); let _layouter = layouter.clone(); let _assigned_cells = assigned_cells.clone(); let _memory_writing_table = memory_writing_table.clone(); s.spawn(move |_| { - let post_image_table: ImageTableLayouter = - self.slice.encode_post_compilation_table_values(config.k); - - let (rest_memory_writing_ops, memory_finalized_set) = - _memory_writing_table.count_rest_memory_finalize_ops(); - - let cells = post_image_chip - .assign( - _layouter, - &image_table_assigner, - post_image_table, - rest_memory_writing_ops, - memory_finalized_set, - ) - .unwrap(); + exec_with_profile!(|| "Assign post image table", { + let post_image_table: ImageTableLayouter = + self.slice.encode_post_compilation_table_values(config.k); + + let (rest_memory_writing_ops, memory_finalized_set) = + _memory_writing_table.count_rest_memory_finalize_ops(); - *_assigned_cells.post_image_table_cells.lock().unwrap() = Some(cells); + let cells = post_image_chip + .assign( + _layouter, + &image_table_assigner, + post_image_table, + rest_memory_writing_ops, + memory_finalized_set, + ) + .unwrap(); + + *_assigned_cells.post_image_table_cells.lock().unwrap() = Some(cells); + }); }); let _layouter = layouter.clone();