The Zero-Copy Pivot: How a CUDA Memory Bottleneck Forced a Deep Rethink of Rust's Ownership Model

Introduction

In the long arc of performance engineering, few moments are as decisive as the one where a team realizes that a seemingly modest optimization—a memcpy here, a buffer registration there—is not merely suboptimal but fundamentally wrong for the problem at hand. Message [msg 3070] in the opencode session marks precisely such a turning point. It is the message where the assistant, having spent the previous chunk of Segment 22 instrumenting GPU code and tracing the root cause of severe GPU underutilization, commits to a radical architectural change: direct synthesis of constraint vectors into CUDA-pinned memory, eliminating the Host-to-Device (H2D) transfer bottleneck at its source.

This message is remarkable not because it implements the solution—that comes in the subsequent messages of the chunk—but because it contains the reasoning that chooses the solution. The assistant's "Agent Reasoning" block is a window into a sophisticated design process: weighing multiple approaches, identifying hidden constraints, reasoning about Rust's ownership model, and ultimately converging on a design that is both elegant and safe. The message also reads a key file to understand the CUDA FFI surface, marking the transition from design to implementation.

To understand why this message matters, one must understand the problem it solves. The cuzk proving daemon was experiencing GPU utilization of only ~50%. After extensive instrumentation—first in Rust with GPU_TIMING and FIN_TIMING macros, then in C++ with timing around GPU mutex acquisition and NTT phases—the team had traced the bottleneck to the execute_ntts_single function. Inside that function, the a/b/c synthesis vectors were being transferred from host to device at 1–4 GB/s instead of the PCIe Gen5 x16 line rate of ~50 GB/s. The root cause: the vectors were allocated as standard heap memory, forcing CUDA to stage transfers through a tiny pinned bounce buffer. The SRS points used in MSM operations, by contrast, benefited from direct DMA because they were allocated with cudaHostAlloc. The fix was clear: make the a/b/c vectors pinned too.

But the fix was anything but simple.

The Fork in the Road: Option A vs. Option B

The immediate predecessor to [msg 3070] is [msg 3069], where the assistant presents the user with a binary choice. Option A is a staging memcpy approach: allocate pinned buffers from a pool, memcpy the a/b/c data into them after synthesis completes, then pass those pinned pointers to the GPU. The overhead is estimated at ~80ms per partition for the memcpy—noticeable but small relative to the 2–9 second H2D transfer it replaces. Option A requires no changes to bellperson's ProvingAssignment struct; it is a surgical intervention at the gpu_prove_start boundary.

Option B is zero-copy: synthesize directly into pinned memory from the start. This eliminates the memcpy entirely but requires changing ProvingAssignment's internal vector types from Vec<Scalar> to something backed by pinned memory. It is a deeper, riskier change that ripples through the codebase.

The user's response, embedded in the conversation data of [msg 3069], chooses Option B. Their reasoning is sharp: the memcpy wouldn't just be 80ms, because under heavy memory pressure—20+ synthesis threads contending for bandwidth across 400+ GiB of RAM—host memory bandwidth is already contested. A staged memcpy would itself be slow. Moreover, the current Vec allocation via malloc already does implicit copies during reallocation (the push → grow → memcpy cycle). Pre-sized pinned buffers eliminate both the reallocation copies and the H2D staging.

This decision sets the stage for [msg 3070], where the assistant internalizes the user's reasoning and begins the deep design work.

The Agent's Reasoning: A Masterclass in Constraint-Driven Design

The "Agent Reasoning" block in [msg 3070] is the heart of the message. It is a stream-of-consciousness design document, roughly 1,500 words of iterative refinement. Let us walk through it in detail, because the thinking process itself is the story.

Phase 1: Accepting the Premise

The reasoning opens with the assistant absorbing the user's point:

The user wants Option B — direct synthesis into pinned memory. And they make a good point: the memcpy wouldn't just be 80ms, because under heavy memory pressure (20+ synthesis threads, 400+ GiB used), host memory bandwidth is already contested. The current Vec allocation via malloc also does implicit copies during reallocation (push → grow → memcpy to new allocation). With pre-sized pinned buffers, we eliminate both the reallocation copies AND the H2D staging.

This is the moment of commitment. The assistant has been weighing options across multiple messages ([msg 3062], [msg 3064]–[msg 3068]), reading code, checking CUDA documentation, and considering trade-offs. Now the user has spoken, and the assistant pivots fully to Option B.

The initial plan is stated in four bullet points:

  1. Create a PinnedPool in cuzk-core/memory.rs that manages cudaHostAlloc'd buffers
  2. Build a PinnedVec<T> wrapper that pulls from the pool, supports push and pointer extraction, and returns its buffer on drop
  3. Swap ProvingAssignment.a/b/c from regular Vecs to use pinned backing
  4. Wire the pool into MemoryBudget This seems straightforward, but the reasoning immediately reveals the complications.

