The Pivot Point: Bridging Memory Management and Proof Synthesis in the CuZK Pinned Memory Pool

Message Overview

In message [msg 3083], the assistant performs a brief but architecturally significant action: reading lib.rs to register the newly written PinnedPool module and signaling the transition to modifying bellperson's ProvingAssignment with a PinnedBacking struct. The full message reads:

[assistant] Now register the module in lib.rs and add PinnedBacking to bellperson: [read] /tmp/czk/extern/cuzk/cuzk-core/src/lib.rs

This is a transitional message — the bridge between two major implementation phases. On one side lies the freshly written PinnedPool in cuzk-core/src/pinned_pool.rs, a CUDA-pinned memory allocator integrated with the budget-based memory management system. On the other side lies the more delicate task: modifying bellperson's ProvingAssignment to safely hold and release pinned memory buffers without triggering undefined behavior from Rust's standard Vec deallocation. The message is short, but it marks a critical architectural pivot point in the implementation of a zero-copy pipeline designed to eliminate a severe GPU underutilization bottleneck.

The Problem That Drives This Work

The broader context is a weeks-long investigation into GPU underutilization in the CuZK proving pipeline. Precise timing instrumentation had revealed that the GPU was idle roughly 50% of the time during proof generation. The root cause, traced through layers of Rust-side timing (GPU_TIMING, FIN_TIMING) and C++ instrumentation of gpu_prove_start, was 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 Vec<Fr> memory at 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 malloc allocations are not page-aligned for direct memory access (DMA) by the GPU. CUDA's cudaMemcpyAsync must stage transfers through a small pinned bounce buffer, creating a severe bandwidth bottleneck. Meanwhile, the SRS points used in MSM operations benefited from direct DMA via cudaHostAlloc, confirming that pinned memory was the correct solution.

The chosen approach — Option B, direct synthesis into pinned memory — eliminates both the reallocation copies that occur during Vec::push growth and the staged H2D transfer. Under heavy memory pressure from 20+ concurrent synthesis threads using 400+ GiB of host memory, even a staged memcpy would be severely bottlenecked by contested memory bandwidth. Direct synthesis into pre-allocated pinned buffers collapses the H2D transfer from seconds to milliseconds.

Why This Message Exists

Message [msg 3083] exists because the assistant has just completed the first major piece of the implementation — the PinnedPool struct in pinned_pool.rs — and now faces a dependency problem. The PinnedPool lives in cuzk-core, which has access to CUDA via the supraseal-c2 crate. But the ProvingAssignment struct that needs to use pinned buffers lives in bellperson, a separate crate in the dependency tree that does not and should not depend on CUDA directly.

The assistant's first action in this message — reading lib.rs — is the mechanical step of adding pub mod pinned_pool; to the crate's module registry. This is straightforward but necessary: without it, the PinnedPool type is invisible to the rest of cuzk-core and to downstream consumers. The second, implied action — "add PinnedBacking to bellperson" — is the architectural challenge.

The PinnedBacking struct is the bridge. It lives in bellperson (which has no CUDA dependency) and holds raw pointers and a return callback function. The callback is registered at startup by the higher-level cuzk-core layer, which has access to CUDA and the PinnedPool. This separation of concerns keeps bellperson clean of GPU-specific dependencies while still allowing it to participate in the zero-copy pipeline.

Architectural Decisions Visible in This Message

Several design choices are crystallized in this transitional moment:

The ManuallyDrop pattern. The assistant has decided that ProvingAssignment.a/b/c will remain Vec<Scalar> fields, not be replaced with a custom PinnedVec type. When pinned backing is provided, these Vecs are constructed via Vec::from_raw_parts(pinned_ptr, 0, capacity). The critical insight is that dropping a Vec constructed this way would call the global allocator's dealloc on the pinned pointer, causing undefined behavior. The solution is a release_abc() method that uses std::mem::take and ManuallyDrop::take to "forget" the Vec's contents, preventing deallocation, before returning the buffer to the pool via the callback.

The Drop safety net. A custom Drop implementation on ProvingAssignment calls release_abc() if the pinned backing hasn't been explicitly released. This ensures that even if the struct is destroyed through an unexpected path (e.g., an error during synthesis), the pinned buffers are properly returned to the pool rather than leaked or deallocated incorrectly.

