Route every decoder device pin through one resolver and two sanctioned wrappers

Signed-off-by: Melody Ren <melodyr@nvidia.com>

Author: melody-ren

libs/qec/lib/hardware_guards.h

@@ -8,6 +8,9 @@
 
 #pragma once
 
+#include "cudaq/qec/decoder.h"
+
+#include <cassert>
 #include <cuda_runtime_api.h>
 #include <stdexcept>
 #include <string>
@@ -30,6 +33,37 @@ inline void set_cuda_device_for_decode(int target) {
                              ") failed: " + cudaGetErrorString(err));
 }
 
+/// Resolve a decoder's device: its cuda_device_id, or 0 when unpinned. For
+/// paths that need a concrete device (graph capture/launch), unlike
+/// set_cuda_device_for_decode() which no-ops on < 0.
+inline int decode_device_for(int cuda_device_id) {
+  return cuda_device_id >= 0 ? cuda_device_id : 0;
+}
+
+/// Pin before dispatch / decode / get_corrections / reset. No-op for an
+/// unpinned decoder (cuda_device_id < 0): a CPU decoder must not be forced onto
+/// a GPU. The sanctioned dispatch pin for every transport.
+inline void pin_decode_device(const cudaq::qec::decoder &dec) {
+  set_cuda_device_for_decode(dec.get_cuda_device_id());
+}
+
+/// Capture a decoder's realtime graph, pinned to its device so capture lands on
+/// the GPU every launch uses. Unpinned resolves to device 0 (a graph needs a
+/// concrete device). The only sanctioned caller of capture_decode_graph().
+inline void *capture_graph_pinned(cudaq::qec::decoder &dec,
+                                  int reserved_sms = 0) {
+  const int device = decode_device_for(dec.get_cuda_device_id());
+  set_cuda_device_for_decode(device);
+  void *raw = dec.capture_decode_graph(reserved_sms);
+#ifndef NDEBUG
+  int current = -1;
+  (void)cudaGetDevice(&current);
+  assert(current == device &&
+         "capture_graph_pinned: capture did not land on the decoder's device");
+#endif
+  return raw;
+}
+
 /// RAII: set the calling thread's CUDA device, restore the previous device on
 /// scope exit. No-op for target < 0. Lib-private and header-only so decoder
 /// plugins built as separate .so files can reuse it (PR2 extends this header

libs/qec/lib/realtime/decoding-server-cqr/DecodingServer.cpp

@@ -8,6 +8,7 @@
 
 #include "DecodingServer.h"
 #include "CpuRoceTransceiver.h"
+#include "../../hardware_guards.h"
 
 #include "cudaq/qec/logger.h"
 #include "cudaq/qec/realtime/decoding_config.h"
@@ -45,7 +46,7 @@ using cudaq::qec::decoding::config::DecoderTransport;
 /// graphs cannot split capture and launch across devices, so the decoder must
 /// be pinned to that device.
 int resolve_decode_device(int decoder_pin) {
-  return decoder_pin >= 0 ? decoder_pin : 0;
+  return detail_affinity::decode_device_for(decoder_pin);
 }
 
 std::unique_ptr<ITransceiver>

libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp

@@ -7,7 +7,6 @@
  ******************************************************************************/
 
 #include "DecodingSession.h"
-#include "DecodingServer.h"
 #include "RpcWireFormat.h"
 #include "../../hardware_guards.h"
 #include "cudaq/qec/logger.h"
