Tracing the Memory Footprint: How SRS and PCE Shape the cuzk Memory Management Architecture

In the design of any high-performance GPU proving system, memory is the invisible governor—the resource that silently constrains throughput, concurrency, and reliability. In an opencode coding session focused on building a robust memory management architecture for the cuzk GPU proving engine, a pivotal exchange occurs when the user asks a seemingly simple question: "How does SRS preload and PCE fit into this?" The assistant's response, spanning a deep forensic analysis of two of the largest memory consumers in the system, transforms that question into a comprehensive architectural insight that fundamentally reshapes the memory management design.

The Context: A System Under Memory Pressure

To understand the significance of this message, we must first appreciate the context. The session had been working through a major architectural overhaul of cuzk's memory management. The assistant had just completed a detailed analysis of the entire memory lifecycle for 32 GiB PoRep proofs, identifying that the existing working_memory_budget configuration option was entirely dead code—never enforced anywhere in the system. The only real throttle was a static partition_workers semaphore, a fragile count-based mechanism that was memory-unaware and could not distinguish between different proof types with vastly different memory footprints (PoRep at ~13.6 GiB per partition vs. WinningPoSt at ~0.4 GiB).

The assistant had proposed a MemoryBudget admission control system: a budget-based gate that would replace the static semaphore with a memory-aware controller tracking estimated allocations in real-time, blocking new synthesis tasks when the working memory budget would be exceeded. The design included a two-phase release mechanism matching the actual lifecycle of proof synthesis (a/b/c freed immediately after GPU prove start, the remainder after proof finalization), and support for heterogeneous proof types.

But this proposal had a gap. The user, demonstrating sharp system-level thinking, asked about the two largest baseline memory consumers: the Structured Reference String (SRS) at ~44 GiB and the Pre-Compiled Constraint Evaluator (PCE) at ~26 GiB. How did they fit into the budget picture? Were they accounted for? Could they cause problems?

The Message: A Forensic Deep Dive

The assistant's response ([msg 2066]) is a masterclass in systematic codebase analysis. It opens by acknowledging the question's importance—"both SRS and PCE are significant memory consumers and have important interactions with any memory management scheme"—and then proceeds to trace every allocation and deallocation point for both systems across the entire codebase.

Tracing the SRS Lifecycle

The SRS analysis begins at the engine startup path (engine.rs:866-906), where Engine::start() calls SrsManager::ensure_loaded() for each circuit ID in the preload configuration. The assistant follows the chain through SrsManager::ensure_loaded (srs_manager.rs:148-208), which checks an in-memory HashMap and then calls SuprasealParameters::new(path), which in turn calls SRS::try_new (supraseal-c2/src/lib.rs:21-38)—a FFI call to the C++ create_SRS() function.

The real insight comes at the C++ level (groth16_srs.cuh:259-299). Here the assistant identifies the critical detail: the SRS is loaded via cudaHostAlloc()—CUDA pinned host memory. This is page-locked RAM registered with the CUDA driver for DMA, meaning it cannot be swapped out. It counts against both RSS and the CUDA pinned memory limit. The loading process involves mmap()ing the 44 GiB .params file from disk, computing the total size across all point arrays (L, A, B1, B2, H), allocating pinned memory, and spawning a reader thread to deserialize points from the mmap into the pinned buffer before munmap()ing the file.

The assistant identifies a critical transient spike: during the load, both the mmap (~44 GiB virtual, demand-paged) and the pinned allocation (~44 GiB) exist simultaneously, creating a peak transient of approximately 88 GiB during SRS load. This is a hidden OOM risk that the existing system does not account for.

Tracing the PCE Lifecycle

The PCE analysis covers two scenarios: warm start (PCE file exists on disk) and cold start (first-ever proof). For warm starts, the assistant traces through preload_pce_from_disk (pipeline.rs:351-372) and the binary deserialization at ~5 GB/s, noting that the on-disk file (~25 GiB) and the in-memory struct coexist briefly but the file buffer is freed as it's consumed since it uses read() not mmap().

For cold starts, the analysis is particularly valuable. The assistant identifies that the first proof uses the slow synthesize_with_hint path, while a background thread fires off extract_and_cache_pce_from_c1() (engine.rs:1394-1406). This extraction builds one partition circuit with RecordingCS, capturing R1CS structure into CSR matrices. The transient spike during cold extraction is approximately 40 GiB (13.6 GiB for synthesis + 26 GiB for the PCE being built). Crucially, the assistant notes that this runs on a background thread with no memory gate—it can overlap with partition synthesis work, potentially causing a dangerous memory spike.

