The Moment of Restructure: A Deep Dive into Phase 10's Two-Lock Design for Groth16 GPU Proving

Introduction

In the long arc of optimizing the SUPRASEAL_C2 Groth16 proof generation pipeline for Filecoin's Proof-of-Replication (PoRep), few messages capture the tension between elegant software abstraction and unforgiving hardware reality as vividly as message 2591. This message, nestled in the twenty-seventh segment of an extensive coding session, represents the precise moment when the assistant transitions from design to implementation of Phase 10's two-lock architecture. It is a message about restructuring code—about taking a known, working single-mutex design and splitting it into two mutexes to unlock overlapping execution. But it is also a message that foreshadows a fundamental collision between the assistant's logical model of concurrency and the physical reality of CUDA device-global synchronization, a collision that will unfold catastrophically in the subsequent messages of this chunk.

The Message in Context

To understand message 2591, we must understand the optimization journey that led to it. The session had been tracking a bottleneck migration across multiple phases. Phase 9 had implemented PCIe transfer optimization, achieving a 14.2% throughput improvement in single-worker mode, but dual-worker mode revealed PCIe bandwidth contention as the next bottleneck ([msg 2576]). Detailed timing analysis showed that the CPU critical path—specifically prep_msm at ~1.9 seconds and b_g2_msm at ~0.48 seconds—now dominated the per-partition wall time at ~2.4 seconds, leaving the GPU idle for ~600ms per partition waiting for the CPU thread ([chunk 27.0]).

The user had proposed a two-lock design to better overlap CPU and GPU work. The assistant, in the messages immediately preceding 2591, had analyzed deadlock risks and designed a simplified protocol using a mem_mtx (for VRAM allocation and upload) and a compute_mtx (for GPU kernel execution). The key insight was that since d_a and d_bc—the two large device buffers totaling 12 GiB—were freed synchronously via cudaFree inside compute_mtx, the next worker could safely allocate in mem_mtx without needing a cooperative free slot. The design document (c2-optimization-proposal-10.md) had been written ([msg 2587]), and the gpu_locks struct with its two mutexes had been defined ([msg 2589]). The fallback mutex handling had been updated ([msg 2590]).

Message 2591 is the bridge: the assistant has the struct defined, the FFI boundary understood, and now needs to perform the actual surgical restructure of the 300+ line generate_groth16_proofs_c function.

What the Message Actually Says

The message opens with a clear, methodical mapping of the current code structure:

Now the big restructure. Currently the code is: 1. Lines ~610-636: cudaHostRegister (before lock) 2. Line 641: std::unique_lock<std::mutex> gpu_lock(*mtx_ptr); 3. Lines 643-777: Pre-staging (alloc + upload) under gpu_lock 4. Lines 779-951: per_gpu kernels under gpu_lock 5. Lines 958-970: per_gpu join + cleanup under gpu_lock 6. Line 975: gpu_lock.unlock() 7. Lines 977-984: hostUnregister + prep_msm join (no lock)

This is a numbered inventory of the existing control flow, each line corresponding to a region of code that must be surgically separated. The assistant then presents the target structure:

I need to change it to: 1. Lines ~610-636: cudaHostRegister (before any lock) — unchanged 2. NEW: mem_lock scope for pre-staging 3. NEW: compute_lock scope for GPU kernels + cleanup 4. Event/stream destroy, hostUnregister, prep_msm join (no lock)

The symmetry is elegant: the old single lock encompassed both pre-staging (memory management) and kernel execution (compute). The new design splits these into two disjoint lock scopes, with a lock-free region after both for cleanup and CPU-only post-processing.

The message then performs a read operation to load the actual code section that needs replacement, and the conversation data shows the beginning of that section (lines 637-644), which includes error handling for cudaHostRegister.

The Reasoning and Motivation

Why was this message written? The surface answer is straightforward: the assistant needed to read the code it was about to edit. But the deeper answer reveals the assistant's cognitive process. The assistant was about to perform a high-risk surgical edit on a critical function—one that, if done incorrectly, could introduce deadlocks, data races, or memory corruption. Before cutting, the assistant needed a complete mental model of the current code's structure.

The numbered list of "Currently the code is" serves as a cognitive map. Each numbered item corresponds to a contiguous block of code with a specific semantic role: host registration, lock acquisition, pre-staging, kernel execution, join/cleanup, unlock, post-processing. By explicitly enumerating these regions and their line numbers, the assistant is creating a spatial representation of the code that will guide the subsequent edits.