Exact-size allocation. The PinnedPool allocates buffers to exact requested sizes rather than using a single max-size pool. This naturally handles mixed workloads — SnapDeals requires approximately 2.59 GiB per buffer while PoRep needs 4.17 GiB — without wasting memory on over-sized allocations.

Budget integration. The pool holds permanent reservations for allocated buffers through the MemoryBudget system. When a new buffer is allocated, try_acquire is called to reserve the budget; when the buffer is freed, release_internal returns the budget. This integrates seamlessly with the existing memory pressure management.

The Thinking Process

The assistant's reasoning is visible in the sequence of actions leading to and following this message. In [msg 3077], the assistant 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

This architecture reveals a clear layered design: the pool manages raw memory, the backing struct bridges the crate boundary, the assignment uses the backing, and the pipeline and engine wire everything together. Message [msg 3083] is the moment where the assistant moves from layer 1 (PinnedPool) to layer 2 (PinnedBacking).

The assistant's earlier reasoning in [msg 3070] shows the deliberation behind this design. The assistant considered using a custom PinnedVec<T> type, an enum wrapper over Vec or PinnedVec, or a thread-local pool accessible globally. Each option was evaluated against criteria: minimal code change, no runtime branching in hot paths (where push is called tens of millions of times per partition), and no undefined behavior on drop. The chosen approach — ManuallyDrop<Vec<T>> with a callback — emerged as the cleanest minimal-change solution.

Assumptions and Risks

The approach makes several assumptions worth examining:

That Fr is Copy and has 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 Box<dyn FnOnce(...)> 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<Vec<...>> free list. Under heavy concurrent synthesis (20+ threads), contention on this mutex could become a bottleneck. The assistant does not address this concern in the visible reasoning, though the pool's design minimizes checkout/checkin frequency (once per partition, not per element).

Input Knowledge Required

To understand this message, one needs:

The GPU utilization investigation: That the H2D transfer of a/b/c vectors was identified as the bottleneck, with measured bandwidth of 1-4 GB/s versus the PCIe Gen5 line rate of ~50 GB/s.
The CUDA memory model: That cudaHostAlloc produces pinned memory suitable for direct DMA, while standard malloc allocations require staged transfers through a bounce buffer.
The crate dependency graph: That cuzk-core depends on CUDA via supraseal-c2, while bellperson is a pure Rust crate with no GPU dependencies.
Rust's allocator model: That Vec::from_raw_parts creates a Vec whose deallocator is the global allocator, and that dropping such a Vec with a non-standard pointer causes undefined behavior.
The MemoryBudget system: The existing budget-based memory manager with try_acquire, release_internal, and the two-phase release pattern for partition synthesis.

Output Knowledge Created

This message produces:

Module registration: The pub mod pinned_pool; line added to cuzk-core/src/lib.rs, making the PinnedPool type part of the crate's public API.
The PinnedBacking struct: A bridge type in bellperson that holds raw pointers and a return callback, enabling CUDA-free code to participate in the pinned memory pipeline.
The release_abc() method: A safe mechanism for extracting and returning pinned buffers from ProvingAssignment without undefined behavior.
The new_with_pinned() constructor: An alternative constructor that accepts pre-allocated pinned buffers for direct synthesis.
The Drop safety net: A custom Drop implementation that calls release_abc() as a fallback, preventing leaks from unexpected destruction paths.
The prove_start update: Modification to the supraseal integration point to call release_abc() after extracting pointers, ensuring buffers are returned to the pool promptly.

Conclusion

Message [msg 3083] is brief in content but pivotal in structure. It marks the transition from implementing the memory management infrastructure (the PinnedPool) to modifying the consumer (bellperson's ProvingAssignment). The assistant's methodical approach — write, register, bridge, integrate — reflects a clear architectural vision for solving one of the most stubborn performance problems in the CuZK proving pipeline. The zero-copy pinned memory solution promises to collapse the H2D transfer bottleneck from seconds to milliseconds, dramatically improving GPU utilization and proof throughput. This message, unassuming as it appears, is the hinge on which that optimization turns.