The Bridge Between Layers: Reading the Rust FFI in a Multi-GPU Fix

Introduction

In the architecture of any complex software system, the boundaries between layers are where the most subtle bugs live and where the most careful engineering is required. Nowhere is this truer than in a GPU-accelerated zero-knowledge proving engine, where Rust orchestration code must hand off computation to C++ CUDA kernels across a foreign function interface (FFI). Message [msg 464] captures a pivotal moment in precisely such a boundary crossing: the assistant, having just completed modifications to the C++ GPU proving code, pauses to read the Rust FFI layer that connects the engine's worker threads to the CUDA implementation. This single read tool call — showing lines 195–205 of supraseal-c2/src/lib.rs — is the hinge point between two halves of a critical architectural fix.

The Context: From Lazy Hack to Proper Architecture

To understand message [msg 464], one must first understand the crisis that precipitated it. The conversation's preceding segment had been consumed with a frustrating bug: partitioned PoRep proofs were failing on multi-GPU systems with out-of-memory errors and data races. The initial diagnosis revealed that the C++ GPU proving code always routed single-circuit proofs to GPU 0, regardless of which Rust worker submitted them. The "fix" applied was a shared mutex — a single std::mutex that serialized all partition proofs onto GPU 0, effectively wasting the second GPU entirely.

The user's response at [msg 444] was pointed and correct: "Why is GPU prove for the second GPU not running.. on the second GPU? That's the whole point. CuZK is meant to be a fairly sophisticated proving engine, so it must support multiple GPUs, gpu memory management, and GPU workers are supposed to be interlocking two phases of data transfer to gpu vs compute. Isn't the shared lock just a lazy hack?"

The assistant acknowledged this immediately at [msg 445]: "You're right. The shared mutex is a hack that serializes everything to GPU 0 and wastes the second GPU entirely." The real problem was identified at line 483 of groth16_cuda.cu: size_t n_gpus = std::min(ngpus(), num_circuits). For partitioned proofs, num_circuits=1, so n_gpus=1, and the GPU selection logic always picked GPU 0.

The proper solution was clear: thread a gpu_index parameter through the entire call chain so that the C++ code uses the GPU assigned by the Rust engine instead of always defaulting to GPU 0. This required changes across multiple layers: the C++ groth16_cuda.cu, the Rust FFI in supraseal-c2/src/lib.rs, the bellperson prover functions, the pipeline layer, and finally the engine's GPU worker code in engine.rs.

What Message 464 Actually Contains

The message itself is deceptively simple — a single read tool call that displays lines 195–205 of the file /tmp/czk/extern/supraseal-c2/src/lib.rs:

