fix: release lock when compiling recursion program

succinctlabs · Feb 26, 2025 · d65cb0f · d65cb0f
1 parent 5a957a5
commit d65cb0f
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 108 deletions.
diff --git a/crates/core/machine/src/cpu/trace.rs b/crates/core/machine/src/cpu/trace.rs
@@ -109,6 +109,10 @@ impl<F: PrimeField32> MachineAir<F> for CpuChip {
             shard.contains_cpu()
         }
     }
+
+    fn num_rows(&self, shard: &Self::Record) -> Option<usize> {
+        Some(shard.cpu_events.len())
+    }
 }
 
 impl CpuChip {

diff --git a/crates/prover/src/lib.rs b/crates/prover/src/lib.rs
@@ -941,9 +941,14 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         &self,
         input: &SP1RecursionWitnessValues<CoreSC>,
     ) -> Arc<RecursionProgram<BabyBear>> {
+        // Check if the program is in the cache.
         let mut cache = self.lift_programs_lru.lock().unwrap_or_else(|e| e.into_inner());
-        cache
-            .get_or_insert(input.shape(), || {
+        let shape = input.shape();
+        let program = cache.get(&shape).cloned();
+        drop(cache);
+        match program {
+            Some(program) => program,
+            None => {
                 let misses = self.lift_cache_misses.fetch_add(1, Ordering::Relaxed);
                 tracing::debug!("core cache miss, misses: {}", misses);
                 // Get the operations.
@@ -971,9 +976,14 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
                 }
                 let program = Arc::new(program);
                 compiler_span.exit();
+
+                // Insert the program into the cache.
+                let mut cache = self.lift_programs_lru.lock().unwrap_or_else(|e| e.into_inner());
+                cache.put(shape, program.clone());
+                drop(cache);
                 program
-            })
-            .clone()
+            }
+        }
     }
 
     pub fn compress_program(