Building the Zero-Copy Bridge: Implementing a Pinned Memory Pool for GPU Proving

Introduction

In the previous chunk of Segment 22, the team had conclusively identified the root cause of a crippling GPU underutilization problem in the cuzk proving daemon. Through precise instrumentation of the C++ gpu_prove_start function, they had traced the bottleneck to the Host-to-Device (H2D) transfer of the a/b/c synthesis vectors inside execute_ntts_single. These vectors—each up to 4.17 GiB for PoRep proofs, totaling 12.5 GiB per partition—were being transferred from standard heap-allocated memory at a paltry 1–4 GB/s instead of the PCIe Gen5 x16 line rate of approximately 50 GB/s. The root cause was mundane but devastating: standard heap allocations forced CUDA to stage transfers through a small pinned bounce buffer, while the SRS points used in MSM operations enjoyed direct DMA via cudaHostAlloc.

The chosen solution was Option B: direct synthesis into pinned memory. Instead of allocating a/b/c vectors on the heap and then copying them to pinned staging buffers (Option A), the synthesis threads would write directly into cudaHostAlloc'd buffers that are directly DMA-able from the moment they are populated. This eliminated both the reallocation copies during synthesis (the Vec::push → realloc → memcpy cycle) and the slow staged H2D transfer. Under heavy memory pressure from 20+ concurrent synthesis threads consuming 400+ GiB of RAM, even a staged memcpy would be severely bottlenecked by contested host memory bandwidth—Option B was not just faster, it was necessary.

This chunk captures the implementation of that solution: the creation of a PinnedPool in cuzk-core, the modification of bellperson's ProvingAssignment to safely hold and release pinned buffers, the wiring of release_abc() into prove_start, and the validation via cargo check that the core integration compiles cleanly. It is a story of careful architectural design, meticulous API verification, and defensive programming against undefined behavior.

The Design Solidified: Why Option B Was the Only Path

The chunk opens with the assistant solidifying the decision to pursue Option B and elaborating on its necessity ([msg 3070]). The reasoning is worth examining because it reveals a deep understanding of the system's performance characteristics under real-world conditions.

The assistant's "Agent Reasoning" block in [msg 3070] is a stream-of-consciousness design document spanning roughly 1,500 words. It walks through the problem from multiple angles: the architectural boundaries between crates, Rust's ownership semantics, CUDA's memory APIs, the budget accounting system, and the size class strategy. The key insight is that the ManuallyDrop<Vec<T>> approach—keeping the field types as Vec<Scalar> but adding an optional PinnedBacking struct—minimizes code changes while avoiding undefined behavior.

The assistant considers and rejects several alternatives before arriving at this design. A custom PinnedVec<T> wrapper type would ripple through the entire codebase, changing access patterns everywhere a/b/c are used. An enum wrapper abstracting over Vec or PinnedVec would introduce runtime branching in the hot path where push is called tens of millions of times per partition. A custom allocator using Rust's nightly Allocator trait would require nightly Rust, which the project does not use. The ManuallyDrop approach emerges as the cleanest minimal-change solution.

The core challenge is the undefined behavior trap. When a Vec is constructed via Vec::from_raw_parts from a pinned pointer, the Vec's Drop implementation will attempt to deallocate that memory through the global allocator. But the memory was allocated by cudaHostAlloc and must be freed by cudaFreeHost. Calling free on a cudaHostAlloc'd pointer is undefined behavior. The release_abc() method solves this by using std::mem::take to replace each Vec with an empty one (whose drop is a no-op), then using ManuallyDrop::take to forget the original Vec's contents before returning the buffer to the pool via a callback. A custom Drop implementation on ProvingAssignment serves as a safety net.

The Build System Reconnaissance

Before writing any code, the assistant needed to understand the existing CUDA FFI infrastructure. A grep for cudaHostAlloc|cudaFreeHost|cudaMallocHost|cuda_host returned no results ([msg 3071]), confirming that no pinned memory allocation bindings existed anywhere in the codebase. This meant the assistant would have to build from scratch.

The next step was understanding how CUDA code was compiled. A glob for **/build.rs found a single result at /tmp/czk/extern/supraseal-c2/build.rs ([msg 3072]). Reading that file ([msg 3073]) revealed that the crate uses sppark::build::ccmd() to compile a single CUDA source file, groth16_cuda.cu. This confirmed that CUDA was already integrated into the build system and that new CUDA API bindings could be added following the same pattern.

The assistant then read cuzk-core/Cargo.toml (<msg id=3074-3075>) to verify the dependency chain. The feature flag cuda-supraseal being the default confirmed that CUDA support was enabled by default. The presence of libc = &#34;0.2&#34; indicated that raw C FFI was already in use, making CUDA FFI additions architecturally consistent. The ec-gpu-gen dependency confirmed that the scalar field types (Fr) are Copy types with no destructors—a critical safety invariant for the ManuallyDrop approach.

With this reconnaissance complete, the assistant declared "OK, now I have enough context" ([msg 3076]) and laid out the concrete architecture ([msg 3077]):

cuzk-core/src/pinned_pool.rs  — PinnedPool (cudaHostAlloc/Free, budget-integrated)
bellperson/src/groth16/prover/mod.rs — PinnedBacking + release_abc on ProvingAssignment
bellperson/src/groth16/prover/supraseal.rs — prove_start uses release_abc
cuzk-core/src/pipeline.rs — synthesis checks out pinned buffers, passes to ProvingAssignment
cuzk-core/src/engine.rs — creates pool, evictor integration

Implementing the PinnedPool

The PinnedPool struct was written to cuzk-core/src/pinned_pool.rs ([msg 3078]). It manages cudaHostAlloc'd buffers with a free list tracking (pointer, size) pairs. It integrates with the existing MemoryBudget system via try_acquire: when allocating a new buffer, the pool acquires budget upfront; when freeing, it releases the budget. The pool uses exact-size allocation rather than pre-sized pools for each circuit type, naturally handling mixed workloads (SnapDeals at ~2.59 GiB per buffer, PoRep at ~4.17 GiB).

However, the initial implementation contained a subtle API mismatch. The try_acquire method on MemoryBudget uses an unusual Rust pattern: pub fn try_acquire(self: &amp;Arc&lt;Self&gt;, amount: u64) -&gt; Option&lt;MemoryReservation&gt;. The assistant had assumed a conventional signature and written code that didn't match. Reading the actual API (<msg id=3079-3080>) revealed the correct pattern, and the assistant applied a fix (<msg id=3081-3082>). This moment of self-correction—catching the mistake before compilation—demonstrates the value of verifying API signatures against source code rather than relying on assumptions.

The read in [msg 3080] also revealed the Two-Phase Release contract for a/b/c memory: after prove_start, the vectors are freed and ~12.5 GiB is returned to the budget; after prove_finish, the remaining ~1.1 GiB is released. This documentation confirmed that the pool's buffer checkout/checkin lifecycle could be mapped onto the existing budget protocol without conflict.

The Module Registration

With the PinnedPool implemented, the assistant registered it in cuzk-core/src/lib.rs (<msg id=3083-3084>). This single-line edit—pub mod pinned_pool;—marked the transition from design to implementation. Before this edit, pinned_pool.rs was a standalone file, invisible to the compiler's module resolution. After this edit, the PinnedPool type became part of the crate's public interface, accessible to engine.rs, pipeline.rs, and any other module that needs to allocate pinned memory.

The Critical Bridge: Modifying ProvingAssignment

The most delicate part of the implementation was modifying bellperson's ProvingAssignment to safely hold and release pinned memory buffers. This required bridging two separate codebases—cuzk-core (where the pool lives) and bellperson (where the vectors are synthesized)—without introducing undefined behavior or breaking the public API.

