The Final Mile: Wiring a Zero-Copy Pinned Memory Pool Through the cuzk Proving Pipeline

Introduction

In the world of high-performance GPU computing, the difference between 50% utilization and 100% utilization can be the difference between a system that works and a system that excels. For the cuzk proving daemon—a GPU-accelerated zero-knowledge proof generator for the Filecoin network—that difference came down to a single bottleneck: the host-to-device (H2D) transfer of the a, b, and c assignment vectors from unpinned Rust heap memory to the GPU. The vectors were being copied at 1–4 GB/s through a tiny CUDA internal bounce buffer instead of the PCIe Gen5 line rate of ~50 GB/s, leaving the GPU idle for 2–8 seconds out of every partition cycle.

The solution was a zero-copy pinned memory pool (PinnedPool), integrated with the existing MemoryBudget system. The core components—the PinnedPool struct with its checkout/checkin/shrink methods, the PinnedBacking wrapper, the release_abc() method, and the new_with_pinned() constructor—had already been implemented and compiled cleanly in earlier work. But the critical task remained: wiring this pool into the engine's complex asynchronous pipeline, threading the Arc<PinnedPool> reference through every layer from Engine::new() to the per-partition synthesis functions, and ensuring a graceful fallback path for pool exhaustion.

This chunk of the opencode session captures that wiring effort in its entirety. It is the story of how a carefully designed performance optimization was methodically integrated into a production-grade proving pipeline, verified through compilation and grep-based consistency checks, and finally packaged into a deployable Docker image. It is a case study in the discipline required to transform a correct design into a working system.

The Architecture of the Wiring

The PinnedPool integration touched three major subsystems: the engine (cuzk-core/src/engine.rs), the synthesis pipeline (pipeline.rs), and the bellperson proving library (bellperson/src/groth16/). Each required different modifications, but the overall pattern was consistent: thread the Arc<PinnedPool> reference through the call chain, check out buffers when a capacity hint is available, create ProvingAssignment instances backed by pinned memory, and return buffers to the pool after GPU kernels complete.

The data flow, as traced by the assistant in [msg 3196], follows eight distinct stages:

Engine startup (Engine::new): Creates the Arc<PinnedPool> backed by the memory budget
Evictor integration: The pinned pool is the first thing tried for memory pressure—fast, no lock needed
Dispatch chain: The pinned_pool flows through Engine::start → dispatcher task → dispatch_batch → process_batch → PartitionWorkItem
Synthesis worker loop: Extracts pool_ref from item.pinned_pool and passes to synthesize_partition
Per-partition synthesis: Passes through synthesize_auto → synthesize_with_hint
Buffer checkout: When pool and capacity hint are available, checks out three pinned buffers and creates ProvingAssignment::new_with_pinned instances
GPU kernel execution: prove_start calls prover.release_abc() which returns buffers to the pool via a callback
Fallback path: If pool checkout fails or no hint is available, falls back to standard heap allocation This trace reveals the architectural philosophy behind the integration: the pinned pool is an optimization that sits transparently in the existing pipeline. When it works, the system benefits from zero-copy transfers. When it doesn't—when the pool is exhausted or the capacity hint is unavailable—the system degrades gracefully to the original behavior. There is no single point of failure.

Threading Through the Engine

The engine (engine.rs) is the heart of the cuzk daemon, a ~3259-line file containing the main event loop, the dispatcher, the synthesis worker, and the GPU worker. Threading the pinned_pool through this file required modifications at multiple levels.

The first modification was in Engine::new(), where the Arc<PinnedPool> is created and stored as a field on the Engine struct. This is the root of the dependency: every other component that needs the pool must receive it from here.

The second modification was in the evictor callback. The SrsManager is behind Arc<tokio::sync::Mutex<SrsManager>>, so the evictor callback cannot use blocking_lock() from async context. Instead, it uses try_lock(). The pool's shrink() method is called during eviction to free pinned memory when SRS or PCE needs more space. The assistant's design makes the pinned pool the first resort for memory pressure because shrinking the pool is fast and requires no lock—a deliberate prioritization.

The third and most extensive modification was threading the pinned_pool through the dispatch chain. The dispatch_batch function, which wraps process_batch, needed the pool parameter added to its signature. Then all five call sites of dispatch_batch within the dispatcher closure needed to pass &pinned_pool. Similarly, process_batch needed the parameter added to its signature, and its two call sites (inside dispatch_batch) needed updating. Finally, the PartitionWorkItem struct—which carries per-partition data to the synthesis workers—needed a new field: pinned_pool: Option<Arc<PinnedPool>>.

