second slice pass

DelphinusLab · Nov 6, 2023 · ccdaa76 · ccdaa76
1 parent 5474fad
commit ccdaa76
Show file tree

Hide file tree

Showing 48 changed files with 640 additions and 289 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -12,3 +12,6 @@ wasmi = { path = "third-party/wasmi" }
 
 [profile.dev]
 opt-level = 3
+
+[patch."https://github.com/DelphinusLab/halo2-gpu-specific.git"]
+halo2_proofs = { path = "./halo2-gpu-specific/halo2_proofs" }
diff --git a/crates/specs/src/etable.rs b/crates/specs/src/etable.rs
@@ -2,6 +2,8 @@ use serde::Serialize;
 
 use super::itable::InstructionTableEntry;
 use crate::host_function::HostPlugin;
+use crate::mtable::AccessType;
+use crate::mtable::MemoryWritingTable;
 use crate::step::StepInfo;
 
 #[derive(Clone, Debug, Serialize)]
@@ -69,3 +71,84 @@ impl EventTable {
             .collect::<Vec<_>>()
     }
 }
+
+use std::cmp::Ordering;
+
+use crate::mtable::MemoryTableEntry;
+
+#[derive(Debug)]
+pub struct MemoryRWEntry {
+    pub entry: MemoryTableEntry,
+    pub start_eid: u32,
+    pub end_eid: u32,
+}
+
+#[derive(Debug)]
+pub struct EventTableEntryWithMemoryInfo<'a> {
+    pub eentry: &'a EventTableEntry,
+    pub memory_rw_entires: Vec<MemoryRWEntry>,
+}
+
+#[derive(Debug)]
+pub struct EventTableWithMemoryInfo<'a>(pub Vec<EventTableEntryWithMemoryInfo<'a>>);
+
+impl<'a> EventTableWithMemoryInfo<'a> {
+    pub fn new(
+        event_table: &'a EventTable,
+        memory_writing_table: &MemoryWritingTable,
+        memory_event_of_step: fn(&EventTableEntry, &mut u32) -> Vec<MemoryTableEntry>,
+    ) -> Self {
+        let lookup = memory_writing_table.build_lookup_mapping();
+
+        let lookup_mtable_eid = |(eid, ltype, offset, is_writing)| {
+            let records = lookup.get(&(ltype, offset)).unwrap();
+
+            if is_writing {
+                let idx = records
+                    .binary_search_by(|(start_eid, _)| start_eid.cmp(eid))
+                    .unwrap();
+                records[idx]
+            } else {
+                let idx = records
+                    .binary_search_by(|(start_eid, end_eid)| {
+                        if eid <= start_eid {
+                            Ordering::Greater
+                        } else if eid > end_eid {
+                            Ordering::Less
+                        } else {
+                            Ordering::Equal
+                        }
+                    })
+                    .unwrap();
+                records[idx]
+            }
+        };
+
+        EventTableWithMemoryInfo(
+            event_table
+                .entries()
+                .iter()
+                .map(|eentry| EventTableEntryWithMemoryInfo {
+                    eentry,
+                    memory_rw_entires: memory_event_of_step(eentry, &mut 1)
+                        .into_iter()
+                        .map(|mentry| {
+                            let (start_eid, end_eid) = lookup_mtable_eid((
+                                &eentry.eid,
+                                mentry.ltype,
+                                mentry.offset,
+                                mentry.atype == AccessType::Write,
+                            ));
+
+                            MemoryRWEntry {
+                                entry: mentry,
+                                start_eid,
+                                end_eid,
+                            }
+                        })
+                        .collect(),
+                })
+                .collect(),
+        )
+    }
+}
diff --git a/crates/specs/src/imtable.rs b/crates/specs/src/imtable.rs
@@ -1,10 +1,11 @@
 use std::cmp::Ordering;
+use std::collections::BTreeMap;
 
 use crate::mtable::LocationType;
 use crate::mtable::VarType;
 use serde::Serialize;
 
-#[derive(Serialize, Debug, Clone)]
+#[derive(Serialize, Debug, Clone, PartialEq)]
 pub struct InitMemoryTableEntry {
     pub ltype: LocationType,
     pub is_mutable: bool,
@@ -13,35 +14,27 @@ pub struct InitMemoryTableEntry {
     pub vtype: VarType,
     /// convert from [u8; 8] via u64::from_le_bytes
     pub value: u64,
+
+    // for stack
+    pub start_eid: u32,
+    pub end_eid: u32,
 }
 
 #[derive(Serialize, Default, Debug, Clone)]
 pub struct InitMemoryTable {
     entries: Vec<InitMemoryTableEntry>,
-    sorted_global_init_entries: Vec<InitMemoryTableEntry>,
+    sorted_global_init_entries: BTreeMap<u32, InitMemoryTableEntry>,
     sorted_heap_init_entries: Vec<InitMemoryTableEntry>,
+    sorted_stack_init_entries: BTreeMap<u32, InitMemoryTableEntry>,
 }
 
 impl InitMemoryTable {
-    pub fn new(entries: Vec<InitMemoryTableEntry>, k: u32) -> Self {
+    pub fn new(entries: Vec<InitMemoryTableEntry>) -> Self {
         let mut imtable = Self {
-            entries: entries
-                .into_iter()
-                .map(|entry| InitMemoryTableEntry {
-                    ltype: entry.ltype,
-                    is_mutable: entry.is_mutable,
-                    start_offset: entry.start_offset,
-                    end_offset: if entry.end_offset == u32::MAX {
-                        (1u32 << (k - 1)) - 1
-                    } else {
-                        entry.end_offset
-                    },
-                    vtype: entry.vtype,
-                    value: entry.value,
-                })
-                .collect(),
-            sorted_global_init_entries: vec![],
+            entries,
+            sorted_global_init_entries: BTreeMap::new(),
             sorted_heap_init_entries: vec![],
+            sorted_stack_init_entries: BTreeMap::new(),
         };
         imtable.sort();
         imtable.merge();
@@ -56,7 +49,13 @@ impl InitMemoryTable {
             .entries
             .iter()
             .filter(|entry| entry.ltype == LocationType::Global)
-            .map(|entry| entry.clone())
+            .map(|entry| (entry.start_offset, entry.clone()))
+            .collect();
+        imtable.sorted_stack_init_entries = imtable
+            .entries
+            .iter()
+            .filter(|entry| entry.ltype == LocationType::Stack)
+            .map(|entry| (entry.start_offset, entry.clone()))
             .collect();
 
         imtable
@@ -93,16 +92,18 @@ impl InitMemoryTable {
                 ));
             }
             LocationType::Global => {
-                for entry in self.sorted_global_init_entries.iter() {
-                    if entry.start_offset == offset {
-                        return Some((offset, offset, entry.value));
-                    }
-                }
+                return self
+                    .sorted_global_init_entries
+                    .get(&offset)
+                    .map(|entry| (offset, offset, entry.value))
+            }
+            LocationType::Stack => {
+                return self
+                    .sorted_stack_init_entries
+                    .get(&offset)
+                    .map(|entry| (offset, offset, entry.value))
             }
-            LocationType::Stack => unreachable!(),
         }
-
-        None
     }
 
     fn sort(&mut self) {
@@ -114,7 +115,7 @@ impl InitMemoryTable {
         let mut merged_entries: Vec<_> = self
             .entries()
             .iter()
-            .filter(|entry| entry.ltype == LocationType::Global)
+            .filter(|entry| entry.ltype != LocationType::Heap)
             .map(|entry| entry.clone())
             .collect();
 
@@ -128,7 +129,10 @@ impl InitMemoryTable {
             let mut scan = 0;
             let mut cursor = scan + 1;
             while scan < heap_initial.len() && cursor < heap_initial.len() {
-                if heap_initial[scan].value == heap_initial[cursor].value {
+                if heap_initial[scan].value == heap_initial[cursor].value
+                    && heap_initial[scan].start_eid == heap_initial[cursor].start_eid
+                    && heap_initial[scan].end_eid == heap_initial[cursor].end_eid
+                {
                     cursor += 1;
                 } else {
                     merged_entries.push(InitMemoryTableEntry {
@@ -138,6 +142,9 @@ impl InitMemoryTable {
                         end_offset: heap_initial[cursor - 1].end_offset,
                         vtype: VarType::I64,
                         value: heap_initial[scan].value,
+
+                        start_eid: heap_initial[scan].start_eid,
+                        end_eid: heap_initial[scan].end_eid,
                     });
 
                     scan = cursor;
@@ -151,6 +158,9 @@ impl InitMemoryTable {
                 end_offset: heap_initial[cursor - 1].end_offset,
                 vtype: VarType::I64,
                 value: heap_initial[scan].value,
+
+                start_eid: heap_initial[scan].start_eid,
+                end_eid: heap_initial[scan].end_eid,
             });
         }
 

diff --git a/crates/specs/src/lib.rs b/crates/specs/src/lib.rs
@@ -16,7 +16,6 @@ use itable::InstructionTable;
 use jtable::JumpTable;
 use jtable::StaticFrameEntry;
 use mtable::AccessType;
-use mtable::LocationType;
 use mtable::MTable;
 use mtable::MemoryTableEntry;
 use rayon::prelude::IntoParallelRefIterator;
@@ -155,14 +154,13 @@ pub struct Tables {
     pub is_last_slice: bool,
 }
 
-impl Tables {
+impl EventTable {
     pub fn create_memory_table(
         &self,
+        imtable: &InitMemoryTable,
         memory_event_of_step: fn(&EventTableEntry, &mut u32) -> Vec<MemoryTableEntry>,
     ) -> MTable {
         let mut memory_entries = self
-            .execution_tables
-            .etable
             .entries()
             .par_iter()
             .map(|entry| memory_event_of_step(entry, &mut 1))
@@ -172,17 +170,9 @@ impl Tables {
         let init_value = memory_entries
             .par_iter()
             .map(|entry| {
-                if entry.ltype == LocationType::Heap || entry.ltype == LocationType::Global {
-                    let (_, _, value) = self
-                        .compilation_tables
-                        .imtable
-                        .try_find(entry.ltype, entry.offset)
-                        .unwrap();
-
-                    Some(value)
-                } else {
-                    None
-                }
+                imtable
+                    .try_find(entry.ltype, entry.offset)
+                    .map(|(_, _, value)| value)
             })
             .collect::<Vec<_>>();
 
@@ -213,13 +203,13 @@ impl Tables {
 
         MTable::new(memory_entries)
     }
+}
 
+impl Tables {
     pub fn is_last_slice(&self) -> bool {
         self.is_last_slice
     }
-}
 
-impl Tables {
     pub fn write_json(&self, dir: Option<PathBuf>) {
         fn write_file(folder: &PathBuf, filename: &str, buf: &String) {
             let mut folder = folder.clone();