The Last Stitch: How a Single Edit Completed the Zero-Copy Pinned Memory Pool Wiring

[assistant] [edit] /tmp/czk/extern/cuzk/cuzk-core/src/pipeline.rs Edit applied successfully.

At first glance, this message appears to be the most mundane of confirmations — a tool call succeeded, a file was modified, the system acknowledged. But in the context of the broader coding session, this single edit represents the culmination of a painstaking, multi-hour effort to solve one of the most vexing performance problems in GPU-accelerated proving: the host-to-device (H2D) memory transfer bottleneck. This message is the final stitch in a surgical refactoring that threaded a zero-copy pinned memory pool through hundreds of lines of production Rust code, touching files across three separate crates and dozens of function signatures.

The Problem That Drove the Work

To understand why this message exists, one must understand the crisis that preceded it. The cuzk proving daemon — a high-performance GPU-accelerated zero-knowledge proof system — was suffering from persistent GPU underutilization hovering around 50%. Through detailed C++ timing instrumentation (CUZK_TIMING, CUZK_NTT_H), the team had conclusively identified the root cause: the GPU mutex was held for 1.6–7.0 seconds per partition, but the GPU was only actively computing for approximately 1.2 seconds. The remaining time was consumed by ntt_kernels — specifically, the H2D transfer of a/b/c vectors from unpinned Rust heap memory (Vec<Scalar>) via cudaMemcpyAsync.

The numbers told a stark story. CUDA's default behavior when copying from unpinned (pageable) host memory is to stage through a tiny internal pinned bounce buffer, achieving transfer rates of only 1–4 GB/s. The PCIe Gen5 link between host and GPU is capable of approximately 50 GB/s. The bottleneck was not the GPU's compute capacity, nor the CPU's synthesis throughput — it was the memory copy path. The solution was to allocate pinned (page-locked) host memory that CUDA could DMA directly to the device, bypassing the bounce buffer entirely.

The Architecture of the Fix

The team designed a PinnedPool — a reusable, budget-aware memory pool that pre-allocates pinned buffers and hands them out to synthesis workers on demand. The pool integrates with the existing MemoryBudget system so that pinned allocations count against the same memory cap as all other allocations, preventing OOM scenarios. When memory pressure triggers the evictor callback, the pool can shrink() to release idle pinned buffers back to the operating system.

The core data structures — PinnedPool, PinnedBacking, release_abc(), and new_with_pinned() — were implemented in a previous session and compiled cleanly. But a data structure is useless if it isn't wired into the execution paths. The work captured in the messages preceding this one (messages 3149 through 3192) was the exhaustive, systematic threading of the Arc<PinnedPool> reference through the entire proving pipeline.

The Wiring Journey

The assistant approached the wiring with methodical precision. Starting from the Engine::new() constructor, it created the Arc<PinnedPool> and stored it on the Engine struct. From there, the pool needed to reach every code path that performs synthesis — because synthesis is where the a/b/c vectors are allocated, and those allocations are what must be pinned.

The chain of propagation is worth tracing because it reveals the architecture of the cuzk engine:

1. Engine startup — Engine::new() creates the pool and stores it as self.pinned_pool.
2. Evictor callback — The pool is cloned into the evictor closure so that memory pressure can trigger pool shrinkage.
3. Dispatcher setup — The dispatcher closure clones self.pinned_pool into its scope.
4. dispatch_batch() — The function signature gains a pinned_pool: &Arc<PinnedPool> parameter.
5. process_batch() — The function signature similarly gains the parameter.
6. PartitionWorkItem — The struct gains an Option<Arc<PinnedPool>> field.
7. Synthesis worker — The worker closure extracts item.pinned_pool.as_ref() and passes it to synthesis functions.
8. synthesize_partition() — The function signature is updated to accept Option<&Arc<PinnedPool>>.
9. synthesize_circuits_batch_with_prover_factory() — A new function is created that checks out pinned buffers and creates ProvingAssignment instances with pinned backing. Each step required updating not just the function declaration but every call site. The assistant used grep to find all occurrences, read each one to understand the context, and applied edits to thread the parameter through. When the first compilation attempt failed due to a missing IndexedParallelIterator import in supraseal.rs, the assistant fixed that too.

The Missed Call Site

After updating all five dispatch_batch() call sites and verifying six occurrences of &pinned_pool, in engine.rs, the assistant ran cargo check again. The compilation succeeded — or so it seemed. But then, in message 3192, the assistant realized:

"There's another call site to synthesize_partition I missed — line 2491 in the old partitioned proving path."

This is the moment our subject message becomes significant. Line 2491 of pipeline.rs contained a direct call to synthesize_partition() in what the assistant called the "old partitioned proving path" — a legacy code path that was not part of the newer pipeline architecture. This call site was not caught by the earlier grep searches because those searches focused on dispatch_batch and process_batch — the newer pipeline entry points. The old path bypasses those functions entirely, calling synthesize_partition directly.

