Skip to content

Commit

Permalink
opt post image table assignment
Browse files Browse the repository at this point in the history
  • Loading branch information
@junyu0312
junyu0312 committed May 28, 2024
1 parent da63a43 commit 458453c
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 39 deletions.
47 changes: 36 additions & 11 deletions crates/zkwasm/src/circuits/post_image_table/continuation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -310,8 +314,8 @@ impl<F: FieldExt> PostImageTableChip<F> {
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 =
Expand All @@ -320,15 +324,40 @@ impl<F: FieldExt> PostImageTableChip<F> {
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::<Vec<_>>()
};

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,
Expand All @@ -338,7 +367,7 @@ impl<F: FieldExt> PostImageTableChip<F> {
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,
Expand All @@ -351,18 +380,14 @@ impl<F: FieldExt> PostImageTableChip<F> {
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::<F>(&address) + *post),
)?;

rest_memory_writing_ops = rest_memory_writing_ops - F::one();
}

offset += 1;

Ok(entry)
})
.collect::<Result<Vec<_>, Error>>()
Expand Down
67 changes: 39 additions & 28 deletions crates/zkwasm/src/circuits/utils/image_table.rs
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -298,33 +299,43 @@ pub(crate) fn encode_compilation_table_values<F: FieldExt>(
{
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::<F>(&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::<F>(&ImageTableEncoder::InitMemory.encode(entry.encode()))
} else {
bn_to_field::<F>(&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::<Vec<_>>()
.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::<F>(&entry);

let addr = address.as_ptr();
unsafe {
let addr = addr as *mut F;

*addr.offset((pos + 1) as isize) = entry;
}
}
});
}

Ok(cells)
Expand Down

0 comments on commit 458453c

Please sign in to comment.