Reduce peak memory during prove by releasing witness shared memory early (#204)

* Add unchecked device memory export, LogUp query count API, lightweight prove, and configurable CPU count

- Add `export_device_memories_unchecked()` for exporting device memories without
  state assertion, enabling memory optimization workflows where context is dropped
  before proving
- Add `prove_lightweight()` to ExpanderNoOverSubscribe, allowing prove without
  holding computation_graph or prover_setup references
- Add `final_check_with_query_count()` to LogUpSingleKeyTable and
  LogUpRangeProofTable for hint-free logup verification with externally provided
  query counts
- Support `ZKML_NUM_CPUS` env var to override physical CPU detection for MPI
  process count

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Reduce peak memory during prove by releasing witness shared memory early

Add witness_ack shared memory signaling between client and server:
- Client resets a 1-byte ack signal before writing witness
- Server signals ack after reading witness into MPI shared memory
- Client polls for ack, then immediately releases witness shared memory
  and calls malloc_trim to return memory to OS
- Prove request runs concurrently via tokio async, so witness memory is
  freed while proving is in progress
- Skip reading PCS setup from shared memory (return default) since the
  client does not need it after setup

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Address review feedback: platform-guard malloc_trim, reduce polling interval

- Wrap malloc_trim calls with #[cfg(all(target_os = "linux", target_env = "gnu"))]
  to avoid linker errors on non-glibc platforms
- Reduce witness_ack polling interval from 500ms to 10ms for faster response

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Address review: add polling timeout, revert debug label, remove redundant malloc_trim

- Add 5-minute timeout to wait_for_witness_read_complete to prevent
  indefinite hang if the server crashes
- Revert timer label from "new setup" back to "setup"
- Remove duplicate malloc_trim inside spawn_blocking (shared memory
  is mmap-managed, not glibc heap)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: restore verifier setup from shared memory to fix verify panic

The previous optimization skipped reading PCS setup from shared memory
and returned empty defaults, which caused verify to panic on v_keys
lookup (unwrap on None).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: hczphn

expander_compiler/src/zkcuda/proving_system/expander_parallelized/client_utils.rs

@@ -140,7 +140,11 @@ where
 
     setup_timer.stop();
 
-    SharedMemoryEngine::read_pcs_setup_from_shared_memory()
+    // Prover setup not needed on client side (server does the proving).
+    // Verifier setup is required for verification, so read it from shared memory.
+    let (_prover_setup, verifier_setup) =
+        SharedMemoryEngine::read_pcs_setup_from_shared_memory::<C::FieldConfig, C::PCSConfig>();
+    (ExpanderProverSetup::default(), verifier_setup)
 }
 
 pub fn client_send_witness_and_prove<C, ECCConfig>(
@@ -152,8 +156,39 @@ where
 {
     let timer = Timer::new("prove", true);
 
+    // Reset ack signal, then write witness
+    SharedMemoryEngine::reset_witness_ack();
     SharedMemoryEngine::write_witness_to_shared_memory::<C::FieldConfig>(device_memories);
-    wait_async(ClientHttpHelper::request_prove());
+
+    #[cfg(all(target_os = "linux", target_env = "gnu"))]
+    {
+        extern "C" {
+            fn malloc_trim(pad: usize) -> i32;
+        }
+        unsafe {
+            malloc_trim(0);
+        }
+    }
+
+    // Async: send prove request + poll for witness ack to release shared memory early
+    let rt = tokio::runtime::Runtime::new().unwrap();
+    rt.block_on(async {
+        let prove_handle = tokio::spawn(async {
+            ClientHttpHelper::request_prove().await;
+        });
+
+        // Poll witness_ack; once server confirms read, release witness shared memory
+        tokio::task::spawn_blocking(|| {
+            SharedMemoryEngine::wait_for_witness_read_complete();
+            unsafe {
+                super::shared_memory_utils::SHARED_MEMORY.witness = None;
+            }
+        })
+        .await
+        .expect("Witness cleanup task failed");
+
+        prove_handle.await.expect("Prove task failed");
+    });
 
     let proof = SharedMemoryEngine::read_proof_from_shared_memory();
 

expander_compiler/src/zkcuda/proving_system/expander_parallelized/server_ctrl.rs

@@ -149,6 +149,9 @@ where
             let mut witness_win = state.wt_shared_memory_win.lock().await;
             S::setup_shared_witness(&state.global_mpi_config, &mut witness, &mut witness_win);
 
+            // Signal client: witness has been read, shared memory can be released
+            SharedMemoryEngine::signal_witness_read_complete();
+
             let prover_setup_guard = state.prover_setup.lock().await;
             let computation_graph = state.computation_graph.lock().await;
 

expander_compiler/src/zkcuda/proving_system/expander_parallelized/shared_memory_utils.rs

@@ -18,12 +18,15 @@ pub struct SharedMemory {
     pub pcs_setup: Option<Shmem>,
     pub witness: Option<Shmem>,
     pub proof: Option<Shmem>,
+    /// 1-byte signal: 0 = witness not read, 1 = server finished reading witness
+    pub witness_ack: Option<Shmem>,
 }
 
 pub static mut SHARED_MEMORY: SharedMemory = SharedMemory {
     pcs_setup: None,
     witness: None,
     proof: None,
+    witness_ack: None,
 };
 
 pub struct SharedMemoryEngine {}
@@ -106,6 +109,56 @@ impl SharedMemoryEngine {
         Self::read_object_from_shared_memory("pcs_setup", 0)
     }
 
+    /// Client: reset witness_ack to 0 (call before writing witness)
+    pub fn reset_witness_ack() {
+        unsafe {
+            Self::allocate_shared_memory_if_necessary(
+                &mut SHARED_MEMORY.witness_ack,
+                "witness_ack",
+                1,
+            );
+            let ptr = SHARED_MEMORY.witness_ack.as_mut().unwrap().as_ptr();
+            std::ptr::write_volatile(ptr, 0u8);
+        }
+    }
+
+    /// Server: set witness_ack to 1 (call after reading witness)
+    pub fn signal_witness_read_complete() {
+        let shmem = ShmemConf::new()
+            .flink("witness_ack")
+            .open()
+            .expect("Failed to open witness_ack shared memory");
+        unsafe {
+            std::ptr::write_volatile(shmem.as_ptr(), 1u8);
+        }
+    }
+
+    /// Client: poll until witness_ack becomes 1, with a timeout to avoid hanging
+    /// if the server crashes.
+    pub fn wait_for_witness_read_complete() {
+        const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
+        let start = std::time::Instant::now();
+        unsafe {
+            let ptr = SHARED_MEMORY
+                .witness_ack
+                .as_ref()
+                .expect("witness_ack not initialized, call reset_witness_ack first")
+                .as_ptr() as *const u8;
+            loop {
+                if std::ptr::read_volatile(ptr) != 0 {
+                    break;
+                }
+                if start.elapsed() > TIMEOUT {
+                    panic!(
+                        "Timed out waiting for server to read witness ({}s)",
+                        TIMEOUT.as_secs()
+                    );
+                }
+                std::thread::sleep(std::time::Duration::from_millis(10));
+            }
+        }
+    }
+
     pub fn write_witness_to_shared_memory<F: FieldEngine>(values: Vec<Vec<F::SimdCircuitField>>) {
         let total_size = std::mem::size_of::<usize>()
             + values