The assistant read the surrounding code to understand the context. The call was inside a scope.spawn() closure:

scope.spawn(move || -> Result<Duration> {
    let part_job_id = format!("{}-part-{}", job_id_owned, partition_idx);
    let synth = synthesize_partition(parsed_ref, partition_idx, &part_job_id)?;
    let duration = synth.synthesis_duration;
    ...
});

The synthesize_partition function had been updated to accept an Option<&Arc<PinnedPool>> parameter, but this call site was still passing only three arguments. The compiler would reject this code with a type error. The fix was to add None as the fourth argument — the fallback path for when the pinned pool is not available in this legacy context.

What the Edit Actually Changed

The edit applied in message 3193 changed line 2491 from:

let synth = synthesize_partition(parsed_ref, partition_idx, &part_job_id)?;

to something like:

let synth = synthesize_partition(parsed_ref, partition_idx, &part_job_id, None)?;

The None signals that this code path does not have access to the pinned pool and will fall back to standard heap allocations. This is the graceful degradation path — the old proving path continues to work correctly, just without the performance benefit of pinned memory. The pool is optional by design; the system degrades gracefully when the pool is exhausted or unavailable.

Why This Matters

This edit is small — a single parameter added to a single function call. But its significance lies in what it represents: the completion of a complete data flow. Every code path that performs synthesis now has the opportunity to use pinned memory. The old path gets None (graceful fallback), while the new pipeline paths get the actual pool reference. The system is now consistent: there is no code path that will fail to compile because of the new parameter, and no code path that will silently use unpinned memory when pinned memory is available.

The assistant's todo list reflected this completion. The first two items — "Wire PinnedPool into pipeline.rs" and "Wire PinnedPool into engine.rs" — were marked as completed. The third item — "Verify cargo check --features cuda-supraseal passes clean" — was in progress. After this edit, the next compilation attempt would succeed.

The Verification That Followed

The assistant's next action after this edit would be to run cargo check --features cuda-supraseal again. The previous attempt (message 3191) had shown only warnings about an unrelated cfg condition — no errors. With this last call site fixed, the compilation would pass cleanly, confirming that every function signature, every call site, every type conversion was correct.

The successful compilation would then enable the next phase: building a Docker image (cuzk-rebuild:pinned1) and deploying the instrumented binary to the remote test machine for performance validation. The primary metric would be the ntt_kernels time in the CUZK timing logs, which should drop from the 2–9 second range to under 100 milliseconds — effectively eliminating the H2D bottleneck and raising GPU utilization toward 100%.

Assumptions and Decisions

The assistant made several assumptions during this work. First, it assumed that the synthesize_partition function in the old path could safely receive None for the pool — that the function handles the None case correctly by falling back to standard Vec::<Scalar> allocations. This assumption was validated by the earlier implementation of the fallback path in synthesize_circuits_batch_with_prover_factory.

Second, the assistant assumed that the old partitioned proving path was still in active use and needed to compile. If this code path had been dead code slated for removal, the edit would have been unnecessary. But the assistant correctly treated it as live code that must be updated.

Third, the assistant assumed that passing None was semantically correct — that the old path should not attempt to use pinned memory. This is a reasonable assumption because the old path does not have access to the PinnedPool (it runs in a different scope without the pool reference), and introducing the pool would require additional refactoring of the legacy code.

Input Knowledge Required

To understand this message, one needs knowledge of:

• The CUDA pinned memory model and why cudaMemcpyAsync from pageable memory is slow
• The architecture of the cuzk proving engine, including the dispatcher, batch collector, synthesis workers, and GPU workers
• The PinnedPool data structure and its integration with MemoryBudget
• The distinction between the new pipeline-based proving path and the old partitioned proving path
• Rust's Arc for shared ownership and Option for optional parameters
• The synthesize_partition function signature and how it was extended

Output Knowledge Created

This message created:

• Confirmation that the PinnedPool wiring is complete across all code paths
• A fixed compilation error that would have blocked the Docker build
• A consistent interface where all synthesize_partition call sites pass the pool parameter
• The foundation for the performance validation deployment

The Broader Lesson

There is a lesson here about the nature of systems programming. The dramatic performance gains — from 1–4 GB/s to 50 GB/s — are enabled not by any single change but by the exhaustive, unglamorous work of threading a reference through every code path that allocates memory. The assistant's systematic approach — grep for all call sites, read each one, apply edits, verify with grep counts, compile, find the missed site, fix it — is the discipline required to make zero-copy work. There are no shortcuts. Every call site must be found. Every path must be updated. The last stitch is as important as the first.