The PCE Loading Transition: From Static Globals to On-Demand Caching in cuzk's Memory Manager

In the middle of a sweeping architectural transformation of the cuzk GPU proving engine's memory management subsystem, a single message from the assistant captures a critical design inflection point. The message reads:

Now update load_pce_from_disk and preload_pce_from_disk functions. They need to be adapted to work with PceCache. Let me update load_pce_from_disk to become a standalone function that returns the PCE (for backward compat with bench tool), and modify preload_pce_from_disk to be removed since loading is now on-demand: [edit] /tmp/czk/extern/cuzk/cuzk-core/src/pipeline.rs Edit applied successfully.

This seemingly small edit — adapting two functions in a single file — is in fact a microcosm of the entire memory manager redesign. It represents the transition from a rigid, process-lifetime caching model to a flexible, budget-aware, on-demand loading architecture. Understanding why this message was written, what decisions it embodies, and what assumptions underpin it requires tracing the full arc of the memory manager implementation.

The Context: A Memory Crisis in the Making

The cuzk daemon is a GPU-accelerated zero-knowledge proving engine used in the Filecoin network. It generates proofs for various circuit types — PoRep (Proof-of-Replication), WindowPoSt, WinningPoSt, and SnapDeals — each of which requires large amounts of memory. The original architecture stored Pre-Compiled Constraint Evaluator (PCE) data in four static OnceLock<PreCompiledCircuit<Fr>> globals, one per circuit type. These were loaded once at daemon startup (or on first use) and never evicted. On a machine running multiple proof types, the PCE data alone could consume over 70 GiB of heap memory, and with SRS (Structured Reference String) pinned memory added on top, baseline RSS could exceed 200 GiB.

The problem was not just the absolute memory consumption, but the lack of any coordination between consumers. The partition_workers config limited concurrent partition synthesis by a static count, not by actual memory pressure. The working_memory_budget and pinned_budget configs existed but were dead code — never checked at runtime. The system had no way to say "we're out of memory, evict something." PCE data, once loaded into a OnceLock, was immortal.

The design document at cuzk-memory-manager.md (written in the preceding sub-session, see [msg 2080]) proposed a radical solution: a single unified memory budget that covers SRS, PCE, and synthesis working set under one byte-level cap, with on-demand loading and LRU eviction. The assistant's task in this chunk was to implement that design.

Why This Message Matters

By the time the assistant reached message [msg 2109], it had already completed several major steps. It had created the new memory.rs module with MemoryBudget, MemoryReservation, and estimation constants ([msg 2093]). It had updated config.rs to replace the old dead-code fields with the new unified budget configuration ([msg 2097]-[msg 2103]). It had rewritten srs_manager.rs to be budget-aware with last_used tracking and eviction support ([msg 2105]-[msg 2106]). And in the immediately preceding message ([msg 2108]), it had replaced the four static OnceLock PCE caches in pipeline.rs with the new PceCache struct.

But replacing the statics was only half the work. The old code had two functions that directly interacted with those statics: load_pce_from_disk and preload_pce_from_disk. These functions embodied the old paradigm — preload everything at startup, cache forever, never evict. To complete the transition, they needed to be adapted or removed.

The Reasoning Behind the Two Changes

The assistant's reasoning, visible in the message text, shows careful consideration of backward compatibility and architectural purity. Let us examine each decision.

`load_pce_from_disk`: Preserving Backward Compatibility

The assistant writes: "Let me update load_pce_from_disk to become a standalone function that returns the PCE (for backward compat with bench tool)." This is a pragmatic decision. The bench tool (cuzk-bench) is a separate binary used for performance testing and benchmarking. It likely calls load_pce_from_disk directly to load PCE data without going through the full engine startup sequence. If the function were simply removed or changed to require a PceCache reference, the bench tool would break.

The solution is elegant: make load_pce_from_disk a pure function that reads a PCE file from disk and returns the deserialized PreCompiledCircuit<Fr>, without inserting it into any cache. The caller — whether the bench tool or the new PceCache::load_from_disk method — is responsible for managing the returned value. This preserves the bench tool's workflow while decoupling the function from the old static globals.

This decision reveals an important assumption: that the bench tool's usage pattern is fundamentally different from the engine's. The bench tool wants to load PCE once, use it for a series of benchmark runs, and then exit. It doesn't need eviction or budget management. The engine, by contrast, needs to manage PCE as one component in a shared memory budget. Making load_pce_from_disk a standalone function cleanly separates these two use cases.

`preload_pce_from_disk`: Removing the Old Paradigm

The assistant writes: "modify preload_pce_from_disk to be removed since loading is now on-demand." This is the more radical change. The preload_pce_from_disk function was called during daemon startup to load all known PCE files into the static caches before any proof requests arrived. It returned a count of successfully loaded circuit types, which was used for logging and health checks.

In the new design, preloading is antithetical to the memory budget philosophy. The budget is a finite resource shared between SRS, PCE, and synthesis working memory. Loading PCE for all circuit types at startup would consume budget that might never be needed (if certain proof types are never requested) and would leave less room for synthesis working memory. The on-demand approach — load PCE only when a proof of that circuit type is first encountered — is strictly more memory-efficient.

The removal of preload_pce_from_disk also eliminates the startup-time latency spike. Preloading PCE for PoRep (26 GiB) and WindowPoSt (26 GiB) at startup could take tens of seconds and consume significant memory bandwidth. Under the new design, the first proof of each type pays the loading cost, but subsequent proofs benefit from the cached data — and if the cache is evicted under pressure, the cost is paid again on reload.

Assumptions and Potential Risks

The assistant's approach rests on several assumptions that deserve scrutiny.

Assumption 1: On-demand loading will not cause unacceptable latency for the first proof. If a PoRep proof arrives and the PCE is not yet loaded, the proof must either wait for PCE extraction (a ~50-second cold-start process that builds the circuit from scratch using RecordingCS) or load from disk (a ~5-second operation if the PCE file exists). The design document addresses this by having the first proof use the slow synthesis path while a background thread performs the extraction. But if the extraction fails or is delayed, the first proof could be significantly slower. The assistant implicitly assumes this trade-off is acceptable.

Assumption 2: The bench tool does not need budget management. This is likely correct — benchmarks typically run in isolation with known memory requirements. But if the bench tool were ever extended to run concurrent benchmarks for multiple circuit types, the standalone load_pce_from_disk function would not provide any memory coordination, potentially leading to OOM.

Assumption 3: No other code depends on preload_pce_from_disk. The assistant checked for call sites via grep (visible in subsequent messages at [msg 2111]) and found only the one in pipeline.rs. The daemon's startup code in engine.rs called preload_pce_from_disk, and that call would need to be removed separately (which the assistant planned for later in the engine.rs changes). The assumption is that removing the function entirely, rather than keeping a stub, is safe.

Assumption 4: The PceCache struct will handle all the responsibilities that load_pce_from_disk and preload_pce_from_disk previously covered. This includes loading from disk, inserting into the cache, acquiring budget, and supporting eviction. The PceCache::load_from_disk method (specified in the design document) would internally call the now-standalone load_pce_from_disk and then acquire budget for the loaded PCE. This separation of concerns is architecturally clean but requires careful implementation to avoid budget leaks or double-counting.

The Input Knowledge Required

To understand this message, one must grasp several layers of context:

The old PCE caching model: Four static OnceLock<PreCompiledCircuit<Fr>> globals, one per circuit type, loaded via load_pce_from_disk (single load) or preload_pce_from_disk (all-at-once), never evicted.
The new PCE caching model: A single PceCache struct with a HashMap<CircuitId, PceCacheEntry>, where each entry tracks last_used and size_bytes, supports eviction when the budget is under pressure, and acquires budget on insert.
The memory budget architecture: MemoryBudget with total_bytes and used_bytes, where reservations are RAII guards that release on drop, and eviction is triggered when acquire() cannot be fulfilled.
The bench tool's existence and its separate code path: The bench binary at cuzk-bench/src/main.rs has its own proving pipeline and references load_pce_from_disk directly, necessitating backward compatibility.
The pipeline.rs file structure: Understanding where the static OnceLocks were declared (around lines 293-305), where load_pce_from_disk was defined, and where preload_pce_from_disk was called from the engine startup code.

The Output Knowledge Created

This message produces a modified pipeline.rs where:

load_pce_from_disk is transformed from a function that inserts into a static OnceLock into a standalone function that returns Result<Option<PreCompiledCircuit<Fr>>>. It reads the PCE file from the parameter cache directory, deserializes it, and returns the value to the caller. The caller — whether PceCache::load_from_disk or the bench tool — decides what to do with it.
preload_pce_from_disk is removed entirely. The function that iterated over all known circuit types and called load_pce_from_disk for each, incrementing a counter, is deleted. The engine startup code that called it will need to be updated separately (in the pending engine.rs changes).
A comment is added noting that "PCE loading is now on-demand via PceCache," serving as documentation for future readers.

The Thinking Process Visible

The assistant's reasoning, as expressed in the message, shows a methodical approach to API migration. The assistant identifies two functions that need attention, evaluates each one's role in the old and new architectures, and applies a different strategy to each:

load_pce_from_disk: Keep the core functionality (reading and deserializing a PCE file) but change the interface to return the value instead of storing it in a global. This is a "refactor" strategy — preserve the useful computation, change the integration point.
preload_pce_from_disk: Delete entirely. This is a "remove and replace" strategy — the function's entire purpose (preloading all PCE at startup) is antithetical to the new design, so it has no place in the new codebase. The assistant also demonstrates awareness of downstream consumers ("backward compat with bench tool"), showing that the migration is not done in isolation but considers the broader codebase.

Conclusion

Message [msg 2109] is a small but pivotal step in the memory manager implementation. It completes the transition from static, process-lifetime PCE caching to a dynamic, budget-aware model by adapting the two loading functions that bridged the old and new worlds. The decision to keep load_pce_from_disk as a standalone function while removing preload_pce_from_disk entirely reflects a careful balance between backward compatibility and architectural progress. The assumptions underlying these changes — that on-demand loading is acceptable, that the bench tool needs a separate path, that no other code depends on the removed function — are reasonable but worth documenting, as they define the boundary conditions of the new design.

This message, like the broader chunk it belongs to, demonstrates a systematic approach to large-scale refactoring: create the new structures first, then migrate the consumers, then remove the old code. The PCE loading functions were the last vestiges of the old static caching model in pipeline.rs, and their adaptation marks the completion of that file's transformation.