The assistant approached this threading task methodically. At [msg 3161], the process_batch signature was updated. At [msg 3162], the two PartitionWorkItem constructions were updated. At [msg 3164], the assistant used grep -n 'process_batch(' to find all call sites. At [msg 3165], grep -n 'dispatch_batch(' found the five call sites of the wrapper. At [msg 3167], the dispatch_batch signature was updated. Then, across messages [msg 3168] through [msg 3178], the assistant methodically updated each call site, reading the surrounding code to ensure the parameter was inserted correctly.

The verification at [msg 3180] was a simple but powerful check: grep -c '&pinned_pool,' /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs returned 6—matching the expected count of 5 dispatch_batch calls plus 1 process_batch call within dispatch_batch. This grep-based verification is a hallmark of the assistant's approach: rather than relying on compilation alone to catch errors, the assistant proactively counts occurrences to ensure completeness.

The Synthesis Path

The synthesis pipeline (pipeline.rs) required a different kind of modification. Here, the pinned_pool parameter needed to be threaded through the synthesis functions themselves—synthesize_partition, synthesize_auto, and synthesize_with_hint—all of which now accept Option<&Arc<PinnedPool>>.

The critical innovation is the synthesize_circuits_batch_with_prover_factory function, introduced to handle the pinned allocation path. When a capacity hint is available and the pool has buffers, this function creates a closure that:

Checks out three pinned buffers (for a, b, c) from the pool
Creates ProvingAssignment::new_with_pinned instances backed by those buffers
Forgets the PinnedAbcBuffers wrapper, transferring ownership to the PinnedBacking struct
Sets up a return callback that checks buffers back into the pool via PinnedBuffer::from_raw after GPU kernels complete The capacity hint system is essential. The SynthesisCapacityHint mechanism caches exact vector sizes after the first synthesis, enabling pre-sized pinned allocation. This is important because pinned memory is a finite resource—allocating oversized buffers would waste the pool, while undersized buffers would require reallocation. When the hint is unavailable (e.g., on the first partition of a new proof type), the system falls back to unpinned allocation. The memory safety of this design is carefully considered. When a Vec is backed by pinned memory, the normal Vec::drop would try to free the memory via the global allocator, which would corrupt the pool. Instead, release_abc() calls mem::forget on the Vecs (leaking the memory intentionally) and returns the pointers to the pool via a callback. The Drop impl on ProvingAssignment calls release_abc() as a safety net if the explicit call was missed. This two-layer safety ensures that pinned memory is never freed by the wrong allocator.

The Bellperson Integration

The bellperson library, which provides the Groth16 proving infrastructure, required modifications in two files: prover/mod.rs and prover/supraseal.rs.

In prover/mod.rs, the ProvingAssignment struct gained a pinned_backing: Option<PinnedBacking> field. The new_with_pinned() constructor creates Vec<Scalar> via Vec::from_raw_parts on pinned pointers, so the Vec's data lives in pinned memory. The release_abc() method uses mem::forget on the Vecs and returns buffers to the pool via a callback.

In prover/supraseal.rs, the prove_start function now calls prover.release_abc() instead of the old cleanup pattern (prover.a = Vec::new()). This ensures pinned buffers are properly returned to the pool after GPU kernels complete, rather than being leaked or freed incorrectly.

The assistant discovered a missed call site during compilation ([msg 3192]). The prove_start function in supraseal.rs had an old partitioned proving path that was still using the old cleanup pattern. The compilation error revealed this oversight, and the assistant fixed it with a targeted edit at [msg 3193]. This is a classic example of why compilation is not just a formality—it is a verification tool that catches real bugs.

Compilation and Verification

The compilation milestone at [msg 3195] was a moment of quiet triumph. After the extensive wiring effort spanning dozens of edits across multiple files, cargo check --features cuda-supraseal passed cleanly. All warnings were pre-existing. The new pinned pool wiring compiled without errors.

But the assistant did not stop at compilation. Before building the Docker image, the assistant performed a comprehensive mental walkthrough of the data flow at [msg 3196], tracing the pinned_pool reference through all eight stages from engine startup to GPU kernel completion. This walkthrough served as a final sanity check, ensuring that every connection was sound and every ownership transfer was correct.

The assistant also performed grep-based verification at multiple points. At [msg 3180], the count of &pinned_pool, occurrences was verified to be 6. At [msg 3185], the type match between item.pinned_pool.as_ref() (producing Option<&Arc<PinnedPool>>) and the synthesis function signature (expecting Option<&Arc<PinnedPool>>) was confirmed. These checks, while simple, provided confidence that the wiring was complete and consistent.