The motivation is rooted in the optimization journey. Phase 9 had proven that the bottleneck had shifted to CPU memory bandwidth contention. The prep_msm thread (CPU-only synthesis work) and the b_g2_msm (CPU-side point arithmetic) now dominated the critical path. The single-mutex design meant that while one worker held the GPU lock, no other worker could even begin its pre-staging—even though pre-staging only took ~18ms and didn't need the GPU kernels to be idle. The two-lock design aimed to allow worker B to allocate VRAM and upload data (under mem_mtx) while worker A was still running kernels (under compute_mtx), effectively hiding the CPU overhead behind GPU execution.

Assumptions Embedded in the Design

The message, and the design it implements, rests on several critical assumptions:

Assumption 1: Memory management and compute can be cleanly separated on a single CUDA device. The entire two-lock architecture depends on the premise that mem_mtx operations (cudaMalloc, cudaMemcpyAsync, cudaEventRecord) do not conflict with compute_mtx operations (kernel launches, cudaFree, stream synchronization). This assumption is reasonable from a software abstraction perspective—different CUDA streams can run concurrently, and memory allocation is logically independent of kernel execution. But it ignores the reality that cudaDeviceSynchronize and cudaMemPoolTrimTo are device-global operations that interact with all streams on the device.

Assumption 2: Pool trim without device synchronization is sufficient to reclaim memory. The design specified cudaMemPoolTrimTo(pool, 0) inside mem_mtx without a preceding cudaDeviceSynchronize, relying on the fact that all stream operations from the previous compute_mtx holder had been synchronized (via gpu.sync() calls within the per-GPU thread). This assumption proved incorrect because cudaFreeAsync (used by internal MSM and batch_add operations) places memory in the stream-ordered pool, and pool trim may not reclaim that memory until the stream has been synchronized.

Assumption 3: Three workers can usefully share a single GPU with 24 GiB VRAM. Each worker needs 12 GiB for pre-staged buffers (d_a at 4 GiB, d_bc at 8 GiB). With three workers, the total VRAM demand is 36 GiB—exceeding the 24 GiB available. The design relied on the fact that only one worker holds compute_mtx at a time, and that worker frees its 12 GiB before releasing the lock, so peak VRAM usage is 12 GiB (compute worker) + 12 GiB (pre-staging worker) = 24 GiB. This is correct in theory, but the assumption that cudaMalloc would reliably succeed when the previous worker's memory had been freed depended on the pool trim behavior.

Assumption 4: The fallback path (non-prestaged) inside compute_mtx would work correctly when pre-staging fails. This was a safety net: if mem_mtx allocation failed, the code would fall back to the original allocation path inside compute_mtx. The assistant assumed this fallback would be reliable because it was the same path used in Phase 8. But the fallback path itself depended on cudaMalloc and cudaMallocAsync calls that could fail if the stream-ordered pool hadn't been trimmed.

The Thinking Process Visible in the Message

The message reveals a remarkably disciplined engineering thought process. The assistant does not dive directly into editing. Instead, it first builds a complete structural model of the existing code, then maps the target structure onto it. This is evident in the parallel numbered lists: seven items for the current structure, four items for the target structure. The mapping is not one-to-one—the old items 2-5 (single lock region) are split into new items 2-3 (two lock regions), and old items 6-7 (unlock + post-processing) are merged into new item 4 (lock-free cleanup).

The assistant also shows awareness of what should remain unchanged: cudaHostRegister before any lock, and the post-processing after all locks. This is important—the assistant is not blindly restructuring everything but is preserving the boundaries that were already correct.

The message ends with a read operation to load the actual code. This is the transition from planning to execution: the assistant has the mental model, now it needs the raw material to operate on.

Input Knowledge Required

To understand this message, one needs:

The Groth16 proving pipeline: Knowledge that proof generation involves CPU-side synthesis (producing circuit assignments), GPU-side NTT (Number Theoretic Transform), MSM (Multi-Scalar Multiplication), and batch addition. Understanding that d_a, d_bc are large device buffers for the A and B/C wire assignments.
The Phase 8 dual-worker architecture: Understanding that the existing code uses a single std::mutex per GPU, acquired before any CUDA operations and released after all kernels complete. The Phase 8 design proved that two workers could share a GPU by serializing kernel execution while allowing CPU preprocessing to overlap.
CUDA memory management: Understanding cudaMalloc (synchronous, returns immediately with dedicated memory), cudaMallocAsync (stream-ordered, uses a memory pool), cudaFreeAsync (returns memory to pool, not immediately available), cudaMemPoolTrimTo (reclaims cached pool memory), and cudaDeviceSynchronize (blocks until all streams complete).
The FFI boundary: Understanding that the Rust side (cuzk-core) passes an opaque void* pointer to the C++ CUDA code, and that this pointer originally pointed to a std::mutex but now points to a gpu_locks struct containing two mutexes.
The specific hardware constraints: 24 GiB VRAM per GPU, DDR5 memory bandwidth limitations, PCIe Gen4 x16 bandwidth (~32 GB/s), and the ~200 GiB peak memory requirement for the full pipeline.