The Critical Insight: Baseline vs. Working Memory

The most important conceptual contribution of this message is the distinction between baseline memory (SRS + PCE, held for process lifetime) and working memory (per-partition synthesis allocations, freed after GPU prove). The assistant calculates:

After startup (warm):
  SRS PoRep 32G ......... 44 GiB  (CUDA pinned)
  PCE PoRep 32G ......... 26 GiB  (heap)
  ────────────────────────────────
  Baseline .............. 70 GiB

Available for working set on a 256 GiB machine:
  256 - 70 - ~10 (OS/stack/misc) = ~176 GiB
  → fits ~12 partitions at 13.6 GiB each

This calculation reveals a fundamental design flaw in the original MemoryBudget proposal: the working_memory_budget cannot simply be "how much total RAM the machine has." It must be the amount after subtracting SRS and PCE. The assistant proposes two options: either the operator manually sets it (error-prone), or the system computes it automatically: working_budget = total_ram - srs_loaded_bytes - pce_memory_bytes() - safety_margin.

Four Concrete Problems Identified

The message identifies four specific issues that the memory management design must address:

  1. The budget must account for baseline: SRS and PCE consume ~70 GiB before any proof work begins. The working memory budget must be computed relative to this baseline, not total RAM.
  2. SRS loading during proof processing is dangerous: If a new proof type arrives (e.g., WinningPoSt while PoRep is running) and its SRS isn't preloaded, ensure_loaded is called inside the partition dispatch path while partition synthesis is already consuming memory. The SRS load's transient spike (mmap + cudaHostAlloc) can cause an OOM. The assistant recommends either pre-reserving space, gating SRS loading behind the budget, or requiring all expected SRS to be preloaded.
  3. PCE cold-start extraction is ungated: The background thread for PCE extraction has ~40 GiB transient cost and no coordination with the partition semaphore or any memory budget. On a first-ever start with pw=12, the first proof could consume 12 × 13.6 = 163 GiB working + 40 GiB PCE extraction + 44 GiB SRS = 247 GiB on a 256 GiB machine—dangerously close to OOM.
  4. PCE does not reduce per-partition working memory: While PCE speeds up synthesis (10-20s vs 50s), the memory footprint of each synthesis output remains ~13.6 GiB. PCE improves throughput, not memory efficiency.

The Recommendations: Evolving the Design

Based on these findings, the assistant proposes four concrete additions to the MemoryBudget design:

The Thinking Process: What Makes This Analysis Exceptional

What distinguishes this message is the assistant's systematic methodology. It does not simply answer the user's question with a high-level description. Instead, it traces the actual code paths—file by file, function by function—identifying every allocation, its lifetime, its memory type (pinned vs. heap), and its transient peaks. The analysis moves from Rust (engine.rs, srs_manager.rs, pipeline.rs) through FFI (supraseal-c2/src/lib.rs) into C++ (groth16_srs.cuh), following the allocation chain across language boundaries.

The assistant also demonstrates a sophisticated understanding of memory management concepts: the distinction between pinned and pageable memory, the implications of mmap vs. read for file I/O, the transient double-allocation during SRS load, and the difference between baseline and working memory. Each finding is traced back to specific line numbers and file paths, making the analysis actionable for implementation.

The message also shows careful reasoning about edge cases: what happens on cold start vs. warm start, what happens when multiple proof types arrive, what happens when SRS loading overlaps with partition synthesis. These are exactly the scenarios that cause real-world production incidents.

Conclusion

This message represents a critical turning point in the memory management architecture design. What began as a proposal for a MemoryBudget admission controller is transformed, through the user's probing question and the assistant's thorough analysis, into a more complete and robust design that accounts for the full memory lifecycle—including the 70 GiB baseline that exists before any proof work begins. The assistant's recommendations for auto-computing the available budget, gating PCE extraction, and gating SRS loading become essential elements of the final architecture specification.

The message is a testament to the power of systematic codebase analysis in systems design. By tracing every allocation path, identifying every transient spike, and reasoning about every edge case, the assistant produces not just an answer to the user's question, but a refined architectural vision that addresses the real complexity of memory management in a GPU proving engine.