The Architecture of Zero-Copy: Designing a Pinned Memory Pool for GPU Proving

Introduction

In the course of optimizing a GPU-based zero-knowledge proving pipeline, a critical bottleneck was uncovered: the Host-to-Device (H2D) transfer of the a/b/c synthesis vectors was running at a mere 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 vectors were standard heap allocations, forcing CUDA to stage transfers through a tiny pinned bounce buffer. The solution chosen by the team was "Option B": direct synthesis into pinned memory, eliminating both the reallocation copies during synthesis and the slow staged H2D transfer.

Message 3077 represents a pivotal moment in this optimization effort. It is the assistant's concrete design blueprint, issued just before implementation begins. In this single message, the assistant lays out the architecture across five files, explains the key mechanism that bridges the gap between a CUDA-dependent pool and a generic constraint system library, and then proceeds to inspect the target directory to begin coding. This article examines that message in depth: the reasoning behind the design, the assumptions made, the knowledge required to understand it, and the thinking process that shaped it.

The Message: A Design Blueprint

The message opens with a clear declaration of intent: "Here's the concrete design before I start coding." This signals a deliberate pause—a moment of planning before execution. The assistant then presents an architecture diagram and a key mechanism description, followed by a directory read command to begin implementation.

The architecture is laid out as a simple file map:

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

This is a remarkably concise architecture. Five files, five responsibilities. The design respects the existing module boundaries: bellperson is the generic constraint system library that knows nothing about CUDA, while cuzk-core is the application-specific engine that does. The PinnedPool lives in cuzk-core where CUDA is available. The PinnedBacking struct and release_abc method live in bellperson but are generic—they don't depend on CUDA directly, only on a callback function pointer.

The key mechanism is then explained in detail:

"ProvingAssignment.a/b/c remain Vec<Scalar>, but when pinned backing is provided, they're constructed via Vec::from_raw_parts(pinned_ptr, 0, capacity). A release_abc() method does mem::forget on the Vecs (preventing global-allocator dealloc of pinned ptrs) and returns buffers to pool via callback. ProvingAssignment::Drop calls release_abc() as safety net."

This is the heart of the design. The insight is that Vec::from_raw_parts can be used to create a Vec that looks like a normal heap allocation but is actually backed by CUDA-pinned memory. The trick is in the cleanup: when the Vec is dropped, its normal destructor would try to deallocate through the global allocator, which would be undefined behavior for a CUDA-allocated pointer. To prevent this, release_abc() uses std::mem::forget to leak the Vec's internal buffer without deallocating it, then returns the buffer to the pool via a callback. A custom Drop implementation on ProvingAssignment provides a safety net—if release_abc() was never called explicitly, it fires automatically when the struct goes out of scope.

Why This Message Was Written

The message was written at a critical juncture. The preceding conversation had just concluded the investigative phase: the team had identified the H2D bottleneck, ruled out alternative suspects (tracker lock contention, malloc_trim overhead), and the user had explicitly chosen Option B over Option A. The assistant now needed to transition from investigation to implementation.

But before writing code, the assistant needed to solidify the design. This message serves several purposes:

1. It documents the agreed-upon architecture. The user had just chosen Option B in response to a question. The assistant is now formalizing what that option means in concrete terms—which files to modify, what each file's responsibility is, and how the pieces fit together.
2. It validates the feasibility of the approach. The key mechanism—using Vec::from_raw_parts with a mem::forget + callback pattern—is non-trivial. The assistant is essentially saying "here's how we'll make this work without undefined behavior." This is a critical reasoning step that demonstrates the design is sound before any code is written.
3. It establishes a shared mental model. By writing the architecture down, the assistant creates a reference point. If something goes wrong during implementation, the team can return to this design and ask "did we deviate from the plan?"
4. It serves as a todo list. The five files listed are essentially implementation tasks. The assistant immediately follows up by reading the target directory, ready to create the first file.

How Decisions Were Made

Several design decisions are encoded in this message, each reflecting deliberate trade-offs:

Decision 1: Keep ProvingAssignment.a/b/c as Vec<Scalar> rather than switching to a custom PinnedVec<T> type.

This was not the first idea the assistant considered. In the reasoning traces visible in the preceding messages, the assistant explored several alternatives: a custom PinnedVec<T> wrapper type, an enum abstracting over Vec and PinnedVec, and even a generic allocator parameter. Each was rejected for different reasons:

• A custom PinnedVec<T> would ripple through the entire codebase, changing access patterns everywhere a/b/c are used.
• An enum wrapper would introduce runtime branching in the hot path where push is called tens of millions of times per partition.
• A generic allocator parameter requires nightly Rust. The chosen approach—using Vec::from_raw_parts with mem::forget—is the minimal-change solution. It keeps the field types unchanged, avoids runtime overhead in the synthesis hot path, and works on stable Rust. The only complexity is in the cleanup, which is isolated to release_abc() and the Drop implementation. Decision 2: Use a callback-based return mechanism rather than direct pool access from bellperson. Since bellperson cannot depend on CUDA, the pool reference cannot be passed directly. Instead, the PinnedBacking struct holds a function pointer (callback) that, when invoked, returns the buffer to the pool. This is a clean separation of concerns: bellperson knows it has externally-managed memory and knows how to give it back, but doesn't know or care about CUDA. Decision 3: Integrate the pool with the existing MemoryBudget system. The PinnedPool doesn't just allocate pinned buffers—it tracks their memory usage against the global budget. When a new buffer is allocated, the pool calls try_acquire on the budget. When a buffer is freed, it calls release_internal. This ensures that pinned memory doesn't silently exceed the configured budget, which is critical in a memory-constrained environment where multiple synthesis threads compete for resources. Decision 4: Exact-size allocation rather than a single max-size pool. Rather than pre-allocating buffers at the maximum size needed (which would waste memory for smaller circuit types), the pool allocates buffers to exact requested sizes. The free list tracks (pointer, size) pairs, and checkout finds a buffer that's large enough or allocates a new one. This naturally handles mixed workloads without needing to know circuit types upfront.

Assumptions Made

The design rests on several assumptions, some explicit and some implicit:

Assumption 1: The Vec::from_raw_parts + mem::forget pattern is safe for scalar types like Fr.

This is a critical assumption. Vec::from_raw_parts creates a Vec from a raw pointer, length, and capacity. When the Vec is later forgotten (via mem::forget), its destructor never runs, so the underlying buffer is never deallocated through the global allocator. This is safe only if the elements stored in the Vec have no destructors that need to run. For Fr (a field element type), which is typically a Copy type backed by a fixed-size array of u64 values, this assumption holds. But if Fr ever gained heap-allocated members, this pattern would leak memory.

Assumption 2: The pinned buffers will not be needed after prove_start returns.

The release_abc() method is called after the C++ prove_start function has extracted raw pointers to the a/b/c data. The assumption is that once those pointers are handed off to the GPU, the host-side buffers are no longer needed. If the C++ code retained references to the host buffers after prove_start returned, returning them to the pool would cause use-after-free bugs.

Assumption 3: The MemoryBudget system can handle the pinned pool's allocation patterns.

The budget system was designed for one-shot allocations (e.g., "allocate 8 GiB for this partition"). The pinned pool introduces a new pattern: long-lived allocations that are checked out and checked in repeatedly. The assistant assumes that converting temporary reservations to permanent ones and using release_internal for cleanup will integrate cleanly with the existing budget machinery.

Assumption 4: CUDA's cudaHostAlloc is available and performant.

The design assumes that cudaHostAlloc will provide pinned memory with direct DMA access at PCIe Gen5 line rates. This is true for modern NVIDIA GPUs with proper driver support, but it assumes a specific hardware and software environment. On systems with older GPUs or driver issues, cudaHostAlloc might fall back to slower paths.

Mistakes and Incorrect Assumptions

While the design is sound overall, there are potential issues worth noting:

Potential issue: The mem::forget pattern leaks the elements' destructors.

As noted above, mem::forget on a Vec prevents the Vec's Drop from running, which means the destructors of any elements in the Vec are also skipped. For Fr this is fine, but if the code is later refactored to use a different scalar type with heap-allocated members, this would cause memory leaks. A safer approach might have been to manually drain the Vec (e.g., by setting length to 0) before forgetting, but that would require iterating over all elements, which is expensive for multi-gigabyte vectors.

Potential issue: The callback-based return mechanism adds indirection.

Every time a pinned buffer is returned to the pool, a function pointer is invoked. In a hot loop processing thousands of partitions, this overhead might become measurable. However, compared to the multi-second H2D transfer being eliminated, the microsecond cost of a function call is negligible.

Potential issue: The design doesn't address the Drop safety net's interaction with normal Vec cleanup.