Phase 2: The Dependency Problem

The first obstacle is architectural. ProvingAssignment lives in bellperson, which is a fork of the bellman proving library. It is generic over Scalar and has no knowledge of CUDA. The PinnedPool lives in cuzk-core, which is the proving engine's core library and does have CUDA access. How do you bridge this gap without creating a circular dependency or pulling CUDA into bellperson?

The assistant considers several options:

The cleanest minimal-change approach is to create a PinnedVec<T> type that manages raw pointers and capacity from a global pool, then either make ProvingAssignment generic over the storage type or use a type alias to swap implementations without touching the rest of the code.

But this still requires bellperson to know about the pool. The assistant realizes that bellperson can't depend on CUDA directly, so the allocation mechanism must be abstracted—perhaps through function pointers registered at startup by cuzk-core.

Phase 3: The Undefined Behavior Trap

The reasoning then takes a sharp turn into Rust's ownership model. The assistant considers using Vec::from_raw_parts to create a Vec that points to pre-allocated pinned memory:

I can work around this by creating a Vec directly from the pinned memory using Vec::from_raw_parts, which lets me point the vector at pre-allocated pinned memory without going through the normal allocator.

>

But there's a catch — when that Vec gets dropped, Rust will try to deallocate through the global allocator, causing undefined behavior.

This is the critical moment. The assistant has identified that Rust's Vec type assumes it owns its backing buffer and will call dealloc on drop. If the backing buffer is actually a cudaHostAlloc'd region, the global allocator's dealloc will try to free it, causing undefined behavior—likely a crash or memory corruption.

The assistant considers several mitigations:

Phase 4: The ManuallyDrop Breakthrough

After cycling through several dead ends, the assistant arrives at the key insight:

Actually, I'm reconsidering the ScalarVec approach—there's a simpler path using ManuallyDrop<Vec<T>> in ProvingAssignment to wrap vectors allocated from pinned memory, which avoids changing the field types and lets us control when deallocation happens.

The ManuallyDrop wrapper prevents Rust from automatically dropping the Vec. Instead, the assistant can manually extract the raw pointer, return the buffer to the pool, and then forget the Vec entirely. The Vec's Drop is never called, so the global allocator never tries to free the pinned memory.

But this raises a new question: how do you safely extract the contents? The assistant works through the mechanics:

std::mem::take for Vec<T> replaces with Vec::new() (which is Default). So 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.

This is the breakthrough. The approach is:

  1. Add a pinned_backing: Option<PinnedBacking> field to ProvingAssignment
  2. Add a release_abc(&mut self) method that uses std::mem::take to empty the Vecs and then calls the pool's return callback
  3. Implement Drop for ProvingAssignment that calls release_abc() as a safety net
  4. In prove_start, call release_abc() after extracting pointers The PinnedBacking struct holds the raw pointers and a callback function that returns the buffers to the pool. When release_abc is called, it takes ownership of the Vecs (via std::mem::take), forgets them (via ManuallyDrop::take), and invokes the callback to return the pinned buffer to the pool.

Phase 5: Budget Integration

The reasoning then turns to how the PinnedPool integrates with the existing MemoryBudget system. The assistant identifies a subtlety:

I'm realizing the budget reservation system doesn't map cleanly to the pinned buffer lifetime—the reservation guard is temporary but the buffer persists. Instead, the pool should hold permanent reservations for allocated buffers, releasing them only when buffers are freed through shrinkage or eviction.

The MemoryBudget system uses a reservation model: you acquire a guard that holds a certain amount of budget, and when the guard is dropped, the budget is released. But pinned buffers persist across synthesis and GPU proving phases—they are checked out at synthesis time and checked back in after GPU transfer completes. A temporary reservation guard would be dropped at the wrong time.

The solution is to have the pool hold permanent reservations for its allocated buffers. When a new buffer is allocated, the pool calls try_acquire on the budget to reserve the memory. When the buffer is freed (returned to the pool and eventually evicted), the pool calls release_internal to release the budget. Checkout and checkin of buffers from the pool do not touch the budget at all—the reservation is already held by the pool.

This is a clean separation of concerns: the pool manages the budget for its buffers, and the pipeline code simply checks buffers in and out.

Phase 6: Size Class Strategy

The final design consideration is how to size the buffers. The pool needs to handle different workloads:

The Memory Safety Challenge: Why This Is Hard

To fully appreciate the design in [msg 3070], one must understand why Rust's ownership model makes pinned memory integration difficult. The core tension is between Rust's safety guarantees and CUDA's memory model.

