The Final Link: Reverting the Shared Mutex Hack and Wiring GPU Index Through the Engine Layer

Subject Message: "Now revert the shared mutex and simplify back to per-GPU mutexes with gpu_ordinal passed as gpu_index:" followed by an edit to /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs.

Introduction

In the architecture of a distributed GPU proving system, the difference between a correct fix and a lazy hack often comes down to whether you solve the problem at the right layer of abstraction. Message <msg id=489> captures the precise moment when an engineering team made exactly that choice. The message is deceptively brief — a single line of narration followed by a file edit — but it represents the culmination of a multi-layered refactoring effort spanning C++, Rust FFI, bellperson prover logic, pipeline orchestration, and finally the engine's GPU worker code. This article examines why this message was written, what decisions it embodies, and what it reveals about the thinking process behind proper multi-GPU fix design.

The Problem: A Shared Mutex Hack

The story begins with a data race. In the CuZK proving engine, multiple Rust worker threads submit Groth16 proofs to a C++ CUDA backend. For multi-circuit proofs (multiple independent proofs bundled together), the C++ code correctly distributes work across available GPUs by setting n_gpus = min(ngpus(), num_circuits) and assigning each circuit to a GPU via select_gpu(tid). But for single-circuit proofs — which is what partitioned proof workloads like PoRep and SnapDeals produce — the C++ code always routed work to GPU 0, regardless of which Rust worker submitted it. On a multi-GPU system, this meant two workers targeting GPU 0 simultaneously would corrupt each other's CUDA state.

The initial "fix" was a shared mutex ([msg 452]). Instead of each GPU having its own mutex (the original design), a single global mutex was introduced that serialized all GPU work onto GPU 0. This was a textbook lazy hack: it prevented the data race by ensuring only one worker could enter the GPU code at a time, but it completely wasted the second GPU. All 16 partitions of a SnapDeals workload would queue up on GPU 0, and the VRAM budget for a single SnapDeals partition was large enough that concurrent kernel execution on the same device caused out-of-memory (OOM) errors on a 20 GB RTX 4000 Ada host.

The Proper Solution: Threading GPU Index

The user articulated the correct architectural vision in [msg 451]: "And the 16 proofs should loadbalance between GPUs sort of obviously." The assistant recognized that the fix needed to be structural, not tactical. The C++ code's select_gpu(tid) call — where tid was always 0 for single-circuit proofs — needed to be replaced with select_gpu(gpu_index) where gpu_index came from the Rust engine's worker assignment.

This required threading a gpu_index parameter through every layer of the call stack, from the C++ entry point up through the Rust FFI, the bellperson prover abstraction, the pipeline orchestration layer, and finally the engine's GPU worker code. The assistant worked bottom-up over approximately 40 messages ([msg 453] through [msg 488]), modifying each layer in turn:

C++ (groth16_cuda.cu): Added int gpu_index parameter to generate_groth16_proofs_start_c and generate_groth16_proofs_c. When gpu_index >= 0, the code forces n_gpus = 1 and uses select_gpu(gpu_index) instead of the old select_gpu(tid) logic.
Rust FFI (supraseal-c2/src/lib.rs): Updated the extern "C" declarations and the public wrapper functions (start_groth16_proof, generate_groth16_proof, generate_groth16_proofs) to accept and forward the gpu_index parameter.
Bellperson (supraseal.rs): Updated prove_start and prove_from_assignments to accept gpu_index and pass it through to the FFI layer.
Pipeline (pipeline.rs): Updated gpu_prove and gpu_prove_start to accept gpu_index and forward it to bellperson.

The Subject Message: Completing the Circuit

Message <msg id=489> is the final piece. The assistant writes:

Now revert the shared mutex and simplify back to per-GPU mutexes with gpu_ordinal passed as gpu_index:

And applies an edit to engine.rs. This is the moment where the entire multi-layer refactoring pays off. The engine's GPU worker code — the code that originally introduced the shared mutex hack — is now rewritten to:

