From c42372f6e6c190ac586e33d532d04fe6256bf0ce Mon Sep 17 00:00:00 2001 From: Melody Ren Date: Tue, 14 Jul 2026 23:24:57 -0700 Subject: [PATCH] Route every decoder device pin through one resolver and two sanctioned wrappers Signed-off-by: Melody Ren --- libs/qec/lib/hardware_guards.h | 34 +++++++++++++++++++ .../decoding-server-cqr/DecodingServer.cpp | 3 +- .../decoding-server-cqr/DecodingSession.cpp | 20 ++--------- .../qec/lib/realtime/qec_realtime_session.cpp | 7 ++-- libs/qec/lib/realtime/realtime_decoding.cpp | 9 ++--- 5 files changed, 44 insertions(+), 29 deletions(-) diff --git a/libs/qec/lib/hardware_guards.h b/libs/qec/lib/hardware_guards.h index e4178ccce..d17033936 100644 --- a/libs/qec/lib/hardware_guards.h +++ b/libs/qec/lib/hardware_guards.h @@ -8,6 +8,9 @@ #pragma once +#include "cudaq/qec/decoder.h" + +#include #include #include #include @@ -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(¤t); + 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 diff --git a/libs/qec/lib/realtime/decoding-server-cqr/DecodingServer.cpp b/libs/qec/lib/realtime/decoding-server-cqr/DecodingServer.cpp index 5e3a7a8a8..a0071f584 100644 --- a/libs/qec/lib/realtime/decoding-server-cqr/DecodingServer.cpp +++ b/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 diff --git a/libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp b/libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp index a12d5b541..1bceabd2c 100644 --- a/libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp +++ b/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 g_busy_sessions{0}; @@ -67,8 +53,7 @@ DecodingSession::create(std::unique_ptr 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()); diff --git a/libs/qec/lib/realtime/qec_realtime_session.cpp b/libs/qec/lib/realtime/qec_realtime_session.cpp index e949bba82..d10413525 100644 --- a/libs/qec/lib/realtime/qec_realtime_session.cpp +++ b/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) + diff --git a/libs/qec/lib/realtime/realtime_decoding.cpp b/libs/qec/lib/realtime/realtime_decoding.cpp index 22ba4b1e4..63baa2774 100644 --- a/libs/qec/lib/realtime/realtime_decoding.cpp +++ b/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 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(); }