The assistant read the ProvingAssignment struct definition (<msg id=3085-3086>), confirming that a, b, and c are declared as plain Vec&lt;Scalar&gt;. This was the "before" snapshot that the entire pinned memory effort sought to transform. The struct is generic over Scalar: PrimeField, which meant the pinned backing mechanism must be generic as well.

The PinnedBacking struct was added in [msg 3087]. It holds the raw pinned buffer pointers and a return callback (Box&lt;dyn FnOnce(Vec&lt;*mut u8&gt;)&gt;). The ProvingAssignment struct was extended with a pinned_backing: Option&lt;PinnedBacking&gt; field. This design keeps bellperson CUDA-free while allowing cuzk-core to provide the pool integration through the callback mechanism.

The assistant then added three methods to ProvingAssignment (<msg id=3088-3092>):

  1. new_with_pinned(): An alternative constructor that accepts pre-allocated pinned buffers and constructs Vec&lt;Scalar&gt; instances via Vec::from_raw_parts(pinned_ptr, 0, capacity). This allows synthesis to push directly into DMA-able memory.
  2. release_abc(): Uses std::mem::take to replace each Vec with an empty one (whose drop is a no-op), then invokes the return callback to recycle the buffers. This prevents the global allocator from attempting to free pinned memory.
  3. Custom Drop implementation: A safety net that calls release_abc() if the pinned backing is still present when the struct is dropped. This ensures that even if release_abc() was never called explicitly (e.g., due to an error path), the pinned buffers are properly returned to the pool rather than being deallocated through the wrong allocator. The assistant also added pinned_backing: None to all existing constructors (<msg id=3090, 3092>), including the ConstraintSystem::new() method. This ensured that existing code paths (native proving, non-CUDA builds) continue to work without modification.

The Final Stitch: Wiring release_abc into prove_start

The final integration step was updating the prove_start function in bellperson/src/groth16/prover/supraseal.rs to call release_abc() after extracting pointers ([msg 3093]). The prove_start function is the Rust-side entry point that invokes the C++ GPU proving code. It extracts raw pointers from the a, b, and c vectors and passes them to the C++ layer for GPU computation.

After the pointers have been extracted and handed off to the GPU, the Rust-side vectors are no longer needed by the proving logic. However, they cannot simply be dropped: because their backing memory was allocated via cudaHostAlloc through the PinnedPool, a normal Vec::drop would attempt to free the memory through the global allocator, causing undefined behavior. The release_abc() method must be called instead, which safely returns the buffers to the pool for reuse.

This edit is the critical handshake between the new pinned-memory infrastructure and the existing GPU proving pathway. Without it, the entire zero-copy architecture would be incomplete—the buffers would never be returned to the pool, and the pool would quickly exhaust its capacity, or worse, the program would crash on the next deallocation.

Validation: The Checkpoint

With all the core components implemented, the assistant ran cargo check --features cuda-supraseal to validate the changes ([msg 3094]). The output revealed only a pre-existing visibility error in engine.rs concerning JobTracker and process_partition_result—entirely unrelated to the newly written pinned memory code. The fact that no new errors appeared from the PinnedPool implementation in cuzk-core, the PinnedBacking modifications in bellperson, or the prove_start changes in supraseal.rs was the quiet confirmation that the core architectural changes compile cleanly.

This checkpoint validated several critical assumptions:

The Architecture in Retrospect

Looking back at the implementation, several design decisions stand out as particularly elegant:

The ManuallyDrop pattern: By keeping a/b/c as Vec&lt;Scalar&gt; and adding an optional PinnedBacking field, the assistant avoided changing field types across the entire codebase. The release_abc() method uses std::mem::take to replace each Vec with an empty one, then ManuallyDrop::take to extract and forget the original contents. This is a minimal-change strategy that reduces the risk of introducing bugs in unrelated code paths.

