The Critical Edit: Bridging CUDA Pinned Memory Across Crate Boundaries

In the investigation of GPU underutilization in the cuzk proving pipeline, a single message marks the transition from design deliberation to concrete implementation. Message [msg 3087] is deceptively brief — just one sentence followed by an edit command — but it represents the culmination of extensive architectural reasoning and the first tangible step toward solving one of the most elusive performance bottlenecks in the system.

The Bottleneck That Wouldn't Be Ignored

The cuzk proving daemon was suffering from a persistent and puzzling problem: GPU utilization hovered around 50%, wasting half the available compute capacity. Earlier rounds of instrumentation had ruled out the usual suspects — lock contention in the status tracker, overhead from malloc_trim, and mutex acquisition in the C++ GPU code path. The true culprit, identified in the preceding chunk ([chunk 22.0]), was the Host-to-Device (H2D) transfer of the a/b/c synthesis vectors inside execute_ntts_single. These transfers were running at 1–4 GB/s instead of the PCIe Gen5 x16 line rate of approximately 50 GB/s.

The root cause traced back to memory allocation: the a/b/c vectors were standard heap allocations (malloc), forcing CUDA to stage transfers through a small pinned bounce buffer. While the SRS points used in MSM operations benefited from direct DMA via cudaHostAlloc, the synthesis vectors had no such privilege. The result was a bottleneck that varied wildly — NTT kernel phases ranged from 287ms to 8918ms depending on memory bandwidth contention from concurrent synthesis threads.

The Design Journey: Option B Emerges

The preceding message ([msg 3070]) contains an extraordinary amount of reasoning that sets the stage for this edit. The assistant considered and rejected multiple approaches before arriving at the final design.

The first instinct was to create a PinnedVec<T> wrapper type that mimics Vec behavior while managing CUDA-allocated buffers internally. But this would ripple through the entire codebase — every access pattern, every pointer extraction in prove_start, every cleanup operation would need updating. The assistant briefly considered an enum wrapper to abstract over Vec or PinnedVec, but recognized that this would introduce runtime branching in the hot path where push is called tens of millions of times per partition, defeating the purpose of the optimization.

Another approach involved function pointer callbacks registered at startup, allowing bellperson to access the pool without directly depending on CUDA. But the assistant realized that using a custom allocator with Vec::from_raw_parts causes undefined behavior on drop, since the vector would try to deallocate through the global allocator.

The breakthrough came with the ManuallyDrop<Vec<T>> insight. The assistant realized that std::mem::take on a Vec<T> replaces it with an empty Vec::new(), whose drop is a no-op. This meant the field types could remain Vec<Scalar> — no codebase-wide refactoring needed — while the pinned memory was safely forgotten and returned to the pool via a callback.

The Message: A Single Sentence of Intent

The subject message itself reads:

Now I'll add the PinnedBacking struct and the release_abc/new_with_pinned methods. I need to be careful — ProvingAssignment is generic over Scalar: PrimeField, but the pinned backing is just raw bytes: [edit] /tmp/czk/extern/bellperson/src/groth16/prover/mod.rs Edit applied successfully.

This brevity is deceptive. The message encapsulates the core architectural challenge of the entire zero-copy initiative: how to bridge the gap between bellperson — a generic proving library that knows nothing about CUDA — and cuzk-core — the engine layer that manages GPU resources and pinned memory pools.

The ProvingAssignment struct is generic over Scalar: PrimeField, but the pinned backing is just raw bytes. This type mismatch is the crux of the design. The assistant's caution — "I need to be careful" — reflects the awareness that any mistake here could introduce undefined behavior, memory corruption, or silent data loss.

The PinnedBacking Mechanism

The design that emerges from this message is elegant in its minimalism. Rather than changing the field types of ProvingAssignment.a, b, and c from Vec<Scalar> to some new wrapper type, the assistant adds a single optional field: pinned_backing: Option<PinnedBacking>. This struct holds the raw pointers to the CUDA-pinned buffers and a callback function (Box<dyn FnOnce(Vec<*mut u8>)>) that returns the buffers to the pool when synthesis is complete.

The lifecycle works as follows:

  1. Before synthesis begins, the pipeline checks out pinned buffers from the PinnedPool (implemented in [msg 3078]).
  2. A new constructor, new_with_pinned, constructs Vec<Scalar> instances via Vec::from_raw_parts(pinned_ptr, 0, capacity), allowing synthesis to push directly into DMA-able memory.
  3. After prove_start extracts the GPU pointers, release_abc() is called. This method uses std::mem::take to replace each Vec with an empty one, then invokes ManuallyDrop::take on the PinnedBacking to safely forget the Vec contents and invoke the return callback.
  4. A custom Drop implementation serves as a safety net: if release_abc was never called explicitly (e.g., due to an error path), the destructor handles cleanup to prevent the undefined behavior that would occur if Vec::drop tried to deallocate pinned memory through the global allocator.

Assumptions and Potential Pitfalls

The design makes several critical assumptions. First, it assumes that the scalar type Scalar is Copy — that its elements have no destructors that need running. For the field types used in Filecoin proofs (e.g., Fr), this is true, but the generic nature of ProvingAssignment means this assumption is encoded implicitly rather than enforced by the type system.

Second, the approach assumes that Vec::from_raw_parts with a pinned pointer is safe when the vector is never dropped through the standard allocator path. The release_abc method and Drop implementation are carefully designed to ensure this invariant holds, but any future code path that accesses self.a, self.b, or self.c after release_abc would operate on empty vectors — a potential source of subtle bugs.

Third, the pool integration assumes that buffer sizes can be matched exactly. The PinnedPool allocates buffers to exact requested sizes and tracks them in a free list. Under mixed workloads (SnapDeals requiring ~2.59 GiB, PoRep requiring ~4.17 GiB), this naturally handles different circuit types without size-class waste.

Input Knowledge Required

To understand this message, one must grasp several layers of context: the GPU utilization investigation that identified the H2D bottleneck; the architecture of the proving pipeline where ProvingAssignment is synthesized and then transferred to the GPU; the memory model of CUDA where cudaHostAlloc provides directly DMA-able memory while standard malloc requires staging through a pinned bounce buffer; and Rust's ownership semantics where Vec::from_raw_parts creates a vector that assumes ownership through the global allocator, making it incompatible with externally-managed memory without careful intervention.

Output Knowledge Created

This message creates the bridge between the PinnedPool (implemented in the previous message) and the ProvingAssignment that consumes its buffers. It establishes the pattern — PinnedBacking + release_abc + Drop safety net — that will be replicated across the integration points. The subsequent messages ([msg 3088] through [msg 3093]) flesh out the implementation, adding the constructor, the release method, the Drop impl, and updating prove_start in the supraseal module to call release_abc after extracting GPU pointers.

The Significance

This single edit is the keystone of the zero-copy pipeline. Without it, the PinnedPool would be an isolated component with no consumer. The cargo check that follows ([msg 3094]) reveals only a pre-existing visibility error unrelated to the new code, confirming that the core integration compiles cleanly. The path is now clear to wire the pool into the engine startup and evictor, pass it through the synthesis dispatch, and finally test whether the H2D bottleneck — the source of that frustrating 50% GPU utilization — has been eliminated.