@@ -20,19 +19,6 @@
 
 namespace cudaq::qec::decoding_server {
 
-namespace {
-
-void set_graph_capture_device(const cudaq::qec::decoder &decoder) {
-  const int device = resolve_decode_device(decoder.get_cuda_device_id());
-  cudaq::qec::detail_affinity::set_cuda_device_for_decode(device);
-  if (device >= 0)
-    CUDA_QEC_INFO(
-        "DecodingSession::create: set CUDA device {} before graph capture",
-        device);
-}
-
-} // namespace
-
 // Busy high-water mark across all sessions (worker threads increment while
 // executing an item).
 static std::atomic<uint64_t> g_busy_sessions{0};
@@ -67,8 +53,7 @@ DecodingSession::create(std::unique_ptr<cudaq::qec::decoder> decoder,
   s->dec = std::move(decoder);
 
   if (s->dec->supports_graph_dispatch()) {
-    set_graph_capture_device(*s->dec);
-    void *gr = s->dec->capture_decode_graph();
+    void *gr = cudaq::qec::detail_affinity::capture_graph_pinned(*s->dec);
     s->graph_resources =
         GraphResourcesPtr(gr, GraphResourcesDeleter{s->dec.get()});
   }
@@ -86,8 +71,7 @@ void DecodingSession::start_worker() {
   auto pin_result = pinned.get_future();
   worker = std::thread([this, &pinned] {
     try {
-      cudaq::qec::detail_affinity::set_cuda_device_for_decode(
-          dec->get_cuda_device_id());
+      cudaq::qec::detail_affinity::pin_decode_device(*dec);
       pinned.set_value();
     } catch (...) {
       pinned.set_exception(std::current_exception());

libs/qec/lib/realtime/qec_realtime_session.cpp

@@ -140,8 +140,7 @@ static void apply_decoder_cuda_device(cudaq::qec::decoder *dec) {
   // Throws on failure (fail fast): host dispatch surfaces it as an error
   // response and graph initialization aborts, rather than continuing on
   // whichever device happened to be current.
-  cudaq::qec::detail_affinity::set_cuda_device_for_decode(
-      dec->get_cuda_device_id());
+  cudaq::qec::detail_affinity::pin_decode_device(*dec);
 }
 
 // Two-ring response writer: the request stays in `rx_slot` (read-only); the
@@ -557,9 +556,9 @@ void qec_realtime_session::capture_decoder_graphs() {
           "qec_realtime_session::initialize: decoder " + std::to_string(i) +
           " does not support graph dispatch in DEVICE mode.");
 
-    apply_decoder_cuda_device(dec);
     // reserved_sms = 0 is intentional for the inproc_rpc desktop / CI path.
-    void *raw = dec->capture_decode_graph(/*reserved_sms=*/0);
+    void *raw =
+        cudaq::qec::detail_affinity::capture_graph_pinned(*dec, /*sms=*/0);
     if (!raw)
       throw std::runtime_error("qec_realtime_session::initialize: decoder " +
                                std::to_string(i) +

libs/qec/lib/realtime/realtime_decoding.cpp

@@ -432,8 +432,7 @@ void enqueue_syndromes(std::size_t decoder_id, uint8_t *syndromes,
   // decoder's pinned device before decoding (set-if-different; throws on
   // failure) -- and before the capture callback, so a pin failure cannot
   // record a round that was never decoded.
-  cudaq::qec::detail_affinity::set_cuda_device_for_decode(
-      decoder->get_cuda_device_id());
+  cudaq::qec::detail_affinity::pin_decode_device(*decoder);
   capture_syndromes();
 
   std::vector<uint8_t> syndrome_u8(syndrome_length);
@@ -501,8 +500,7 @@ void get_corrections(std::size_t decoder_id, uint8_t *corrections,
 #endif
 
   // clear_corrections may touch device memory in some plugins.
-  cudaq::qec::detail_affinity::set_cuda_device_for_decode(
-      decoder->get_cuda_device_id());
+  cudaq::qec::detail_affinity::pin_decode_device(*decoder);
   auto ret = decoder->get_obs_corrections();
   for (std::size_t i = 0; i < correction_length; ++i) {
     corrections[i] = ret[i];
@@ -538,8 +536,7 @@ void reset_decoder(std::size_t decoder_id) {
   }
 #endif
 
-  cudaq::qec::detail_affinity::set_cuda_device_for_decode(
-      decoder->get_cuda_device_id());
+  cudaq::qec::detail_affinity::pin_decode_device(*decoder);
   decoder->reset_decoder();
 }