perf: use vec_map and paging for executor memory (#1461)

crates/core/executor/Cargo.toml

@@ -43,6 +43,7 @@ log = "0.4.22"
 hex = "0.4.3"
 bytemuck = "1.16.3"
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
+vec_map = { version = "0.8.2", features = ["serde"] }
 
 [dev-dependencies]
 sp1-zkvm = { workspace = true }

crates/core/executor/src/executor.rs

@@ -4,8 +4,7 @@ use std::{
     sync::Arc,
 };
 
-use hashbrown::{hash_map::Entry, HashMap};
-use nohash_hasher::BuildNoHashHasher;
+use hashbrown::HashMap;
 use serde::{Deserialize, Serialize};
 use sp1_stark::SP1CoreOpts;
 use thiserror::Error;
@@ -18,6 +17,7 @@ use crate::{
         MemoryWriteRecord,
     },
     hook::{HookEnv, HookRegistry},
+    memory::{Entry, PagedMemory},
     record::{ExecutionRecord, MemoryAccessRecord},
     report::ExecutionReport,
     state::{ExecutionState, ForkState},
@@ -100,7 +100,7 @@ pub struct Executor<'a> {
 
     /// Memory addresses that were touched in this batch of shards. Used to minimize the size of
     /// checkpoints.
-    pub memory_checkpoint: HashMap<u32, Option<MemoryRecord>, BuildNoHashHasher<u32>>,
+    pub memory_checkpoint: PagedMemory<Option<MemoryRecord>>,
 }
 
 /// The different modes the executor can run in.
@@ -216,7 +216,7 @@ impl<'a> Executor<'a> {
             hook_registry,
             opts,
             max_cycles: context.max_cycles,
-            memory_checkpoint: HashMap::default(),
+            memory_checkpoint: PagedMemory::new_preallocated(),
         }
     }
 
@@ -255,7 +255,7 @@ impl<'a> Executor<'a> {
         let mut registers = [0; 32];
         for i in 0..32 {
             let addr = Register::from_u32(i as u32) as u32;
-            let record = self.state.memory.get(&addr);
+            let record = self.state.memory.get(addr);
 
             // Only add the previous memory state to checkpoint map if we're in checkpoint mode,
             // or if we're in unconstrained mode. In unconstrained mode, the mode is always
@@ -283,7 +283,7 @@ impl<'a> Executor<'a> {
     #[must_use]
     pub fn register(&mut self, register: Register) -> u32 {
         let addr = register as u32;
-        let record = self.state.memory.get(&addr);
+        let record = self.state.memory.get(addr);
 
         if self.executor_mode == ExecutorMode::Checkpoint || self.unconstrained {
             match record {
@@ -306,7 +306,7 @@ impl<'a> Executor<'a> {
     #[must_use]
     pub fn word(&mut self, addr: u32) -> u32 {
         #[allow(clippy::single_match_else)]
-        let record = self.state.memory.get(&addr);
+        let record = self.state.memory.get(addr);
 
         if self.executor_mode == ExecutorMode::Checkpoint || self.unconstrained {
             match record {
@@ -1174,7 +1174,6 @@ impl<'a> Executor<'a> {
     pub fn execute_state(&mut self) -> Result<(ExecutionState, bool), ExecutionError> {
         self.memory_checkpoint.clear();
         self.executor_mode = ExecutorMode::Checkpoint;
-        self.print_report = true;
 
         // Take memory out of state before cloning it so that memory is not cloned.
         let memory = std::mem::take(&mut self.state.memory);
@@ -1194,7 +1193,7 @@ impl<'a> Executor<'a> {
                     if let Some(record) = record {
                         checkpoint.memory.insert(addr, record);
                     } else {
-                        checkpoint.memory.remove(&addr);
+                        checkpoint.memory.remove(addr);
                     }
                 });
             } else {
@@ -1212,8 +1211,8 @@ impl<'a> Executor<'a> {
         self.state.channel = 0;
 
         tracing::debug!("loading memory image");
-        for (addr, value) in &self.program.memory_image {
-            self.state.memory.insert(*addr, MemoryRecord { value: *value, shard: 0, timestamp: 0 });
+        for (&addr, value) in &self.program.memory_image {
+            self.state.memory.insert(addr, MemoryRecord { value: *value, shard: 0, timestamp: 0 });
         }
     }
 
@@ -1351,7 +1350,7 @@ impl<'a> Executor<'a> {
 
         // We handle the addr = 0 case separately, as we constrain it to be 0 in the first row
         // of the memory finalize table so it must be first in the array of events.
-        let addr_0_record = self.state.memory.get(&0u32);
+        let addr_0_record = self.state.memory.get(0);
 
         let addr_0_final_record = match addr_0_record {
             Some(record) => record,
@@ -1365,27 +1364,30 @@ impl<'a> Executor<'a> {
             MemoryInitializeFinalizeEvent::initialize(0, 0, addr_0_record.is_some());
         memory_initialize_events.push(addr_0_initialize_event);
 
-        self.report.touched_memory_addresses = self.state.memory.len() as u64;
+        // Count the number of touched memory addresses manually, since `PagedMemory` doesn't
+        // already know its length.
+        self.report.touched_memory_addresses = 0;
         for addr in self.state.memory.keys() {
-            if *addr == 0 {
+            self.report.touched_memory_addresses += 1;
+            if addr == 0 {
                 // Handled above.
                 continue;
             }
 
             // Program memory is initialized in the MemoryProgram chip and doesn't require any
             // events, so we only send init events for other memory addresses.
-            if !self.record.program.memory_image.contains_key(addr) {
-                let initial_value = self.state.uninitialized_memory.get(addr).unwrap_or(&0);
+            if !self.record.program.memory_image.contains_key(&addr) {
+                let initial_value = self.state.uninitialized_memory.get(&addr).unwrap_or(&0);
                 memory_initialize_events.push(MemoryInitializeFinalizeEvent::initialize(
-                    *addr,
+                    addr,
                     *initial_value,
                     true,
                 ));
             }
 
             let record = *self.state.memory.get(addr).unwrap();
             memory_finalize_events
-                .push(MemoryInitializeFinalizeEvent::finalize_from_record(*addr, &record));
+                .push(MemoryInitializeFinalizeEvent::finalize_from_record(addr, &record));
         }
     }
 

crates/core/executor/src/lib.rs

@@ -26,6 +26,7 @@ mod executor;
 mod hook;
 mod instruction;
 mod io;
+mod memory;
 mod opcode;
 mod program;
 #[cfg(any(test, feature = "programs"))]