perf: minor optimiazation for image table assignment

DelphinusLab · May 28, 2024 · f65d1cd · f65d1cd
1 parent fa64c11
commit f65d1cd
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 59 deletions.
diff --git 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<F: FieldExt>(
     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::<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,
+        // 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::<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)))
                 };
 
-                cells.push(bn_to_field::<F>(
-                    &ImageTableEncoder::InitMemory.encode(entry.encode()),
-                ));
-            } else {
-                cells.push(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;
+                }
+            });
+        }
 
         Ok(cells)
     };

diff --git a/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs b/crates/zkwasm/src/circuits/zkwasm_circuit/mod.rs
@@ -348,39 +348,40 @@ impl<F: FieldExt> Circuit<F> for ZkWasmCircuit<F> {
             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<F> =
-                    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<F> =
+                        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();