From 458453cb6eb1f2ab8e440543c4b7d197f49cb1cb Mon Sep 17 00:00:00 2001 From: Zhang Junyu Date: Tue, 28 May 2024 07:33:35 +0000 Subject: [PATCH] opt post image table assignment --- .../circuits/post_image_table/continuation.rs | 47 ++++++++++--- .../zkwasm/src/circuits/utils/image_table.rs | 67 +++++++++++-------- 2 files changed, 75 insertions(+), 39 deletions(-) diff --git a/crates/zkwasm/src/circuits/post_image_table/continuation.rs b/crates/zkwasm/src/circuits/post_image_table/continuation.rs index efdff1e80..02bdfe281 100644 --- a/crates/zkwasm/src/circuits/post_image_table/continuation.rs +++ b/crates/zkwasm/src/circuits/post_image_table/continuation.rs @@ -14,6 +14,10 @@ use halo2_proofs::plonk::ConstraintSystem; use halo2_proofs::plonk::Error; use halo2_proofs::plonk::Fixed; use num_bigint::BigUint; +use rayon::iter::IndexedParallelIterator; +use rayon::iter::IntoParallelIterator; +use rayon::iter::IntoParallelRefIterator; +use rayon::iter::ParallelIterator; use specs::encode::init_memory_table::encode_init_memory_table_address; use specs::encode::init_memory_table::MEMORY_ADDRESS_OFFSET; use specs::jtable::STATIC_FRAME_ENTRY_NUMBER; @@ -310,8 +314,8 @@ impl PostImageTableChip { Some(rest_memory_finalized_count_cell); } - let entries = { - // start from 'base_offset" because 'encode_compilation_table_values' have inserted an empty at the beginning. + let rest_memory_writing_ops = { + // start from 'base_offset" instead of 'base_offset + 1' because 'encode_compilation_table_values' have inserted an empty at the beginning. let mut offset = base_offset; let mut rest_memory_writing_ops = @@ -320,15 +324,40 @@ impl PostImageTableChip { post_image_table .init_memory_entries .iter() - .map(|post| { - let entry = ctx.borrow_mut().region.assign_advice( + .map(|_| { + let v = rest_memory_writing_ops; + + if memory_finalized_set + .contains(&image_table_offset_to_memory_location(offset)) + { + rest_memory_writing_ops = rest_memory_writing_ops - F::one(); + } + + offset += 1; + + v + }) + .collect::>() + }; + + let entries = { + post_image_table + .init_memory_entries + .par_iter() + .zip(rest_memory_writing_ops.into_par_iter()) + .enumerate() + .map(|(offset, (post, rest_memory_writing_ops))| { + // start from 'base_offset" instead of 'base_offset + 1' because 'encode_compilation_table_values' have inserted an empty at the beginning. + let offset = base_offset + offset; + + let entry = region.assign_advice( || "post image table: init memory", self.config.post_image_table, offset, || Ok(*post), )?; - ctx.borrow_mut().region.assign_advice( + region.assign_advice( || "post image table: updated memory count", self.config.rest_memory_finalized_count, offset, @@ -338,7 +367,7 @@ impl PostImageTableChip { let position = image_table_offset_to_memory_location(offset); if memory_finalized_set.contains(&position) { - ctx.borrow_mut().region.assign_advice( + region.assign_advice( || "post image table: init memory", self.config.update, offset, @@ -351,18 +380,14 @@ impl PostImageTableChip { position.1.into(), ) * MEMORY_ADDRESS_OFFSET; - ctx.borrow_mut().region.assign_advice( + region.assign_advice( || "post image table: init memory lookup", self.config.memory_finalized_lookup_encode, offset, || Ok(bn_to_field::(&address) + *post), )?; - - rest_memory_writing_ops = rest_memory_writing_ops - F::one(); } - offset += 1; - Ok(entry) }) .collect::, Error>>() diff --git a/crates/zkwasm/src/circuits/utils/image_table.rs b/crates/zkwasm/src/circuits/utils/image_table.rs index bb5020205..3af965938 100644 --- a/crates/zkwasm/src/circuits/utils/image_table.rs +++ b/crates/zkwasm/src/circuits/utils/image_table.rs @@ -1,8 +1,9 @@ use anyhow::Error; use halo2_proofs::arithmetic::FieldExt; use num_bigint::BigUint; -use rayon::iter::IntoParallelIterator; +use num_traits::Zero; use rayon::iter::ParallelIterator; +use rayon::prelude::ParallelSlice; use specs::brtable::BrTable; use specs::brtable::ElemTable; use specs::encode::image_table::ImageTableEncoder; @@ -298,33 +299,43 @@ pub(crate) fn encode_compilation_table_values( { 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))) - }; - - let addr = address.as_ptr(); - unsafe { - let addr = addr as *mut F; - - *addr.offset((pos + 1) as isize) = entry; - } - }); + const THREADS: usize = 4; + let chunk_size = (layouter.len() + THREADS - 1) / THREADS; + + (0..layouter.len()) + .collect::>() + .par_chunks(chunk_size) + .for_each(|chunk| { + for pos in chunk { + let (ltype, offset) = layouter.memory_location_from_offset(*pos); + + let entry = if let Some(entry) = init_memory_table.try_find(ltype, offset) { + 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, + }; + + ImageTableEncoder::InitMemory.encode(entry.encode()) + } else { + ImageTableEncoder::InitMemory.encode(BigUint::zero()) + }; + + let entry = bn_to_field::(&entry); + + let addr = address.as_ptr(); + unsafe { + let addr = addr as *mut F; + + *addr.offset((pos + 1) as isize) = entry; + } + } + }); } Ok(cells)