195: pub fn alloc_gpu_mutex() -> *mut core::ffi::c_void {
196:     unsafe { create_gpu_mutex() }
197: }
198: 
199: /// Free a C++ `std::mutex` previously created by `alloc_gpu_mutex`.
200: ///
201: /// # Safety
202: /// `ptr` must be a pointer returned by `alloc_gpu_mutex` and must not
203: /// have been freed already.
204: pub unsafe fn free_gpu_mutex(ptr: *mut core::ffi::c_void) {
205:     destroy_gpu_mutex(ptr);

This is a fragment of the Rust FFI wrapper that bridges the gap between Rust's memory-safe world and the raw C++ mutex implementation. The alloc_gpu_mutex function creates a C++ std::mutex on the heap and returns an opaque pointer to it. The free_gpu_mutex function destroys it. The safety documentation is notable — it explicitly warns that the pointer must be valid and not already freed, a reminder of the contract between the two language worlds.

Why This Message Was Written: The Reasoning and Motivation

The assistant's motivation for reading this file at this exact moment is deeply strategic. The implementation of the multi-GPU fix was proceeding bottom-up: first the C++ layer (messages [msg 457] through [msg 462]), then the Rust FFI layer, then the bellperson prover functions, then the pipeline, and finally the engine. At message [msg 462], the assistant had just completed all C++ changes and marked that todo item as completed. The next item on the list was: "Add gpu_index to supraseal-c2 Rust FFI wrappers (start_groth16_proof, generate_groth16_proof)."

The assistant needed to understand the existing FFI code before modifying it. The read at message [msg 464] is not the first read of this file — the assistant had already read the beginning of lib.rs at [msg 463] to see the generate_groth16_proof function signature. Now it was reading further down to understand the mutex management functions and the overall structure of the FFI layer.

This is a pattern visible throughout the conversation: the assistant reads code in targeted chunks, building a mental model of the system before making changes. The read at message [msg 464] specifically targets the mutex-related functions because the multi-GPU fix involves not just adding a gpu_index parameter but also reverting the shared mutex hack and restoring per-GPU mutexes. Understanding how mutexes are allocated and freed is essential to getting the FFI changes right.

Input Knowledge Required

To fully understand this message, one needs knowledge of several domains:

The CuZK proving engine architecture: The system uses a layered design where a Rust engine manages GPU worker threads, each worker calls into a pipeline layer, which calls bellperson prover functions, which call into the supraseal-c2 Rust FFI library, which finally calls C++ CUDA code. The gpu_index parameter must flow through all these layers.

The FFI pattern: The Rust code uses extern "C" declarations to call C++ functions compiled into a shared library. The C++ functions return RustError structs and use opaque pointers for state. The alloc_gpu_mutex function wraps a C++ function create_gpu_mutex() that heap-allocates a std::mutex and returns a raw pointer.

The shared mutex hack: The previous fix had introduced a single shared mutex that all GPU workers contended on, serializing all GPU work onto GPU 0. This was a temporary band-aid that the assistant was now replacing with proper per-GPU mutexes.

The SnapDeals workload: The user had revealed at [msg 450] that SnapDeals consists of 16 identical partitions. This meant that with proper GPU routing, the 16 proofs would naturally load-balance across both GPUs — 8 on GPU 0 and 8 on GPU 1 — eliminating the OOM issue that occurred when two workers were crammed onto the same GPU.

Output Knowledge Created

This message creates knowledge in several forms:

For the assistant: The read provides the specific function signatures and patterns needed to add the gpu_index parameter. The assistant can now see how alloc_gpu_mutex and free_gpu_mutex work, which informs how the per-GPU mutexes should be managed after reverting the shared mutex hack.

For the reader of the conversation: The message reveals the structure of the FFI layer — the pattern of wrapping C++ functions in Rust-safe wrappers, the use of raw pointers, the safety contracts. It shows that the system uses per-GPU mutexes (created by alloc_gpu_mutex) rather than a single global mutex, which is the architecture that should have been used all along.

For the todo list: The message represents progress on the second todo item. After this read, the assistant will proceed to edit lib.rs to add the gpu_index parameter to the extern "C" declarations, the start_groth16_proof function, the generate_groth16_proof function, and the generate_groth16_proofs function.

The Thinking Process Visible in the Message

While the message itself is just a tool call, the thinking process is visible in the context. The assistant is working through a systematic plan:

  1. Start from the bottom (C++) and work up — this ensures that each layer has the right interface before the layer above it is modified.
  2. Read before writing — the assistant reads the existing code before making changes, building understanding incrementally.
  3. Trace the full call chain — earlier, at message [msg 449], the assistant had spawned a sub-agent task to trace the complete call chain from the Rust engine worker through to the C++ function. This systematic approach ensured no call site was missed.
  4. Use -1 as a sentinel — the assistant decided that gpu_index = -1 would mean "auto / use all GPUs like before," preserving backward compatibility for callers that don't specify a GPU. The todo list visible in the message header shows the assistant's progress tracking: the C++ changes are marked "completed," the supraseal-c2 Rust FFI changes are "in_progress," and the remaining layers (bellperson, pipeline, engine) are still "pending." This structured approach to multi-layer refactoring is a hallmark of disciplined systems engineering.

Assumptions and Potential Issues

The assistant makes several assumptions in this message:

That the d_a_cache global singleton won't cause problems: At [msg 456], the assistant noted that d_a_cache is a single global cache that will thrash between GPUs when two workers are active concurrently. The assistant decided to make it a per-GPU array but deferred that change — the immediate focus was on the gpu_index plumbing. This is a reasonable prioritization: get the routing correct first, then optimize the caching.

That the -1 sentinel convention is sufficient: The assumption is that existing callers (non-engine paths) will pass -1 and get the old behavior. This requires that all call sites be updated, which the assistant plans to do systematically.

That the per-GPU mutexes already exist: The code shows alloc_gpu_mutex and free_gpu_mutex as existing functions, suggesting the system was designed for per-GPU mutexes but the shared mutex hack bypassed that design. The assistant is restoring the original architecture.

The Broader Significance

Message [msg 464] is, on its surface, a mundane read operation. But it represents a crucial architectural pivot: the rejection of a quick hack in favor of a proper fix. The shared mutex was the easy path — it would have "worked" in the sense of preventing data races, but at the cost of wasting half the GPU hardware. The proper fix required touching five layers of code across two languages, understanding the full call chain, and carefully threading a new parameter through every intermediate function.

This is the kind of engineering discipline that separates robust systems from fragile ones. The assistant could have declared victory with the shared mutex — after all, it fixed the crash. But the user's challenge at [msg 444] pushed for a better solution, and the assistant responded by systematically implementing the right architecture.

The message also illustrates a key principle of multi-layer refactoring: work bottom-up. By starting with the C++ layer and moving upward through the Rust FFI, bellperson, pipeline, and engine, the assistant ensures that each layer's interface is stable before the layer above depends on it. This minimizes the risk of cascading changes and makes the refactoring predictable and safe.

Conclusion

Message [msg 464] is a quiet moment in a storm of activity — a single read call that bridges the C++ and Rust halves of a critical multi-GPU fix. It captures the assistant's methodical approach to systems engineering: read before writing, work bottom-up, trace the full call chain, and never settle for a lazy hack when a proper architectural fix is within reach. The message itself is just a fragment of code, but the context transforms it into a case study in disciplined refactoring across language boundaries.