The Docker Build and Deployment Readiness

With compilation verified, the next step was to produce a deployable artifact. The assistant first verified the existence of the Dockerfile at [msg 3197] with a simple ls command—a pre-flight check before committing to a multi-minute build.

The Docker build at [msg 3198] used DOCKER_BUILDKIT=1 docker build --no-cache -f Dockerfile.cuzk-rebuild -t cuzk-rebuild:pinned1 .. The --no-cache flag ensured a completely fresh build, essential after code changes. The build context was substantial—112.40 MB of source code and dependencies. The base image was nvidia/cuda:13.0.2-devel-ubuntu24.04, providing the CUDA 13.0 development environment needed to compile the GPU kernels.

After the build completed, the assistant verified the result at [msg 3199] by querying docker images cuzk-rebuild:pinned1. The image existed at 27.7 MB, built at 20:24:30 CET on March 13, 2026. This verification was not redundant—it was a deliberate check against the possibility of silent build failure. The assistant understood that "the build ran" does not guarantee "the build succeeded."

The user's response at [msg 3200] was an empty message—a silent signal that marked the transition from implementation to deployment. In the context of the session, this emptiness carried meaning: the user had reviewed the work, saw that the build was successful, and implicitly consented to proceed with deployment. The assistant's next step would be to extract the binary from the Docker image, copy it to the remote test machine, and run it against real proving workloads.

The Fallback Path and Graceful Degradation

Throughout the integration, the assistant maintained a consistent design principle: the pinned pool is an optimization, not a requirement. If the pool is exhausted (budget full) or no capacity hint is available, the system falls back to standard new_with_capacity with unpinned heap allocations. This graceful degradation ensures that the system never crashes or deadlocks due to pool exhaustion—it simply runs at the old slower speed for that partition.

This fallback path is critical for production reliability. Pinned memory is a finite resource, bounded by the MemoryBudget system. Under heavy load, multiple partitions may compete for pinned buffers. The pool's checkout method may return None if the budget is full. In that case, the synthesis functions silently fall back to heap allocation, and the partition proceeds without pinned memory benefits.

The evictor integration also respects this fallback principle. When the evictor fires to free memory for SRS or PCE, it calls pool.shrink() to free idle pinned buffers. But it never forcibly evicts buffers that are currently checked out. The shrink() method only frees buffers that have been checked back in (idle), ensuring that in-flight operations are never disrupted.

Lessons in Systematic Integration

The work captured in this chunk offers several lessons for systematic integration of performance optimizations into complex systems:

Thread the dependency from the root. The Arc<PinnedPool> is created once at Engine::new() and propagated downward through every layer. This avoids the complexity of global state or lazy initialization, and makes the dependency explicit in every function signature.

Verify with multiple tools. The assistant used compilation (cargo check), grep-based occurrence counting, type matching, and mental walkthroughs to verify the integration. Each tool catches different kinds of errors: compilation catches type mismatches, grep catches missing call sites, and walkthroughs catch logical errors in data flow.

Design for graceful degradation. The fallback path ensures that the optimization never becomes a single point of failure. This is especially important for memory optimizations, where resource exhaustion is a realistic possibility under production load.

Check prerequisites before heavyweight operations. The ls check on the Dockerfile before the Docker build is a small but important discipline. It costs milliseconds and prevents a multi-minute build from failing due to a missing file.

Verify build artifacts independently. The docker images check after the build confirmed that the image actually existed, with the correct size and timestamp. This is a defense against silent build failures that might otherwise go unnoticed until deployment.

Conclusion

The wiring of the PinnedPool through the cuzk proving pipeline represents the culmination of a deep engineering investigation into GPU underutilization. From the initial discovery that cudaMemcpyAsync from unpinned memory was bottlenecking at 1–4 GB/s, through the design of a zero-copy pinned memory pool, to the methodical threading of the pool through every layer of the proving pipeline, the team demonstrated a disciplined approach to performance optimization.

The final artifact—a 27.7 MB Docker image tagged cuzk-rebuild:pinned1—is ready for deployment to the remote test machine. The primary verification metric will be the ntt_kernels time in the CUZK timing logs, expected to drop from the current 2–9 second range down to under 100 milliseconds. If the fix works as designed, GPU utilization will rise from ~50% toward 100%, and the proving daemon will achieve the throughput that its hardware is capable of delivering.

But regardless of the outcome of that test, the work captured in this chunk stands as a model for systematic integration: start from a clear root cause, design a targeted solution, thread the dependency through every layer, verify with multiple tools, design for graceful degradation, and build a deployable artifact. It is engineering discipline at its finest.