fix allocators recycling

Author: robik75

gpu_prover/src/execution/cpu_worker.rs

@@ -54,7 +54,7 @@ pub trait NonDeterminism: NonDeterminismCSRSource<Ram> + Clone {}
 
 impl<T> NonDeterminism for T where T: NonDeterminismCSRSource<Ram> + Clone {}
 
-const SNAPSHOT_PERIOD: usize = 1 << 20;
+const SNAPSHOT_PERIOD: usize = 1 << 22;
 
 pub(crate) fn run_split_simulator(
     batch_id: u64,

gpu_prover/src/execution/prover.rs

@@ -85,7 +85,7 @@ impl Default for ExecutionProverConfiguration {
             prover_context_config: Default::default(),
             host_allocator_backing_allocation_size: 1 << 26, // 64 MB
             replay_worker_threads_count: 4,
-            host_allocators_per_worker_count: 16, // 1 GB
+            host_allocators_per_worker_count: 64, // 4 GB
             host_allocators_per_device_count: 64, // 4 GB
         }
     }
@@ -450,12 +450,12 @@ impl<K: Clone + Debug + Eq + Hash> ExecutionProver<K> {
                         } = commitment;
                         trace!("BATCH[{batch_id}] PROVER received memory commitment for circuit {circuit_type:?}[{sequence_id}]");
                         if let Some(inits_and_teardowns) = inits_and_teardowns {
-                            if let Some(allocator) = inits_and_teardowns.get_allocator() {
+                            for allocator in inits_and_teardowns.into_allocators() {
                                 self.free_allocators_sender.send(allocator).unwrap();
                             }
                         }
                         if let Some(tracing_data) = tracing_data {
-                            if let Some(allocator) = tracing_data.get_allocator() {
+                            for allocator in tracing_data.into_allocators() {
                                 self.free_allocators_sender.send(allocator).unwrap();
                             }
                         }
@@ -476,12 +476,12 @@ impl<K: Clone + Debug + Eq + Hash> ExecutionProver<K> {
                         } = proof;
                         trace!("BATCH[{batch_id}] PROVER received proof for circuit {circuit_type:?}[{sequence_id}]");
                         if let Some(inits_and_teardowns) = inits_and_teardowns {
-                            if let Some(allocator) = inits_and_teardowns.get_allocator() {
+                            for allocator in inits_and_teardowns.into_allocators() {
                                 self.free_allocators_sender.send(allocator).unwrap();
                             }
                         }
                         if let Some(tracing_data) = tracing_data {
-                            if let Some(allocator) = tracing_data.get_allocator() {
+                            for allocator in tracing_data.into_allocators() {
                                 self.free_allocators_sender.send(allocator).unwrap();
                             }
                         }
@@ -1406,9 +1406,9 @@ mod tests {
             &binary_image,
             &text_section,
         );
-        let non_determinism_source = QuasiUARTSource::new_with_reads(vec![1 << 20, 1 << 16]);
+        let non_determinism_source = QuasiUARTSource::new_with_reads(vec![1 << 26, 0]);
         let (_, result) =
-            prover.get_results(false, 0, &0usize, 1 << 30, non_determinism_source, None);
+            prover.get_results(false, 0, &0usize, 1 << 36, non_determinism_source, None);
         // match result {
         //     ExecutionProverResult::MemoryCommitment(memory_commitment) => {
         //         for (circuit_type, commitment) in memory_commitment

gpu_prover/src/prover/tracing_data.rs

@@ -39,11 +39,11 @@ pub(crate) enum DelegationTracingDataHost<A: GoodAllocator> {
 }
 
 impl<A: GoodAllocator> DelegationTracingDataHost<A> {
-    pub fn get_allocator(self) -> Option<A> {
+    pub fn into_allocators(self) -> Vec<A> {
         match self {
-            DelegationTracingDataHost::BigIntWithControl(trace) => trace.get_allocator(),
-            DelegationTracingDataHost::Blake2WithCompression(trace) => trace.get_allocator(),
-            DelegationTracingDataHost::KeccakSpecial5(trace) => trace.get_allocator(),
+            DelegationTracingDataHost::BigIntWithControl(trace) => trace.into_allocators(),
+            DelegationTracingDataHost::Blake2WithCompression(trace) => trace.into_allocators(),
+            DelegationTracingDataHost::KeccakSpecial5(trace) => trace.into_allocators(),
         }
     }
 }
@@ -82,11 +82,11 @@ pub(crate) enum UnrolledTracingDataHost<A: GoodAllocator> {
 }
 
 impl<A: GoodAllocator> UnrolledTracingDataHost<A> {
-    pub fn get_allocator(self) -> Option<A> {
+    pub fn into_allocators(self) -> Vec<A> {
         match self {
-            UnrolledTracingDataHost::Memory(trace) => trace.get_allocator(),
-            UnrolledTracingDataHost::NonMemory(trace) => trace.get_allocator(),
-            UnrolledTracingDataHost::Unified(trace) => trace.get_allocator(),
+            UnrolledTracingDataHost::Memory(trace) => trace.into_allocators(),
+            UnrolledTracingDataHost::NonMemory(trace) => trace.into_allocators(),
+            UnrolledTracingDataHost::Unified(trace) => trace.into_allocators(),
         }
     }
 }
@@ -98,10 +98,10 @@ pub(crate) enum TracingDataHost<A: GoodAllocator> {
 }
 
 impl<A: GoodAllocator> TracingDataHost<A> {
-    pub fn get_allocator(self) -> Option<A> {
+    pub fn into_allocators(self) -> Vec<A> {
         match self {
-            TracingDataHost::Delegation(trace) => trace.get_allocator(),
-            TracingDataHost::Unrolled(trace) => trace.get_allocator(),
+            TracingDataHost::Delegation(trace) => trace.into_allocators(),
+            TracingDataHost::Unrolled(trace) => trace.into_allocators(),
         }
     }
 }

gpu_prover/src/witness/trace.rs

@@ -31,12 +31,10 @@ impl<T, A: GoodAllocator> ChunkedTraceHolder<T, A> {
         panic!("Index out of bounds");
     }
 
-    pub fn get_allocator(self) -> Option<A> {
-        let chunks = self
-            .chunks
+    pub fn into_allocators(self) -> Vec<A> {
+        self.chunks
             .into_iter()
-            .map(|c| Arc::into_inner(c).unwrap())
-            .collect_vec();
-        chunks.get(0).map(|c| c.allocator().clone())
+            .map(|c| Arc::into_inner(c).unwrap().allocator().clone())
+            .collect()
     }
 }