diff --git a/docs/sphinx/examples_rst/qec/realtime_decoding.rst b/docs/sphinx/examples_rst/qec/realtime_decoding.rst index 814c9c8f2..2f5cff664 100644 --- a/docs/sphinx/examples_rst/qec/realtime_decoding.rst +++ b/docs/sphinx/examples_rst/qec/realtime_decoding.rst @@ -155,6 +155,7 @@ arguments: decoders: - id: 0 type: pymatching + cuda_device_id: 0 # optional: pin this decoder to a CUDA device block_size: 3 syndrome_size: 3 H_sparse: [ 0, -1, 1, -1, 2, -1 ] @@ -164,6 +165,14 @@ arguments: error_rate_vec: [ 0.1, 0.1, 0.1 ] merge_strategy: smallest_weight +``cuda_device_id`` pins a GPU-accelerated decoder (e.g. ``nv-qldpc-decoder`` +or ``trt_decoder``) to a specific CUDA device. The same knob is available as +a construction parameter in C++ and Python +(``qec.get_decoder("trt_decoder", H, cuda_device_id=1)``). The thread that +creates a decoder is pinned to that device and is expected to drive its +decode calls; create each pinned decoder on its own thread to place several +decoders on different GPUs. + Here is how to create and save a decoder configuration: .. tab:: Python diff --git a/libs/qec/include/cudaq/qec/decoder.h b/libs/qec/include/cudaq/qec/decoder.h index 0a2f4c0c0..7c1d2e351 100644 --- a/libs/qec/include/cudaq/qec/decoder.h +++ b/libs/qec/include/cudaq/qec/decoder.h @@ -260,6 +260,12 @@ class decoder /// depends on D_sparse, so you must have called set_D_sparse() first. uint32_t get_num_msyn_per_decode() const; + /// @brief The CUDA device this decoder was pinned to at construction via + /// the "cuda_device_id" parameter, or -1 when no pin was requested. + /// Construction pins the constructing thread persistently (the thread that + /// creates a decoder is the thread expected to drive its decode calls). + int get_cuda_device_id() const { return cuda_device_id_; } + /// @brief Set the observable matrix. void set_O_sparse(const std::vector> &O_sparse); @@ -357,6 +363,10 @@ class decoder /// @brief The decoder's D matrix in sparse format std::vector> D_sparse; + /// @brief CUDA device id consumed from the construction parameters by + /// decoder::get(); -1 = unpinned. See get_cuda_device_id(). + int cuda_device_id_ = -1; + private: decode_result_type result_type_ = decode_result_type::decode_to_errs; }; diff --git a/libs/qec/include/cudaq/qec/realtime/decoding_config.h b/libs/qec/include/cudaq/qec/realtime/decoding_config.h index 7d4a3f2d2..3608ec43e 100644 --- a/libs/qec/include/cudaq/qec/realtime/decoding_config.h +++ b/libs/qec/include/cudaq/qec/realtime/decoding_config.h @@ -183,6 +183,11 @@ struct decoder_config { /// Defaults to cpu_roce. Set to gpu_roce for decoders where syndrome bits /// are DMA'd directly to GPU VRAM (e.g. nv_qldpc_decoder with RelayBP). DecoderTransport transport = DecoderTransport::cpu_roce; + /// CUDA device this decoder is pinned to at construction (see the + /// "cuda_device_id" decoder parameter). Placement knob common to any + /// GPU-accelerated decoder, hence at this level rather than inside the + /// per-decoder custom args. Unset = unpinned. + std::optional cuda_device_id; uint64_t block_size = 0; uint64_t syndrome_size = 0; std::vector H_sparse; diff --git a/libs/qec/lib/CMakeLists.txt b/libs/qec/lib/CMakeLists.txt index f9a52e1d6..0a8fad91e 100644 --- a/libs/qec/lib/CMakeLists.txt +++ b/libs/qec/lib/CMakeLists.txt @@ -150,6 +150,10 @@ target_link_libraries(${DECODERS_LIBRARY_NAME} PUBLIC $ fmt::fmt-header-only + PRIVATE + # decoder::get() consumes the "cuda_device_id" construction parameter + # (validation via cudaGetDeviceCount + persistent cudaSetDevice pin). + CUDA::cudart ) target_link_libraries(${LIBRARY_NAME} diff --git a/libs/qec/lib/decoder.cpp b/libs/qec/lib/decoder.cpp index 796427a54..03f35cbe5 100644 --- a/libs/qec/lib/decoder.cpp +++ b/libs/qec/lib/decoder.cpp @@ -11,7 +11,9 @@ #include "cudaq/qec/logger.h" #include "cudaq/qec/plugin_loader.h" #include "cudaq/qec/version.h" +#include "hardware_guards.h" #include +#include #include #include #include @@ -123,8 +125,33 @@ std::string decoder::get_version() const { std::future decoder::decode_async(const std::vector &syndrome) { - return std::async(std::launch::async, - [this, syndrome] { return this->decode(syndrome); }); + // Captured by value: the worker must not dereference decoder members to + // find its device. The std::async thread is brand-new and unpinned, so it + // guards itself for the duration of the call (the one exception to the + // one-thread-owns-one-decoder persistent pin). + const int cuda_id = cuda_device_id_; + return std::async(std::launch::async, [this, syndrome, cuda_id] { + cudaq::qec::detail_affinity::CudaDeviceGuard dev(cuda_id); + return this->decode(syndrome); + }); +} + +/// Reads "cuda_device_id" from the construction parameters. Absent -> -1. +/// Negative or >= cudaGetDeviceCount() -> std::runtime_error (fail fast: +/// never silently decode on the wrong GPU). +static int read_cuda_device_id(const cudaqx::heterogeneous_map ¶ms) { + if (!params.contains("cuda_device_id")) + return -1; + const int value = params.get("cuda_device_id"); + if (value < 0) + throw std::runtime_error("cuda_device_id must be >= 0 (got " + + std::to_string(value) + ")"); + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess || value >= count) + throw std::runtime_error( + "cuda_device_id " + std::to_string(value) + " is out of range: " + + std::to_string(count) + " CUDA device(s) visible"); + return value; } std::unique_ptr @@ -138,7 +165,27 @@ decoder::get(const std::string &name, const decoder_init &init, "invalid decoder requested: " + name + ". Run with CUDAQ_LOG_LEVEL=info (environment variable) to see " "additional plugin diagnostics at startup."); - return iter->second(init, param_map); + const int cuda_device_id = read_cuda_device_id(param_map); + if (cuda_device_id < 0) + return iter->second(init, param_map); + // Pin the constructing thread persistently (no restore): one thread owns + // one decoder, so every later allocation and kernel launch on this thread + // -- including lazy allocations inside a plugin's decode() -- lands on the + // requested device with no per-call machinery. + cudaError_t err = cudaSetDevice(cuda_device_id); + if (err != cudaSuccess) + throw std::runtime_error("cudaSetDevice(" + + std::to_string(cuda_device_id) + + ") failed: " + cudaGetErrorString(err)); + // The key is consumed here; strip it so plugins that strictly validate + // their parameter keys do not reject it. + cudaqx::heterogeneous_map plugin_params; + for (const auto &kv : param_map) + if (kv.first != "cuda_device_id") + plugin_params.insert(kv.first, kv.second); + auto d = iter->second(init, plugin_params); + d->cuda_device_id_ = cuda_device_id; + return d; } namespace details { diff --git a/libs/qec/lib/hardware_guards.h b/libs/qec/lib/hardware_guards.h new file mode 100644 index 000000000..347681cbb --- /dev/null +++ b/libs/qec/lib/hardware_guards.h @@ -0,0 +1,57 @@ +/******************************************************************************* + * Copyright (c) 2022 - 2026 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#pragma once + +#include +#include +#include + +namespace cudaq::qec::detail_affinity { + +/// 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 +/// with NUMA guards; the nv-qldpc follow-up mirrors its use). +/// +/// This guard is for threads that do NOT follow the one-thread-owns-one- +/// decoder persistent pin (e.g. the fresh worker spawned by decode_async). +class CudaDeviceGuard { +public: + explicit CudaDeviceGuard(int target) { + if (target < 0) + return; + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess || target >= count) + throw std::runtime_error( + "cuda_device_id " + std::to_string(target) + " is out of range: " + + std::to_string(count) + " CUDA device(s) visible"); + // If the current device is unreadable, skip restoration rather than + // restore to a guessed device; the set below still applies. + if (cudaGetDevice(&prev_) != cudaSuccess) + prev_ = -1; + cudaError_t err = cudaSetDevice(target); + if (err != cudaSuccess) + throw std::runtime_error("CudaDeviceGuard: cudaSetDevice(" + + std::to_string(target) + + ") failed: " + cudaGetErrorString(err)); + restore_ = (prev_ >= 0 && prev_ != target); + } + ~CudaDeviceGuard() { + if (restore_) + (void)cudaSetDevice(prev_); + } + CudaDeviceGuard(const CudaDeviceGuard &) = delete; + CudaDeviceGuard &operator=(const CudaDeviceGuard &) = delete; + +private: + int prev_ = -1; + bool restore_ = false; +}; + +} // namespace cudaq::qec::detail_affinity diff --git a/libs/qec/lib/realtime/config.cpp b/libs/qec/lib/realtime/config.cpp index a568e6673..9b100319f 100644 --- a/libs/qec/lib/realtime/config.cpp +++ b/libs/qec/lib/realtime/config.cpp @@ -717,6 +717,7 @@ struct MappingTraits { io.mapRequired("type", config.type); io.mapOptional("transport", config.transport, cudaq::qec::decoding::config::DecoderTransport::cpu_roce); + io.mapOptional("cuda_device_id", config.cuda_device_id); io.mapRequired("block_size", config.block_size); io.mapRequired("syndrome_size", config.syndrome_size); io.mapRequired("H_sparse", config.H_sparse); diff --git a/libs/qec/lib/realtime/qec_realtime_session.cpp b/libs/qec/lib/realtime/qec_realtime_session.cpp index dbe35c2f5..6ed6391f1 100644 --- a/libs/qec/lib/realtime/qec_realtime_session.cpp +++ b/libs/qec/lib/realtime/qec_realtime_session.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -126,6 +127,26 @@ cudaq::qec::decoder *get_decoder_or_throw(std::int64_t decoder_id) { return (*decoders)[static_cast(decoder_id)].get(); } +// Point the calling thread at the decoder's pinned CUDA device before work +// that allocates or launches on it. Set-and-leave (no restore): one +// dispatcher thread serves all decoders, so the thread simply converges to +// the device of the decoder it is currently serving; cudaSetDevice on an +// already-current device is a cheap no-op. +static void apply_decoder_cuda_device(cudaq::qec::decoder *dec) { + if (!dec) + return; + const int id = dec->get_cuda_device_id(); + if (id < 0) + return; + int cur = -1; + if (cudaGetDevice(&cur) == cudaSuccess && cur == id) + return; + cudaError_t err = cudaSetDevice(id); + if (err != cudaSuccess) + CUDA_QEC_WARN("apply_decoder_cuda_device: cudaSetDevice({}) failed: {}", + id, cudaGetErrorString(err)); +} + // Two-ring response writer: the request stays in `rx_slot` (read-only); the // response is written into the distinct `tx_slot`. The preserved header fields // (request_id, ptp_timestamp) must be echoed explicitly from rx to tx. The @@ -178,6 +199,7 @@ void enqueue_syndromes_host(const void *rx_slot, void *tx_slot, } auto *decoder = get_decoder_or_throw(body->decoder_id); + apply_decoder_cuda_device(decoder); // Reject requests larger than this decoder's per-decode window. The slot // is sized for the largest decoder in the session, so an oversized request // for a smaller decoder can still fit the slot; without this guard it would @@ -537,6 +559,7 @@ 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); if (!raw) diff --git a/libs/qec/lib/realtime/realtime_decoding.cpp b/libs/qec/lib/realtime/realtime_decoding.cpp index 23c0ad4c0..d423712b3 100644 --- a/libs/qec/lib/realtime/realtime_decoding.cpp +++ b/libs/qec/lib/realtime/realtime_decoding.cpp @@ -156,6 +156,10 @@ namespace cudaq::qec::decoding::host { cudaqx::heterogeneous_map prepare_decoder_params( const cudaq::qec::decoding::config::decoder_config &decoder_config) { auto params = decoder_config.decoder_custom_args_to_heterogeneous_map(); + // Placement knob: surfaced for every decoder type (deliberately before the + // trt-only early return below); consumed by decoder::get() at construction. + if (decoder_config.cuda_device_id.has_value()) + params.insert("cuda_device_id", decoder_config.cuda_device_id.value()); if (decoder_config.type != "trt_decoder") return params; diff --git a/libs/qec/python/bindings/py_decoding_config.cpp b/libs/qec/python/bindings/py_decoding_config.cpp index 283c90af1..98a4f9aa8 100644 --- a/libs/qec/python/bindings/py_decoding_config.cpp +++ b/libs/qec/python/bindings/py_decoding_config.cpp @@ -279,6 +279,7 @@ void bindDecodingConfig(nb::module_ &mod) { .def(nb::init<>()) .def_rw("id", &decoder_config::id) .def_rw("type", &decoder_config::type) + .def_rw("cuda_device_id", &decoder_config::cuda_device_id) .def_rw("block_size", &decoder_config::block_size) .def_rw("syndrome_size", &decoder_config::syndrome_size) .def_rw("H_sparse", &decoder_config::H_sparse) diff --git a/libs/qec/python/tests/test_decoder.py b/libs/qec/python/tests/test_decoder.py index 3201de1c6..8f00900a8 100644 --- a/libs/qec/python/tests/test_decoder.py +++ b/libs/qec/python/tests/test_decoder.py @@ -1002,6 +1002,32 @@ def test_get_decoder_rejects_unknown_decoder_for_stim_dem_text(): qec.get_decoder("__no_such_decoder__", "error(0.1) D0 L0\n") +def test_decoder_cuda_device_id_invalid_raises(): + H = create_test_matrix() + # A negative id is rejected before reaching the C++ guard: the kwargs + # marshalling layer stores Python ints as size_t, so nanobind refuses the + # negative value with a bare RuntimeError ("std::bad_cast"). + with pytest.raises(RuntimeError): + qec.get_decoder("single_error_lut", H, cuda_device_id=-2) + # An out-of-range id flows through kwargs to decoder::get(), which raises + # a runtime_error naming the offending parameter. + with pytest.raises(RuntimeError, match="cuda_device_id"): + qec.get_decoder("single_error_lut", H, cuda_device_id=1 << 20) + + +def test_decoder_cuda_device_id_valid(): + H = create_test_matrix() + try: + d = qec.get_decoder("single_error_lut", H, cuda_device_id=0) + except RuntimeError as e: + if "out of range" in str(e): + pytest.skip("no CUDA device visible") + raise + syndrome = create_test_syndrome() + result = d.decode(syndrome) + assert len(result.result) == H.shape[1] + + def test_get_decoder_user_O_wins_over_dem_derived(): dem_text = ("error(0.1) D0 L0\n" "error(0.1) D1 L0\n" diff --git a/libs/qec/unittests/CMakeLists.txt b/libs/qec/unittests/CMakeLists.txt index 344d05449..211ba17e2 100644 --- a/libs/qec/unittests/CMakeLists.txt +++ b/libs/qec/unittests/CMakeLists.txt @@ -35,7 +35,7 @@ find_package(CUDAToolkit REQUIRED) add_compile_options(-Wno-attributes) add_executable(test_decoders test_decoders.cpp decoders/sample_decoder.cpp) -target_link_libraries(test_decoders PRIVATE GTest::gtest_main cudaq-qec-decoders libstim) +target_link_libraries(test_decoders PRIVATE GTest::gtest_main cudaq-qec-decoders libstim CUDA::cudart) add_dependencies(CUDAQXQECUnitTests test_decoders) gtest_discover_tests(test_decoders) diff --git a/libs/qec/unittests/test_decoders.cpp b/libs/qec/unittests/test_decoders.cpp index 0f2f4c8d5..32fb4cf00 100644 --- a/libs/qec/unittests/test_decoders.cpp +++ b/libs/qec/unittests/test_decoders.cpp @@ -10,12 +10,15 @@ #include "cudaq/qec/decoder.h" #include "cudaq/qec/detector_error_model.h" #include "cudaq/qec/pcm_utils.h" +#include #include #include +#include #include #include #include #include +#include namespace { class ScopedEnv { @@ -1190,3 +1193,215 @@ TEST(SlidingWindowDecoder, BaseStreamingCopiesFirstRoundDetectors) { << "First-round detector copy runs, but the sliding window is not full " "yet so no final correction is committed."; } + +namespace { + +int cuda_device_count() { + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess) + return 0; + return count; +} + +/// Restores the caller's CUDA device on scope exit so the persistent pin +/// made by one test does not leak into the next (gtest shares the process). +class ScopedDeviceRestore { +public: + ScopedDeviceRestore() { + if (cudaGetDevice(&prev_) != cudaSuccess) + prev_ = -1; + } + ~ScopedDeviceRestore() { + if (prev_ >= 0) + (void)cudaSetDevice(prev_); + } + +private: + int prev_ = -1; +}; + +/// A decoder that rejects any construction parameter it does not know, +/// proving decoder::get() strips cuda_device_id before the plugin ctor. +class strict_keys_decoder : public cudaq::qec::decoder { +public: + strict_keys_decoder(const cudaq::qec::sparse_binary_matrix &H, + const cudaqx::heterogeneous_map ¶ms) + : decoder(H) { + auto invalid = + cudaq::qec::validate_config_parameters(params, {"decode_to_obs"}); + if (!invalid.empty()) + throw std::runtime_error("strict_keys_decoder: unexpected key " + + invalid.front()); + } + cudaq::qec::decoder_result + decode(const std::vector &syndrome) override { + cudaq::qec::decoder_result r; + r.converged = true; + r.result = std::vector(block_size, 0.0); + return r; + } + CUDAQ_EXTENSION_CUSTOM_CREATOR_FUNCTION( + strict_keys_decoder, + static std::unique_ptr create( + const cudaq::qec::decoder_init &init, + const cudaqx::heterogeneous_map ¶ms) { + return cudaq::qec::make_pcm_decoder(init, params); + }) +}; +CUDAQ_EXT_PT_REGISTER_TYPE(strict_keys_decoder) + +cudaq::qec::sparse_binary_matrix make_test_H() { + cudaqx::tensor H({std::size_t{4}, std::size_t{10}}); + return cudaq::qec::sparse_binary_matrix(H); +} + +/// Records the CUDA device current on the thread that runs decode(), so a +/// test can observe which device an async worker thread actually used. +class device_recording_decoder : public cudaq::qec::decoder { +public: + std::atomic last_decode_device{-2}; + device_recording_decoder(const cudaq::qec::sparse_binary_matrix &H, + const cudaqx::heterogeneous_map &) + : decoder(H) {} + cudaq::qec::decoder_result + decode(const std::vector &) override { + int dev = -1; + if (cudaGetDevice(&dev) != cudaSuccess) + dev = -1; + last_decode_device.store(dev); + cudaq::qec::decoder_result r; + r.converged = true; + r.result = std::vector(block_size, 0.0); + return r; + } + CUDAQ_EXTENSION_CUSTOM_CREATOR_FUNCTION( + device_recording_decoder, + static std::unique_ptr create( + const cudaq::qec::decoder_init &init, + const cudaqx::heterogeneous_map ¶ms) { + return cudaq::qec::make_pcm_decoder(init, + params); + }) +}; +CUDAQ_EXT_PT_REGISTER_TYPE(device_recording_decoder) + +} // namespace + +TEST(DecoderCudaDeviceId, AbsentKeyIsNoOp) { + auto d = cudaq::qec::decoder::get("sample_decoder", make_test_H()); + EXPECT_EQ(d->get_cuda_device_id(), -1); + std::vector syndrome(4); + auto r = d->decode(syndrome); + EXPECT_EQ(r.result.size(), 10); +} + +TEST(DecoderCudaDeviceId, NegativeIdThrows) { + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", -2); + try { + auto d = cudaq::qec::decoder::get("sample_decoder", make_test_H(), params); + FAIL() << "expected std::runtime_error"; + } catch (const std::runtime_error &e) { + EXPECT_NE(std::string(e.what()).find("cuda_device_id"), std::string::npos); + } +} + +TEST(DecoderCudaDeviceId, OutOfRangeIdThrows) { + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", 1 << 20); + try { + auto d = cudaq::qec::decoder::get("sample_decoder", make_test_H(), params); + FAIL() << "expected std::runtime_error"; + } catch (const std::runtime_error &e) { + EXPECT_NE(std::string(e.what()).find("cuda_device_id"), std::string::npos); + } +} + +TEST(DecoderCudaDeviceId, PersistentPinAtConstruction) { + if (cuda_device_count() < 2) + GTEST_SKIP() << "needs >= 2 CUDA devices"; + ScopedDeviceRestore restore; + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", 1); + auto d = cudaq::qec::decoder::get("sample_decoder", make_test_H(), params); + EXPECT_EQ(d->get_cuda_device_id(), 1); + // The pin is persistent: the constructing thread is still on device 1 + // after decoder::get() returns. + int cur = -1; + ASSERT_EQ(cudaGetDevice(&cur), cudaSuccess); + EXPECT_EQ(cur, 1); + // Decode entry points need no per-call guard on this thread. + std::vector syndrome(4); + auto r = d->decode(syndrome); + EXPECT_EQ(r.result.size(), 10); + ASSERT_EQ(cudaGetDevice(&cur), cudaSuccess); + EXPECT_EQ(cur, 1); +} + +TEST(DecoderCudaDeviceId, KeyStrippedFromPluginParams) { + if (cuda_device_count() < 1) + GTEST_SKIP() << "needs >= 1 CUDA device"; + ScopedDeviceRestore restore; + // Sanity: strict_keys_decoder does reject unknown keys. + cudaqx::heterogeneous_map bogus; + bogus.insert("bogus_key", 1); + EXPECT_THROW( + cudaq::qec::decoder::get("strict_keys_decoder", make_test_H(), bogus), + std::runtime_error); + // cuda_device_id must be consumed by the base and never reach the plugin, + // while permitted keys (decode_to_obs) must survive the rebuild-and-strip. + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", 0); + params.insert("decode_to_obs", true); + EXPECT_NO_THROW( + cudaq::qec::decoder::get("strict_keys_decoder", make_test_H(), params)); +} + +TEST(DecoderCudaDeviceId, AsyncWorkerPinsItself) { + if (cuda_device_count() < 2) + GTEST_SKIP() << "needs >= 2 CUDA devices"; + ScopedDeviceRestore restore; + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", 1); + auto d = cudaq::qec::decoder::get("device_recording_decoder", make_test_H(), + params); + auto *rec = dynamic_cast(d.get()); + ASSERT_NE(rec, nullptr); + // decode_async spawns a brand-new std::async thread whose current device + // defaults to 0, NOT the owning thread's device 1. The worker must pin + // itself to the decoder's device for the duration of the call. + std::vector syndrome(4); + auto r = d->decode_async(syndrome).get(); + EXPECT_EQ(r.result.size(), 10); + EXPECT_EQ(rec->last_decode_device.load(), 1); + // The owning (calling) thread's device is untouched by the async call. + int cur = -1; + ASSERT_EQ(cudaGetDevice(&cur), cudaSuccess); + EXPECT_EQ(cur, 1); +} + +TEST(DecoderCudaDeviceId, TwoThreadsTwoDevices) { + if (cuda_device_count() < 2) + GTEST_SKIP() << "needs >= 2 CUDA devices"; + auto worker = [](int id, int &observed_device, bool &decode_ok) { + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", id); + auto d = + cudaq::qec::decoder::get("sample_decoder", make_test_H(), params); + std::vector syndrome(4); + auto r = d->decode(syndrome); + decode_ok = (r.result.size() == 10); + int cur = -1; + observed_device = (cudaGetDevice(&cur) == cudaSuccess) ? cur : -1; + }; + int dev0 = -1, dev1 = -1; + bool ok0 = false, ok1 = false; + std::thread t0(worker, 0, std::ref(dev0), std::ref(ok0)); + std::thread t1(worker, 1, std::ref(dev1), std::ref(ok1)); + t0.join(); + t1.join(); + EXPECT_TRUE(ok0); + EXPECT_TRUE(ok1); + EXPECT_EQ(dev0, 0); + EXPECT_EQ(dev1, 1); +} diff --git a/libs/qec/unittests/test_decoders_yaml.cpp b/libs/qec/unittests/test_decoders_yaml.cpp index dc8be3e41..1b3b0178a 100644 --- a/libs/qec/unittests/test_decoders_yaml.cpp +++ b/libs/qec/unittests/test_decoders_yaml.cpp @@ -720,3 +720,36 @@ TEST(DecoderConfigTest, SimulationHostPointerWrappersForwardToHostRuntime) { EXPECT_EQ(corrections, (std::vector{0})); finalize_decoders(); } + +TEST(DecoderYAMLTest, CudaDeviceIdRoundTrip) { + cudaq::qec::decoding::config::multi_decoder_config multi_config; + auto config = create_test_empty_decoder_config(0); + config.cuda_device_id = 2; + multi_config.decoders.push_back(config); + test_decoder_yaml_roundtrip(multi_config); +} + +TEST(DecoderYAMLTest, PrepareDecoderParamsSurfacesCudaDeviceId) { + // Non-trt type: the insert must happen before prepare_decoder_params()'s + // trt-only early return, so the knob reaches every decoder type. + auto config = create_test_empty_decoder_config(0); + config.cuda_device_id = 3; + auto params = cudaq::qec::decoding::host::prepare_decoder_params(config); + ASSERT_TRUE(params.contains("cuda_device_id")); + EXPECT_EQ(params.get("cuda_device_id"), 3); + + // Absent -> key absent (decoder::get() treats absence as unpinned). + auto config2 = create_test_empty_decoder_config(1); + auto params2 = cudaq::qec::decoding::host::prepare_decoder_params(config2); + EXPECT_FALSE(params2.contains("cuda_device_id")); + + // trt type: still surfaced on the trt branch. + auto config3 = create_test_empty_decoder_config(2); + config3.type = "trt_decoder"; + config3.decoder_custom_args = + cudaq::qec::decoding::config::trt_decoder_config{}; + config3.cuda_device_id = 1; + auto params3 = cudaq::qec::decoding::host::prepare_decoder_params(config3); + ASSERT_TRUE(params3.contains("cuda_device_id")); + EXPECT_EQ(params3.get("cuda_device_id"), 1); +}