The Final Sweep: Tracing Every Call Site in a Multi-GPU Refactor

Introduction

In a sprawling refactor that threads a gpu_index parameter from C++ CUDA kernels up through five layers of Rust abstractions, the most dangerous moment is not the first edit — it is the last. After modifying the C++ GPU entry point, the Rust FFI wrappers, the bellperson prover functions, the pipeline layer, and the engine's worker dispatch logic, a single overlooked call site can silently route work to the wrong GPU, undoing the entire fix. Message [msg 496] captures this exact moment: a lone read tool call, unassuming in isolation, that reveals the assistant's disciplined approach to ensuring completeness in a cross-cutting architectural change.

The Message

The subject message is brief and procedural:

[assistant] [read] /tmp/czk/extern/cuzk/cuzk-core/src/pipeline.rs
<path>/tmp/czk/extern/cuzk/cuzk-core/src/pipeline.rs</path>
<type>file</type>
<content>2318:                         "GPU received partition, starting prove"
2319:                     );
2320: 
2321:                     let gpu_result = gpu_prove(slot.synth, params_ref, std::ptr::null_mut())?;
2322:                     total_gpu += gpu_result.gpu_duration;
2323: 
2324:                     info!(
2325:                         partition = pidx,

(File has more lines. Use 'offset' parameter to read...

There is no reasoning block, no commentary, no todo update — just a focused read of a specific file at a specific line offset. The assistant is hunting for remaining callers of gpu_prove that still use the old three-argument signature (std::ptr::null_mut() as the mutex parameter) rather than the new four-argument signature that includes gpu_index.

Context and Background

To understand why this single read matters, one must appreciate the scale of the refactor that precedes it. The session's overarching problem was a GPU data race on multi-GPU systems. The C++ GPU proving code in groth16_cuda.cu always routed single-circuit proofs to GPU 0, regardless of which Rust worker submitted them. The initial "fix" — a shared mutex that serialized all partition proofs onto GPU 0 — was a lazy hack that wasted the second GPU entirely. When a SnapDeals workload with 16 identical partitions OOM'd on a 20 GB RTX 4000 Ada host, it became clear that the shared mutex was insufficient: two workers still entered the GPU code simultaneously, and the VRAM budget for a single SnapDeals partition was too large to allow concurrent kernel execution on the same device.

The proper solution was to 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:

  1. C++ (groth16_cuda.cu): Add gpu_index parameter to generate_groth16_proofs_start_c and generate_groth16_proofs_c; when gpu_index &gt;= 0, force n_gpus = 1 and use select_gpu(gpu_index) instead of select_gpu(tid).
  2. Rust FFI (supraseal-c2/src/lib.rs): Update extern declarations and public wrapper functions (start_groth16_proof, generate_groth16_proof, generate_groth16_proofs) to accept and pass through the new parameter.
  3. Bellperson (bellperson/src/groth16/prover/supraseal.rs): Update prove_start and prove_from_assignments to accept gpu_index and forward it to the FFI layer.
  4. Pipeline (cuzk-core/src/pipeline.rs): Update gpu_prove and gpu_prove_start to accept gpu_index and pass it to bellperson.
  5. Engine (cuzk-core/src/engine.rs): Revert the shared mutex hack, restore per-GPU mutexes, and pass gpu_ordinal as gpu_index at each call site. By message [msg 493], the assistant had completed all the primary edits and verified that the two engine call sites now pass gpu_ordinal as i32. But a grep revealed 15 matches for gpu_prove( and gpu_prove_start(, indicating there were additional callers outside the engine that still needed updating.## The Hunt for Stale Call Sites The grep result at [msg 493] was the trigger. It showed 15 matches across engine.rs and pipeline.rs. The two in engine.rs had already been updated. But the remaining 13 — all in pipeline.rs — were still using the old three-parameter signature. The assistant had already updated some of them in [msg 495], where a call at line 1815 (gpu_prove(synth, params, std::ptr::null_mut())) was changed to include -1 as the gpu_index parameter, signaling "auto" mode for non-engine paths. But the grep output showed more matches than just line 1815. The assistant needed to verify that every call site was caught. Message [msg 496] is the direct consequence of that grep: the assistant is reading around line 2321 of pipeline.rs to find another caller that still uses std::ptr::null_mut() without a gpu_index. The line in question reads:
let gpu_result = gpu_prove(slot.synth, params_ref, std::ptr::null_mut())?;

This is a call inside what appears to be a partition-processing loop — likely the monolithic (non-partitioned) GPU proving path that handles individual slots. The std::ptr::null_mut() here means this caller was using the old three-argument signature, which would fail to compile after the refactor changed gpu_prove to require four arguments.

Why This Matters

The message is a testament to the difficulty of cross-cutting refactors in large codebases. When a parameter is added to a function deep in the call stack, the compiler will catch direct callers of that function — but only if they are in the same compilation unit or if the function signature change breaks the build. However, in a Rust project with conditional compilation features (#[cfg(feature = &#34;cuda-supraseal&#34;)]), some call sites may be behind feature gates that are not exercised in the current build configuration. A cargo check on the developer's machine might not flag every instance.

More subtly, the assistant is not just looking for compilation errors. It is looking for semantic correctness. A caller that passes std::ptr::null_mut() for the mutex and omits gpu_index entirely would, under the old signature, rely on the C++ internal fallback mutex — which always routes to GPU 0. If such a caller were left untouched, it would silently bypass the entire multi-GPU fix, continuing to serialize work onto GPU 0 while the engine's properly routed workers use both GPUs. The result would be a confusing partial fix: some proofs would use both GPUs, while others (those going through the monolithic path) would pile onto GPU 0, potentially causing OOMs or performance degradation.

Assumptions and Reasoning

The assistant's reasoning in this message is implicit but clear. The assumption is that every call site of gpu_prove and gpu_prove_start must be updated to pass gpu_index, and that non-engine paths (those not dispatched by the GPU worker in engine.rs) should use -1 to indicate "auto-select" behavior. This is a reasonable default: the C++ code, when receiving gpu_index = -1, falls back to its original behavior of distributing circuits across all available GPUs based on thread ID. For non-engine callers — such as the monolithic proving path in pipeline.rs that handles a single partition at a time — this preserves backward compatibility while still benefiting from multi-GPU distribution when multiple circuits are present.

The assistant also assumes that gpu_prove and gpu_prove_start are the only entry points to the GPU proving path that need updating. This is validated by the grep at [msg 493], which found exactly 15 matches for these two functions. If there were other entry points (e.g., direct calls to bellperson::groth16::prove_from_assignments from outside the pipeline), they would not appear in this grep and would remain broken. The assistant's confidence in this grep is an assumption worth noting: it assumes that all GPU proving flows go through either gpu_prove or gpu_prove_start.

Input and Output Knowledge

The input knowledge required to understand this message includes: