The Last Five Call Sites: Completing the Memory Manager Integration in cuzk's GPU Proving Engine

In a single, deceptively short message at index 2174 of the conversation, the assistant performed what appears to be a routine find-and-replace operation: updating five dispatch_batch call sites in engine.rs to replace partition_workers, &partition_semaphore, with &budget, plus a pce_cache reference. The message reads in full:

Now let me update all 5 call sites at once. They all follow the same pattern — replace partition_workers, &partition_semaphore, with &budget, + pce_cache: [edit] /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs Edit applied successfully.

Beneath this terse surface lies the culmination of a deep architectural transformation spanning three conversation segments (14–16) and dozens of messages. The change is not cosmetic. It is the final mechanical step in replacing a fragile, static concurrency-limiting mechanism with a dynamic, memory-aware admission control system for the cuzk GPU proving engine — a system that generates zero-knowledge proofs for Filecoin storage proofs in production.

The Architecture Being Dismantled

To understand why this message exists, one must understand what it removes. The old engine.rs used a partition_workers configuration value — a static integer — to limit how many proof partitions could be processed concurrently. This was enforced through a partition_semaphore, a classic concurrency primitive that simply capped the number of in-flight partitions. The approach had a critical flaw: it knew nothing about actual memory consumption. A 32 GiB PoRep proof and a small SnapDeals proof consumed the same "slot," and the system had no mechanism to prevent memory exhaustion when multiple large proofs ran simultaneously. The semaphore was a crude throttle, not a resource manager.

Worse, the system preloaded all SRS (Structured Reference String) parameters and all PCE (Pre-Compiled Constraint Evaluator) circuits at startup, consuming tens of gigabytes of GPU-pinned memory before any proof work began. This "load everything and hope it fits" strategy worked only on machines with abundant RAM and failed unpredictably on memory-constrained hosts. The working_memory_budget field in the configuration existed but was dead code — never wired into any actual allocation decision.

The New Architecture: MemoryBudget and PceCache

The assistant had designed and implemented a unified memory manager across the preceding messages. The centerpiece is MemoryBudget ([msg 2159]), a byte-level budget auto-detected from system RAM that tracks all major memory consumers — SRS, PCE, and synthesis working sets — under a single atomic counter. Instead of a static semaphore, the budget provides acquire() and release() methods that reserve and free byte amounts. When the budget is exhausted, callers block until memory becomes available.

The PceCache ([msg 2162]) replaces the old static OnceLock-based PCE storage. It is eviction-aware: when the budget is under pressure, idle PCE entries can be dropped to free memory for new work. The SrsManager ([msg 2160]) was similarly rewritten to support budget-aware loading with ensure_loaded() taking an optional MemoryReservation, and an evictable_entries() method for the eviction callback.

These components had been implemented across memory.rs, srs_manager.rs, pipeline.rs, and config.rs in the preceding messages. What remained was wiring them into engine.rs — the central coordinator that owns the scheduler, GPU workers, and the public API for submitting proofs.

The engine.rs Transformation

The subject message is the seventh edit applied to engine.rs in this session, and it is the one that completes the dispatcher loop migration. The earlier edits had already:

Removed the SRS and PCE preload blocks ([msg 2165]), replacing them with evictor callback wiring that connects the budget to the SrsManager and PceCache after GPU detection.
Replaced channel capacity sizing ([msg 2166]), switching from partition_workers-derived values to budget-derived values that reflect actual available memory rather than an arbitrary concurrency count.
Removed the partition_semaphore setup ([msg 2167]), deleting the semaphore creation and its associated log message.
Updated the dispatcher spawn block ([msg 2168]) to pass budget and pce_cache into the dispatcher closure instead of the old partition_workers and partition_semaphore.
Updated the dispatch_batch function signature ([msg 2170]) to accept &MemoryBudget and &PceCache instead of usize (worker count) and &Semaphore. But the function signature change alone does nothing — the call sites must be updated to pass the new arguments. That is precisely what the subject message accomplishes.

The Five Call Sites

The assistant had identified the five call sites via a grep for dispatch_batch( ([msg 2172]), which returned matches at lines 1175, 1189, 1231, 1246, and 1263 of engine.rs. Each represented a different dispatch path in the dispatcher loop:

Line 1175: The force-flush path when the dispatcher receives a shutdown signal and needs to drain any remaining batch before exiting.
Line 1189: The normal batch dispatch for PoRep proofs in partition mode.
Line 1231: The batch dispatch for SnapDeals proofs.
Line 1246: Another dispatch path, likely for single-sector or fallback cases.
Line 1263: The monolithic (non-partitioned) dispatch path. All five followed the same parameter pattern: partition_workers, &partition_semaphore, followed by the batch-specific arguments. The assistant's edit replaced these two parameters with &budget, and the appropriate pce_cache reference, completing the migration.

Why This Message Matters

On its surface, this is a mechanical edit — five lines changed in a pattern that the assistant correctly identified as uniform. But the message represents the moment when the old architecture is fully disconnected and the new one is live. Before this edit, the dispatch_batch function had a new signature but was still being called with old-style arguments (which would cause a compilation error). After this edit, the dispatcher loop compiles against the new API, and the partition_semaphore — the last vestige of the static concurrency model — has no remaining references.

The message also demonstrates a key engineering judgment: the assistant recognized that all five call sites followed the same pattern and applied a single bulk edit rather than updating them individually. This was not blind optimism — the assistant had first read the surrounding context ([msg 2173]) to verify the pattern held, and confirmed via grep that no other dispatch_batch call sites existed. The decision to batch the edits reflects an understanding that mechanical transformations with verified uniformity are safer when applied atomically, as it avoids the risk of editing four sites, missing the fifth, and leaving the code in a broken state.

Assumptions and Risks

The assistant made several assumptions in this edit. It assumed that budget and pce_cache were in scope at all five call sites — a reasonable assumption since they had been passed into the dispatcher closure in the earlier edit ([msg 2168]). It assumed that the pce_cache variable name was consistent across all sites (it was, as they all lived within the same async move block). And it assumed that no other code paths called dispatch_batch — verified by the grep that returned exactly five non-definition matches.

The risk of a bulk find-and-replace on a pattern like partition_workers, &partition_semaphore, is that the same token sequence might appear in comments, log messages, or unrelated code. The assistant mitigated this by targeting only the specific edit context within engine.rs and by having read the surrounding lines to confirm the pattern's meaning.

The Broader Significance

This message is a case study in how large architectural changes are landed in complex systems. The memory manager integration touched six files across the cuzk codebase: memory.rs, config.rs, srs_manager.rs, pipeline.rs, engine.rs, and cuzk-bench/src/main.rs. Each file was updated in a dependency order — foundational types first, consumers last. engine.rs, being the central coordinator that depends on all the others, was naturally the final file to be updated.

The five dispatch_batch call sites were the last remaining references to the old API within the dispatcher loop. After this edit, the assistant would go on to update the PoRep and SnapDeals per-partition dispatch paths, the monolithic synthesis path, and the GPU worker loop — all in subsequent edits within the same segment. But the subject message marks the inflection point where the dispatcher, the core routing mechanism of the proving engine, fully transitions to the new memory-aware regime.

In production, this change means that the cuzk daemon can now safely run on machines with varying memory configurations without manual tuning of partition_workers. The system auto-detects available RAM, derives a budget, and dynamically admits or queues work based on actual memory pressure. SRS and PCE entries are loaded on demand and evicted when idle, rather than preloaded unconditionally. The result is a proving engine that is more robust, more portable, and less prone to out-of-memory crashes — all enabled by the removal of a static integer and a semaphore, five call sites at a time.