The PinnedPool: A Zero-Copy Turning Point in GPU Utilization

On the surface, message [msg 3078] appears trivial — a single line confirming that a file was written successfully. But this message marks the precise moment when an abstract performance investigation crystallized into concrete code. The assistant wrote /tmp/czk/extern/cuzk/cuzk-core/src/pinned_pool.rs, the foundational component of a zero-copy pinned memory pool designed to eliminate a crippling Host-to-Device (H2D) transfer bottleneck that had been limiting GPU utilization to roughly 50% in the cuzk proving pipeline.

The Investigation That Led Here

To understand why this message matters, one must trace back through the preceding investigation. The team had been battling a perplexing GPU underutilization problem. The GPU worker loop showed idle gaps of seconds at a time, yet the actual compute phases (MSM, batch_add, tail_msm) were completing in a stable ~1.2 seconds per partition. The culprit was not lock contention or malloc_trim overhead, as initially suspected. Precise instrumentation using GPU_TIMING and FIN_TIMING markers had ruled those out.

The true bottleneck was discovered inside the C++ gpu_prove_start function, specifically in the execute_ntts_single phase where the a/b/c synthesis vectors were being transferred from host to device. These transfers were running at a mere 1–4 GB/s instead of the PCIe Gen5 x16 line rate of approximately 50 GB/s. The root cause was mundane but devastating: the a/b/c vectors were allocated as standard heap memory via malloc, forcing CUDA to stage transfers through a tiny pinned bounce buffer. Meanwhile, the SRS points used in MSM operations enjoyed direct DMA access because they were allocated with cudaHostAlloc.

The nvtop observations told the story plainly: RX bandwidth would spike to 50 GB/s during compute phases but collapse to 1–4 GB/s during the H2D transfer gaps. The ntt_kernels phase varied wildly from 287ms to 8918ms depending on memory bandwidth contention from concurrent synthesis threads. With 20+ synthesis threads running simultaneously and 400+ GiB of host memory in use, the memory bus was already contested — making a simple memcpy-based fix insufficient.

The Design Decision: Option B

The assistant had presented two approaches in [msg 3069]. Option A was a staging memcpy: check out three pinned buffers from a pool, copy a/b/c into them (~80ms for 8 GiB at DDR5 bandwidth), and pass pinned pointers to the C++ function. Option B was direct synthesis into pinned memory — zero-copy from the start, but requiring changes to bellperson's ProvingAssignment struct.

The user chose Option B, and the assistant's reasoning in [msg 3070] elaborated on why this was the right call. Under heavy memory pressure, even the ~80ms memcpy of Option A would be severely degraded because host memory bandwidth was already contested by concurrent synthesis threads. Moreover, the current Vec allocation via malloc already performed implicit copies during reallocation — every push that triggered a grow operation caused a memcpy to a new allocation. With pre-sized pinned buffers, both the reallocation copies and the H2D staging would be eliminated.

The Architecture Takes Shape

Message [msg 3077] laid out the concrete architecture:

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

The key mechanical insight was elegant: ProvingAssignment.a/b/c would remain Vec<Scalar>, but when pinned backing was provided, they would be constructed via Vec::from_raw_parts(pinned_ptr, 0, capacity). A release_abc() method would use mem::forget on the Vecs — preventing the global allocator from deallocating pinned pointers — and return the buffers to the pool via a callback. A custom Drop implementation on ProvingAssignment would call release_abc() as a safety net, ensuring that even if the explicit release was missed, the pinned buffers would not leak.

What the PinnedPool File Contained

The file written in [msg 3078] implemented the PinnedPool struct. Based on the assistant's reasoning and subsequent edits, we can reconstruct its design. The pool managed cudaHostAlloc'd buffers with a free list tracking (pointer, size) pairs. It integrated with the existing MemoryBudget system via try_acquire: when allocating a new buffer, the pool would acquire budget upfront; when freeing, it would release the budget. This kept the memory accounting consistent without requiring separate checkout/checkin operations.

The pool used exact-size allocation rather than pre-sized pools for each circuit type. SnapDeals required approximately 2.59 GiB per vector (81M × 32B), totaling 7.8 GiB per partition, while PoRep required 4.17 GiB per vector (130M × 32B), totaling 12.5 GiB. A free list with best-fit or first-fit allocation naturally handled mixed workloads without needing to know circuit types upfront.## The Thinking Process Behind the Implementation

The assistant's reasoning in [msg 3070] reveals a remarkably thorough design journey. The initial instinct was to create a PinnedVec<T> wrapper type that would mimic Vec behavior while managing CUDA-allocated buffers internally. But this approach ran into a fundamental problem: switching ProvingAssignment.a/b/c from Vec<Scalar> to PinnedVec<Scalar> would ripple through the entire codebase, requiring updates to all access patterns, pointer extraction logic in prove_start, and cleanup operations.