Output Knowledge Created

This message creates several forms of knowledge:

A structural map of the critical code section: The numbered list serves as documentation of the control flow, making explicit what was implicit in the code. Future readers (including the assistant itself in subsequent messages) can refer to this map when reasoning about the restructured code.
A validated design-to-implementation bridge: The message confirms that the two-lock design from the proposal document can be mapped onto the existing code structure. The mapping is clean: the old single lock region (items 2-5) splits cleanly into two lock regions (items 2-3) with the same internal code.
A boundary definition: The message explicitly identifies what stays outside locks (host registration, event/stream cleanup, host unregister, prep_msm join) and what goes inside each lock. This boundary definition is critical for correctness—putting the wrong operation in the wrong lock region could introduce deadlocks or data races.
A testable hypothesis: The message implicitly defines the expected behavior: that splitting the lock will allow memory operations to overlap with compute operations, reducing per-partition wall time from ~2.4s to ~1.8s. This hypothesis will be tested in subsequent messages.

Mistakes and Incorrect Assumptions

The most significant incorrect assumption in this message is not stated explicitly but is embedded in the structure itself: the assumption that mem_mtx operations can proceed independently of compute_mtx operations on the same CUDA device. The message shows the assistant planning to call cudaMemPoolTrimTo and cudaMemGetInfo inside mem_mtx, and cudaDeviceSynchronize was explicitly removed from the mem region in the design ([msg 2584]). The assistant believed that pool trim alone would suffice because "by the time compute_mtx was released by the previous worker, all stream ops were synced."

This assumption would prove catastrophically wrong. In subsequent messages ([msg 2613], [msg 2618], [msg 2619]), the assistant discovers that cudaDeviceSynchronize inside mem_mtx blocks while another worker holds compute_mtx and runs kernels, effectively serializing the two locks. The cudaMemPoolTrimTo call, being device-global, also interacts with the other worker's streams. The result is that only the first worker successfully pre-stages VRAM; all subsequent workers find insufficient free memory and fall back to the slow path, causing prove time to balloon to 102 seconds.

A second mistake is the assumption that the fallback path would work reliably. The fallback path calls cudaMalloc for 8 GiB buffers inside compute_mtx, but if the stream-ordered pool hasn't been trimmed (because cudaDeviceSynchronize was removed from mem_mtx), this allocation can fail with OOM. The assistant had to add a cudaDeviceSynchronize + pool trim at the start of compute_mtx for the fallback path ([msg 2620]), which further serialized the design.

The Deeper Lesson

Message 2591 is a study in the tension between logical concurrency and physical concurrency. The two-lock design is logically sound: memory allocation is independent of kernel execution, so they should be able to proceed in parallel protected by different locks. But CUDA's device-global operations—cudaDeviceSynchronize, cudaMemPoolTrimTo, even cudaMemGetInfo—create implicit dependencies across all streams on a device. You cannot isolate memory management from compute on a single GPU because the driver's memory subsystem is a shared resource.

The assistant's eventual solution ([msg 2626]) was to remove device-wide synchronization from mem_mtx entirely and rely on the fallback path inside compute_mtx when pre-staging fails. This is a humbling admission: after all the design work, the two-lock architecture could not achieve its intended overlap because the hardware itself resists such clean separation. The lesson is that on a single CUDA device, the only reliable isolation is temporal: one worker at a time, fully.

Conclusion

Message 2591 captures the precise moment of transition from design to implementation in a complex optimization effort. It shows the assistant's disciplined engineering approach: building a structural model of existing code, mapping the target structure onto it, and identifying boundaries. But it also reveals the assumptions that would later prove incorrect—assumptions about the independence of memory management and compute on a single CUDA device. The message is a testament to the value of explicit structural reasoning before surgery, and a cautionary tale about the gap between logical concurrency models and the physical reality of GPU hardware. In the end, the two-lock design would be simplified to a one-lock design with a smarter fallback, and the optimization journey would continue into Phase 11 with a deeper understanding of CUDA's device-global constraints.