The Drop implementation calls release_abc(), which does mem::forget on the Vecs. But what if release_abc() was already called explicitly? The method would need to be idempotent—after the first call, the Vecs would already be empty (via mem::take), so the second call would be a no-op. The assistant's reasoning traces show awareness of this: "after take, self.a is an empty Vec. Then when the struct is dropped, the field drop for self.a calls Vec::drop on the empty Vec, which is free."

Input Knowledge Required

To understand this message, a reader needs:

1. Knowledge of CUDA memory management. Specifically, the difference between pageable (standard heap) and pinned (page-locked) host memory, and how cudaMemcpyAsync performance depends on the source memory type.
2. Understanding of the proving pipeline. The a/b/c vectors are the constraint system's evaluation matrices—they represent the linear combinations that define the proof's structure. They are synthesized (populated) during the "synthesis" phase and then transferred to the GPU for NTT (Number Theoretic Transform) and MSM (Multi-Scalar Multiplication) operations.
3. Familiarity with Rust's Vec internals. The from_raw_parts, mem::forget, and ManuallyDrop patterns are advanced Rust techniques that require understanding of ownership, destructors, and memory layout.
4. Knowledge of the codebase structure. The distinction between bellperson (generic constraint system library) and cuzk-core (application-specific proving engine) is crucial to understanding why the callback-based separation of concerns is necessary.
5. Understanding of the MemoryBudget system. The budget system manages memory allocation across concurrent synthesis threads, and the pinned pool must integrate with it to avoid exceeding system memory limits.

Output Knowledge Created

This message creates several kinds of knowledge:

1. A concrete implementation plan. Before this message, the solution was an abstract concept ("direct synthesis into pinned memory"). After this message, it's a five-file architecture with specific responsibilities and mechanisms.
2. A validated approach to the Vec/pinned-memory bridge. The from_raw_parts + mem::forget pattern is a reusable technique that could be applied to other scenarios where external memory needs to masquerade as a standard Vec.
3. A template for the PinnedPool API. The design specifies checkout/checkin semantics, budget integration, and exact-size allocation. This becomes the specification that the implementation must satisfy.
4. A risk assessment. By explicitly stating the Drop safety net and the mem::forget pattern, the message implicitly acknowledges the risks of undefined behavior and provides a mitigation strategy.

The Thinking Process

The reasoning visible in this message and its surrounding context reveals a careful, methodical thought process:

Step 1: Problem identification. The investigation had already identified the H2D transfer as the bottleneck. The assistant had ruled out tracker lock contention and malloc_trim overhead through precise instrumentation.

Step 2: Solution exploration. The assistant considered two approaches: Option A (staging memcpy) and Option B (direct synthesis into pinned memory). The user chose Option B, citing concerns about memory bandwidth contention from concurrent synthesis threads.

Step 3: Design refinement. The assistant then explored multiple implementation strategies:

• A custom PinnedVec<T> type (rejected: too many code changes)
• An enum wrapper (rejected: runtime overhead in hot path)
• A generic allocator parameter (rejected: requires nightly Rust)
• The from_raw_parts + mem::forget pattern (chosen: minimal changes, stable Rust) Step 4: Architecture definition. The assistant mapped the chosen approach onto the existing codebase structure, identifying exactly which files need to be modified and what each modification entails. Step 5: Validation. By writing the design down and then immediately proceeding to read the target directory, the assistant is implicitly validating that the plan is complete enough to begin implementation. The thinking process is characterized by a strong preference for minimal changes, careful attention to undefined behavior risks, and respect for existing module boundaries. The assistant consistently chooses the simplest approach that meets the requirements, even when more architecturally elaborate alternatives exist.

Conclusion

Message 3077 is a masterclass in design communication. In a few concise paragraphs, the assistant captures an architecture that spans five files, explains a non-trivial memory management mechanism, and establishes a shared understanding before any code is written. The design respects existing codebase boundaries, minimizes changes to hot paths, and provides safety nets for edge cases.

The message also reveals the assistant's thinking style: methodical, risk-aware, and pragmatic. Rather than diving straight into implementation, the assistant pauses to document the plan, validate the approach, and ensure alignment with the user's chosen direction. This is the kind of message that, in a complex engineering project, can save hours of miscommunication and rework.

The implementation that follows will build on this foundation, creating a PinnedPool that eliminates the H2D bottleneck and dramatically improves GPU utilization. But the real intellectual work—the design decisions, the trade-off analysis, the risk assessment—happens right here, in this single message.