Restore per-GPU mutexes: Instead of a single shared_gpu_mutex that all workers contend for, each GPU gets its own mutex, allowing true concurrent proving across devices.
Pass gpu_ordinal as gpu_index: The Rust worker already knows which GPU it was assigned to (via the gpu_ordinal variable set by the engine's worker assignment logic). This value is now passed as the gpu_index parameter through the entire call chain, ensuring the C++ code selects the correct GPU. The edit is described as "applied successfully," but the significance is immense. The shared mutex hack — which had been a placeholder while the proper plumbing was built — is gone. The per-GPU mutexes, which were the original design intent, are restored. And crucially, the gpu_ordinal value that the Rust engine already computed for worker assignment is now actually used by the C++ proving code, rather than being ignored in favor of a hardcoded GPU 0.

Assumptions and Reasoning

The assistant made several key assumptions in this message:

The lower layers are correct: The edit to engine.rs assumes that the C++ code, FFI layer, bellperson functions, and pipeline functions have all been correctly updated to accept and propagate the gpu_index parameter. If any of those layers had a bug, the engine edit would compile but produce incorrect behavior. The assistant verified this by working bottom-up and testing each layer's changes before moving up.
gpu_ordinal is the right value: The assistant assumes that the gpu_ordinal variable in the engine worker code correctly identifies the target GPU. This is a reasonable assumption because the engine's worker assignment logic already uses gpu_ordinal to set CUDA_VISIBLE_DEVICES for each worker process — it's the same value that determines which GPU the worker should use.
Per-GPU mutexes are sufficient: The assistant assumes that restoring per-GPU mutexes, combined with correct GPU routing, eliminates the data race. This is correct because each worker now acquires the mutex for its assigned GPU, and no two workers target the same GPU unless they're assigned to it (which the engine's load-balancing logic prevents).
The d_a_cache issue is separate: Earlier in the session ([msg 456]), the assistant identified that the global d_a_cache singleton would thrash between GPUs when workers on different GPUs use it concurrently. The assistant correctly recognized this as a separate issue and deferred it, focusing on the GPU routing fix first.

Input Knowledge Required

To understand this message, one needs:

Multi-GPU proving architecture: Knowledge that the CuZK engine uses a worker-per-GPU model where each Rust worker thread is assigned to a specific GPU, acquires that GPU's mutex, and submits proving work to the C++ CUDA backend.
The call chain: Understanding that proving work flows from engine.rs → pipeline.rs → bellperson/supraseal.rs → supraseal-c2/src/lib.rs (Rust FFI) → groth16_cuda.cu (C++ CUDA), and that each layer must forward the GPU index.
The shared mutex hack: Knowledge that a previous fix introduced a single global mutex to prevent data races, at the cost of serializing all GPU work onto GPU 0.
Partitioned proofs: Understanding that PoRep and SnapDeals workloads split proving into multiple independent partitions, each submitted as a single-circuit proof to the GPU backend.

Output Knowledge Created

This message produces:

Correct multi-GPU load balancing: The engine's GPU worker now passes the assigned GPU ordinal to the C++ proving code, ensuring workers on GPU 1 actually prove on GPU 1. This enables the 16 SnapDeals partitions to naturally load-balance across both GPUs.
Restored per-GPU mutex semantics: Each GPU has its own mutex, allowing concurrent proving on different devices. The shared mutex hack is reverted.
A consistent parameter contract: The gpu_index parameter now flows from the top of the stack (engine) to the bottom (C++), with a convention of -1 meaning "auto" (for non-engine callers that don't have a specific GPU assignment).

The Thinking Process

The assistant's reasoning is visible in the progression from [msg 452] to [msg 489]. In [msg 452], the assistant states the plan clearly: "thread a gpu_index parameter through the entire call chain so the C++ code uses select_gpu(gpu_index) instead of select_gpu(0) for single-circuit proofs." This is followed by a todo list that precisely enumerates the layers to modify, bottom-up.

The assistant then executes this plan methodically over the next 37 messages, reading each file, understanding its role in the call chain, making the edit, and updating the todo list. The todo list serves as both a planning document and a progress tracker — each layer is marked "in_progress" while being edited and "completed" when done.

In [msg 488], the assistant reads the current state of engine.rs to understand what needs to change. The shared mutex hack is visible in the code (lines 2149-2155 creating a shared_gpu_mutex). Message [msg 489] then performs the revert, completing the refactoring.

The assistant also demonstrates awareness of edge cases. In [msg 455], it designs the gpu_index parameter with -1 meaning "auto" to preserve backward compatibility for callers that don't have a specific GPU assignment. In [msg 456], it identifies the d_a_cache thrashing issue as a separate concern and defers it, showing disciplined scope management.

Conclusion

Message <msg id=489> is the final keystone in a carefully engineered multi-layer refactoring. It transforms a lazy shared-mutex hack into a proper architectural solution where GPU assignment information flows from the Rust engine's worker scheduler all the way down to the C++ CUDA kernel selection code. The brevity of the message belies the depth of the work it completes — nearly 40 preceding messages of reading, analysis, and modification across five layers of a heterogeneous C++/Rust codebase. The result is a system where 16 SnapDeals partitions can truly load-balance across multiple GPUs, eliminating both the data race and the OOM errors that plagued the temporary fix.