The callback-based abstraction: The PinnedBacking struct stores a Box&lt;dyn FnOnce(Vec&lt;*mut u8&gt;)&gt; callback rather than a direct pool reference. This keeps bellperson CUDA-free while allowing cuzk-core to provide the pool integration. The callback is registered at startup by the higher-level engine layer, which has access to CUDA and the PinnedPool.

The Drop safety net: The custom Drop implementation on ProvingAssignment transforms a fragile manual protocol into a robust RAII pattern. If release_abc() was never called explicitly (e.g., due to an error path, a refactor, or a logic bug), the Drop impl ensures the pinned buffers are returned to the pool rather than being deallocated through the wrong allocator.

The budget integration: The pool holds permanent reservations for allocated buffers, releasing them only when buffers are freed through shrinkage or eviction. Checkout and checkin of buffers from the pool do not touch the budget at all—the reservation is already held by the pool. This cleanly separates buffer lifecycle from budget accounting.

Assumptions and Risks

The implementation makes several assumptions worth examining:

That Fr (the scalar field element) is Copy with no destructor: The ManuallyDrop::take approach "forgets" the elements stored in the Vec. For scalar field elements like Fr, which are plain data with no destructors, this is safe. But if the generic Scalar type were ever instantiated with a type that has drop glue, this would leak resources. The assistant explicitly notes this assumption in the reasoning.

That the callback mechanism is reliable: The PinnedBacking struct stores a callback that returns buffers to the pool. If the callback is never called (e.g., due to a panic during synthesis that bypasses the Drop impl), the buffer leaks. The Drop safety net mitigates this for normal control flow, but double-panics or abort-on-panic configurations could still cause leaks.

That Vec::from_raw_parts with a CUDA-allocated pointer is valid: The Rust specification requires that pointers passed to Vec::from_raw_parts come from the correct allocator. CUDA's cudaHostAlloc returns host-pinned memory that is also valid for CPU access, so reading and writing through the Vec's pointer is safe. But the allocator mismatch is precisely why the ManuallyDrop trick is needed—the Vec must never be allowed to call its own dealloc.

That the pool's free list is efficient under contention: The pool uses a Mutex&lt;Vec&lt;...&gt;&gt; free list. Under heavy concurrent synthesis (20+ threads), contention on this mutex could become a bottleneck. The pool's design minimizes checkout/checkin frequency (once per partition, not per element), which mitigates this concern.

The Road Ahead

With the core pinned memory integration compiling cleanly, the team is positioned to complete the remaining integration work. The PinnedPool must be created at engine startup, integrated with the evictor to handle memory pressure, and passed through the synthesis dispatch so that every partition's a/b/c vectors are allocated from pinned memory. The pre-existing JobTracker visibility error must also be fixed, but it is a straightforward visibility change unrelated to the pinned memory work.

The zero-copy pipeline promises to collapse the H2D transfer from seconds to milliseconds, dramatically improving GPU utilization by allowing the NTT setup to overlap cleanly with the compute phases of other partitions. When fully integrated and deployed, this optimization could double the proving throughput of the entire cuzk system—a transformative improvement born from careful instrumentation, rigorous design reasoning, and meticulous implementation.

Conclusion

This chunk captures the transition from investigation to implementation in one of the most critical performance optimizations of the cuzk proving engine. The team had identified the H2D transfer of a/b/c vectors as the root cause of ~50% GPU underutilization. The solution—a zero-copy pinned memory pool—required careful coordination across multiple crates, deep understanding of Rust's ownership semantics, and defensive programming against undefined behavior.

The implementation was methodical and iterative: design the architecture, verify the build system, implement the pool, fix API mismatches, register the module, modify the core data structure, add safety nets, wire the integration point, and validate with compilation. Each step built on the previous one, and each verification caught potential issues before they could propagate.

The result is a cleanly-compiling core integration that positions the team to complete the wiring and finally test whether the H2D bottleneck has been eliminated. The zero-copy pipeline is no longer a design document—it is a working implementation awaiting integration.