The assistant then considered an enum wrapper to abstract over Vec or PinnedVec, but rejected it because runtime branching in the hot path — where push gets called tens of millions of times per partition — would defeat the performance purpose. This shows a deep understanding of the performance characteristics of the synthesis loop.

The next consideration was using Rust's Allocator trait on nightly to create Vec<Fr, PinnedAllocator>, but this was rejected because it required nightly Rust, which was not an option for this production codebase.

The breakthrough came with the ManuallyDrop<Vec<T>> approach. By wrapping vectors allocated from pinned memory in ManuallyDrop, the assistant could control exactly when deallocation happened — or more precisely, prevent it from happening at all through the standard Vec::drop path. The release_abc() method would use std::mem::take to replace each Vec with an empty one (since Vec::new() is the default), then ManuallyDrop::take to extract the inner Vec and forget it, preventing the global allocator from freeing the pinned pointers. The pinned buffer itself would be returned to the pool via a callback stored in the PinnedBacking struct.

This design is notable for its minimal invasiveness. The a, b, and c fields remain Vec<Scalar> — no type changes needed. The only addition is an optional pinned_backing field. Existing code paths that don't use pinned memory continue to work unchanged. The Drop implementation serves as a safety net, ensuring that even if release_abc() is never called explicitly, the pinned buffers are properly returned.

Assumptions and Their Validity

The design makes several assumptions worth examining. First, it assumes that Fr (the scalar field element type) is Copy and has no custom Drop implementation. This is critical because ManuallyDrop::take bypasses the destructor of the Vec's elements. If Fr had heap-allocated internals, forgetting its elements would cause memory leaks. For the BLS12-381 scalar field used in Filecoin proofs, this assumption holds — Fr is a simple 32-byte struct with no destructor.

Second, the design assumes that Vec::from_raw_parts(pinned_ptr, 0, capacity) is safe when the pinned memory is properly aligned and sized. This is a reasonable assumption for cudaHostAlloc'd memory, which returns page-aligned buffers, but it relies on the caller ensuring that the capacity doesn't exceed the allocated size.

Third, the design assumes that the pool's budget integration via try_acquire is sufficient to prevent out-of-memory conditions. The pool holds permanent reservations for allocated buffers, releasing them only when buffers are freed. When checking out a buffer, no additional budget is needed since it's already reserved. This is sound as long as the initial budget calculation accounts for the maximum number of pinned buffers that will be in flight simultaneously.

Input Knowledge Required

To fully understand this message, one needs knowledge of several domains. The CUDA programming model is essential — specifically the distinction between pageable host memory (which requires staged transfers through a pinned bounce buffer) and pinned host memory (which enables direct DMA access via cudaMemcpyAsync). Understanding PCIe Gen5 bandwidth characteristics (~50 GB/s for x16) provides context for why the observed 1–4 GB/s transfer rate was so problematic.

Rust's memory model is equally important. The Vec::from_raw_parts function, ManuallyDrop, std::mem::take, and the interaction between custom allocators and Vec::drop are all subtle features that must be used correctly to avoid undefined behavior. The assistant's careful reasoning about what happens when a Vec constructed from a pinned pointer is dropped — namely, that the global allocator would attempt to free the pinned memory, causing undefined behavior — shows a deep understanding of Rust's ownership and allocation semantics.

Knowledge of the Filecoin proof pipeline is also necessary. The a/b/c vectors represent evaluations of the A, B, and C polynomials from the Groth16 proving system. Their sizes — 81M elements for SnapDeals, 130M for PoRep — determine the memory requirements and explain why the H2D transfer dominates the proving time. The ProvingAssignment struct in bellperson is the central data structure that carries these vectors from synthesis to GPU proving.

Output Knowledge Created

This message created the PinnedPool module — the foundational building block for the zero-copy pipeline. Subsequent messages would build on this foundation: adding PinnedBacking and release_abc to bellperson's ProvingAssignment ([msg 3087]), updating prove_start to use release_abc ([msg 3093]), and registering the module in lib.rs ([msg 3084]).

The cargo check in [msg 3094] confirmed that the core pinned memory integration compiled successfully, with only a pre-existing visibility error unrelated to the new code. This validated that the design was sound and positioned the team to wire the pool into the engine startup, evictor, and synthesis dispatch.

The Significance of This Single Line

A single tool call — [write] /tmp/czk/extern/cuzk/cuzk-core/src/pinned_pool.rs — may seem unremarkable in isolation. But this message represents the culmination of an extensive investigation into GPU underutilization, the distillation of multiple design iterations into a coherent architecture, and the first concrete step toward eliminating a bottleneck that was costing seconds per partition. It is the moment when analysis became action, when the abstract understanding of the H2D transfer bottleneck was translated into the first line of code that would eventually solve it.

The PinnedPool file written here would become the backbone of a zero-copy pipeline that promised 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. In the broader narrative of the cuzk proving engine optimization, this message is the turning point.