In Rust, Vec<T> owns its backing buffer. When a Vec is dropped, it calls dealloc on the allocator that created it. This is safe because the allocator and the deallocator are always matched—the global allocator that allocated the buffer is the same one that frees it.

CUDA-pinned memory breaks this assumption. cudaHostAlloc allocates memory that is page-locked and accessible for DMA transfers. This memory must be freed with cudaFreeHost, not with the global allocator's free. If a Vec pointing to pinned memory is dropped normally, Rust will call the global allocator's dealloc, which will attempt to free a CUDA allocation—undefined behavior.

The standard workaround is to use a custom allocator that implements Rust's Allocator trait. But this trait is only available on nightly Rust, and the project targets stable Rust. The assistant therefore must use unsafe code to manually manage the lifecycle.

The ManuallyDrop<Vec<T>> approach is elegant because it prevents the automatic drop while still allowing the Vec API to be used during synthesis. The release_abc method is the controlled teardown path that correctly returns the buffer to the pool. The custom Drop implementation on ProvingAssignment is a safety net: if release_abc was not called explicitly (e.g., due to an early return or error), the Drop impl calls it automatically.

But this safety net has a subtle requirement: after release_abc empties the Vecs via std::mem::take, the Vec fields are empty Vec::new(). When the struct is dropped, Vec::drop on an empty Vec is a no-op—it checks self.buf.ptr.is_null() and returns immediately. So the safety net is truly safe: either the buffers were returned to the pool explicitly, or they were emptied and the empty Vecs are dropped harmlessly.

The CUDA FFI Surface

The second part of [msg 3070] is a file read of /tmp/czk/extern/supraseal-c2/src/lib.rs. This is the assistant's first concrete implementation step: understanding what CUDA FFI bindings are already available.

The file reveals that supraseal-c2 already has:

Assumptions and Knowledge

The reasoning in [msg 3070] makes several assumptions that are worth examining:

Assumption 1: The capacity hint system provides exact sizes before synthesis. The assistant assumes that after the first synthesis run, the system knows exactly how large the a/b/c vectors will be. This is based on the capacity hint cache described in [msg 3057], which caches num_constraints, num_aux, and num_inputs after the first synthesis for a circuit type. This assumption is well-founded but has a failure mode: if the circuit structure changes between runs (e.g., different circuit parameters), the cached sizes could be wrong, leading to either wasted space or buffer overflow.

Assumption 2: Fr (the scalar field element) is Copy. The assistant notes that for scalar types like Fr that are Copy, forgetting the Vec's contents is safe because there are no destructors to run. This is correct for the BLS12-381 scalar field used in Filecoin proofs, but it means the design is specific to scalar types—it would not work for types with custom Drop implementations.

Assumption 3: The global allocator's dealloc on an empty Vec is a no-op. This is correct for standard Rust allocators but could theoretically differ with custom global allocators. The project uses the system allocator, so this is safe.

Assumption 4: cudaHostAlloc is the right API. The assistant considered cudaHostRegister (which registers existing heap memory as pinned) but rejected it because registration is expensive. This decision is validated by the CUDA benchmarks the assistant read in [msg 3061]. However, cudaHostAlloc has its own costs (~60–200μs per allocation) and requires that the allocation size be known upfront. The pool amortizes this cost by reusing buffers.

Assumption 5: The pool can hold permanent budget reservations. The assistant assumes that the MemoryBudget system can support reservations that outlive the immediate scope. Looking at the release_internal method referenced in the reasoning, this appears to be supported—the budget tracks used bytes and notifies waiters when memory is released. The pool would call try_acquire on allocation and release_internal on deallocation, with no intermediate reservation guards.

What This Message Creates

The primary output of [msg 3070] is a design decision and a plan. But more importantly, it creates:

1. A validated architecture. The assistant has traced through the entire lifecycle—from synthesis (where a/b/c are built) through prove_start (where pointers are extracted) to GPU transfer (where H2D copies happen)—and verified that the ManuallyDrop approach works at each stage.

2. A safety strategy. The combination of ManuallyDrop, release_abc, and a custom Drop impl creates a layered safety net that prevents undefined behavior even in error paths.

3. A budget integration model. The pool-as-permanent-reservation-holder model cleanly separates buffer lifecycle from budget accounting.

4. An implementation roadmap. The assistant now knows exactly what files to create and modify:

What This Message Requires

To fully understand [msg 3070], a reader needs:

Knowledge of the problem domain. The reader must understand that GPU proving involves transferring large vectors (a/b/c) from host memory to device memory, and that this transfer is bandwidth-limited. They must understand why pinned memory is faster (direct DMA vs. staged through bounce buffer).

