The Zero-Copy Wiring: Integrating a Pinned Memory Pool into the cuzk Proving Pipeline

Introduction

In the previous segment of this opencode session, the team had diagnosed a critical performance bottleneck in the cuzk GPU proving daemon: GPU utilization was stuck at roughly 50% because the host-to-device (H2D) transfer of assignment vectors from Rust heap memory to the GPU was bottlenecking at 1–4 GB/s through CUDA's internal pinned bounce buffer, rather than achieving the PCIe Gen5 line rate of ~50 GB/s. The solution—a zero-copy pinned memory pool (PinnedPool) integrated with the existing MemoryBudget system—had been designed and its core components implemented. The PinnedPool struct with its checkout/checkin/shrink methods, the PinnedBacking wrapper, the release_abc() method, and the new_with_pinned() constructor all compiled cleanly.

But a design on paper is not a working system. 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 segment captures that wiring effort in its entirety—the methodical integration of a performance optimization into a production-grade proving pipeline, verified through compilation and grep-based consistency checks, and finally packaged into a deployable Docker image.

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 pool is initialized with a capacity derived from the memory budget, ensuring that pinned memory allocations are bounded by the same budget that governs SRS and PCE memory usage.

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—it would risk deadlock. 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 that minimizes latency in the memory pressure path.

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 that even careful manual review can miss.

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. This small discipline of checking prerequisites before heavyweight operations is a hallmark of robust engineering practice: it costs milliseconds and prevents a multi-minute build from failing due to a missing file.

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 that touched multiple source files across the dependency tree. 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." A build that finishes with exit code 0 but produces no image (due to a misconfigured Dockerfile or a build stage that was accidentally skipped) would be a silent failure that could waste hours of debugging time during deployment.

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. This is a critical safety property: the memory pressure path can reclaim resources from the pool, but it cannot interfere with active proving work.

The fallback path is also important for the initial deployment phase. When the pinned pool binary is first deployed to the remote test machine, there will be a period of observation and tuning. The pool's capacity may need adjustment based on real-world workload patterns. During this tuning phase, the fallback path ensures that the system remains operational even if the pool configuration is suboptimal. The team can observe the fallback rate (how often partitions fall back to heap allocation) as a diagnostic signal to guide capacity tuning.

The Capacity Hint System

The SynthesisCapacityHint mechanism deserves special attention because it is the key enabler of efficient pinned allocation. The hint system works as follows:

First partition of a proof type: No hint is available. The system synthesizes with standard heap allocation, measuring the exact sizes of the a, b, and c vectors.
Hint caching: After the first synthesis completes, the measured sizes are cached as a SynthesisCapacityHint associated with the proof type.
Subsequent partitions: The hint is available. The synthesis functions use it to check out pre-sized pinned buffers from the pool, avoiding any reallocation or waste.
Pool exhaustion: If the pool cannot satisfy the requested sizes, the system falls back to heap allocation for that partition. This design ensures that pinned memory is used efficiently. Without the hint system, the synthesis functions would need to allocate pinned buffers conservatively (oversized) or dynamically (requiring reallocation). The hint system eliminates both problems by providing exact size information after a one-partition warmup period. The hint system also interacts cleanly with the memory budget. The budget determines the total amount of pinned memory available. The hint system determines how much of that budget each partition needs. Together, they ensure that the pool is sized appropriately for the workload: the budget sets the upper bound, and the hints ensure that the bound is used efficiently.

Lessons in Systematic Integration

The work captured in this segment 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. Every function that needs the pool declares it as a parameter, making the data flow visible and auditable.

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. The combination of tools provides a safety net that no single verification method could achieve.

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. The system should always have a path forward, even if that path is slower than optimal.

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. This principle extends to any operation with a high cost of failure: verify the prerequisites first.

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. A build that produces no artifact is indistinguishable from a successful build without this verification step.

Compilation catches what review misses. The missed prove_start call site in supraseal.rs ([msg 3192]) is a telling example. The assistant had carefully reviewed the code and believed all call sites were updated. But the compiler found one that was missed—an old partitioned proving path using the old cleanup pattern. This demonstrates that even the most careful manual review cannot replace compilation as a verification tool. The compiler is the ultimate authority on whether the code is consistent.

The Broader Context: From Diagnosis to Deployment

This segment sits within a larger arc that spans multiple segments of the opencode session. The journey began in [segment 21], where the team added precise timing instrumentation to the GPU worker loop and finalizer to understand why GPU utilization was stuck at ~50%. The timing data revealed that the GPU mutex was held for 1.6–7.0 seconds per partition, but the GPU was only actively computing for ~1.2 seconds. The remaining time was spent in ntt_kernels (H2D transfer) copying a/b/c vectors from unpinned Rust heap memory.

In [segment 22], the team identified the root cause: cudaMemcpyAsync from unpinned memory forced CUDA to stage through a tiny internal pinned bounce buffer at only 1–4 GB/s, instead of the PCIe Gen5 line rate of ~50 GB/s. The solution was designed: a zero-copy pinned memory pool (PinnedPool) integrated with the MemoryBudget system. The core components were implemented.

Now, in [segment 23], the team completed the wiring, integrating the pool into every layer of the proving pipeline, and built a deployable Docker image. The next step will be deployment to the remote test machine, where the primary verification metric will be the ntt_kernels time in the CUZK timing logs. If the fix works as designed, this time should drop from the current 2–9 second range down to under 100 milliseconds, effectively eliminating the H2D bottleneck and raising GPU utilization toward 100%.

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 segment 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—the kind of methodical, verification-driven work that transforms a good design into a working system.## References

[1] Chunk article: "The Final Mile: Wiring a Zero-Copy Pinned Memory Pool Through the cuzk Proving Pipeline" — [chunk 23.0]

[2] Segment 21: GPU utilization investigation with timing instrumentation

[3] Segment 22: Root cause identification and PinnedPool design

[4] Segment 23: PinnedPool wiring and Docker build (this segment)