Knowledge of CUDA memory APIs. The distinction between cudaHostAlloc (allocates pinned memory) and cudaHostRegister (pins existing memory) is central to the design. The performance characteristics of each (registration cost, allocation cost) inform the pool-based approach.

Knowledge of Rust's ownership model. The Vec drop issue, ManuallyDrop, std::mem::take, and Vec::from_raw_parts are all Rust-specific concepts. A reader unfamiliar with Rust's allocator model would not understand why the "obvious" approach (just use pinned memory) is fraught with undefined behavior.

Knowledge of the codebase architecture. The distinction between bellperson (generic proving library, no CUDA) and cuzk-core (engine core, CUDA-aware) is essential. The dependency inversion—having bellperson hold a callback to a cuzk-core pool—is a non-trivial architectural pattern.

Knowledge of the prior investigation. The reader should know that the team spent the earlier part of Segment 22 instrumenting GPU code, ruling out lock contention and malloc_trim as causes, and ultimately tracing the bottleneck to the H2D transfer. This context explains why the team is willing to make such a deep change.

Mistakes and Risks

The reasoning in [msg 3070] is thorough, but it contains some potential blind spots:

The Drop safety net may not be reachable. The assistant implements Drop for ProvingAssignment as a safety net, calling release_abc if the buffers haven't been released yet. But in the normal flow, prove_start takes ownership of the ProvingAssignment (it receives Vec<ProvingAssignment>), extracts the pointers, and then... what? If prove_start drops the ProvingAssignment after extracting pointers, the Drop impl would call release_abc, which would try to return already-extracted buffers. The assistant addresses this by having release_abc use std::mem::take to empty the Vecs first, so the second call is a no-op. But this requires careful sequencing.

The pool's free list could fragment. The assistant plans to allocate buffers to exact sizes and return them to a free list. Over time, with mixed SnapDeals and PoRep workloads, the free list could accumulate buffers of different sizes. A checkout request for a 2.59 GiB buffer might find only a 4.17 GiB buffer available, wasting 1.58 GiB. The assistant acknowledges this but does not implement a defragmentation strategy.

The budget integration assumes try_acquire will succeed. If the budget is exhausted when a new buffer needs to be allocated, try_acquire will fail. The assistant does not discuss what happens in this case—does the pool block and wait? Does it evict an existing buffer? Does synthesis fail? The subsequent implementation (in the rest of the chunk) likely addresses this, but the reasoning in [msg 3070] does not.

The ManuallyDrop approach leaks if release_abc is never called and Drop is not implemented. The assistant does implement Drop, so this is mitigated. But if a future developer removes the Drop impl or the ProvingAssignment is moved in a way that bypasses Drop (e.g., into a ManuallyDrop wrapper itself), the pinned buffer would never be returned to the pool.

The Broader Significance

[msg 3070] is more than just a design document for a pinned memory pool. It is a case study in how performance engineering interacts with language-level safety guarantees. The GPU bottleneck was identified through careful instrumentation; the solution required navigating Rust's ownership model to safely integrate with CUDA's memory management.

The message also illustrates the value of the opencode session format. The assistant's reasoning is not hidden in a design doc or a private notebook—it is visible, quotable, and reviewable. The user can see the assistant working through the problem, catch errors, and provide guidance (as they did in choosing Option B). This transparency is a significant advantage over traditional AI interaction models where the reasoning is opaque.

For the project, this message marks the transition from investigation to implementation. The team has spent the early part of Segment 22 instrumenting code, collecting timing data, and tracing the bottleneck. With [msg 3070], they commit to a solution and begin building. The subsequent messages in the chunk will implement the PinnedPool, modify ProvingAssignment, and validate the changes with cargo check.

Conclusion

Message [msg 3070] is the hinge point of Segment 22's investigation into GPU underutilization. It is where the team stops asking "why" and starts asking "how." The assistant's reasoning navigates a complex design space—architectural boundaries between crates, Rust's ownership semantics, CUDA memory APIs, budget accounting, and size class strategy—to converge on a solution that is both performant and safe.

The design is not perfect. It makes assumptions about circuit stability, scalar types, and budget availability that could fail in edge cases. But it is a reasoned, defensible approach to a hard problem. The ManuallyDrop<Vec<T>> trick is particularly elegant: it uses Rust's type system to prevent undefined behavior while preserving the familiar Vec API for synthesis code.

When the implementation is complete—when the PinnedPool is wired into the engine, when synthesis writes directly into pinned buffers, when the H2D transfer collapses from seconds to milliseconds—it will be because of the thinking captured in this single message. The code that follows is execution; the design is the message.