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/core/include/cuda-qx/core/kwargs_utils.h b/libs/core/include/cuda-qx/core/kwargs_utils.h index 784896b20..2297e5955 100644 --- a/libs/core/include/cuda-qx/core/kwargs_utils.h +++ b/libs/core/include/cuda-qx/core/kwargs_utils.h @@ -1,5 +1,5 @@ /****************************************************************-*- C++ -*-**** - * Copyright (c) 2024 - 2025 NVIDIA Corporation & Affiliates. * + * Copyright (c) 2024 - 2026 NVIDIA Corporation & Affiliates. * * All rights reserved. * * * * This source code and the accompanying materials are made available under * @@ -13,6 +13,7 @@ #include #include #include +#include namespace nb = nanobind; @@ -41,7 +42,14 @@ inline heterogeneous_map hetMapFromKwargs(const nb::kwargs &kwargs) { if (nb::isinstance(value)) { result.insert(key, nb::cast(value)); } else if (nb::isinstance(value)) { - result.insert(key, nb::cast(value)); + std::size_t integer_value = 0; + try { + integer_value = nb::cast(value); + } catch (...) { + throw std::runtime_error("Integer keyword argument '" + key + + "' must be non-negative and fit in size_t"); + } + result.insert(key, integer_value); } else if (nb::isinstance(value)) { result.insert(key, nb::cast(value)); } else if (nb::isinstance(value)) { diff --git a/libs/qec/include/cudaq/qec/decoder.h b/libs/qec/include/cudaq/qec/decoder.h index 0a2f4c0c0..a1dd92e9e 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; + /// @brief Set the observable matrix. void set_O_sparse(const std::vector> &O_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..fc13bcfac 100644 --- a/libs/qec/lib/decoder.cpp +++ b/libs/qec/lib/decoder.cpp @@ -8,13 +8,19 @@ #include "cudaq/qec/decoder.h" #include "cuda-qx/core/library_utils.h" +#include "hardware_guards.h" #include "cudaq/qec/logger.h" #include "cudaq/qec/plugin_loader.h" #include "cudaq/qec/version.h" +#include #include +#include +#include #include #include #include +#include +#include #include INSTANTIATE_REGISTRY(cudaq::qec::decoder, const cudaq::qec::decoder_init &, @@ -54,6 +60,11 @@ struct decoder::rt_impl { /// The id of the decoder (for instrumentation) uint32_t decoder_id = 0; + /// CUDA device selected by decoder::get(), or -1 when no placement was + /// requested. Keep this in the existing pimpl so adding placement does not + /// change the public decoder object layout used by external plugins. + int cuda_device_id = -1; + bool is_sliding_window = false; /// The number of syndromes per round. Only used for sliding window decoder. @@ -121,12 +132,139 @@ std::string decoder::get_version() const { return ss.str(); } +int decoder::get_cuda_device_id() const { return pimpl->cuda_device_id; } + 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 = get_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). +template +static bool try_read_cuda_device_id(const std::any &raw_value, int &result) { + const auto *value = std::any_cast(&raw_value); + if (!value) + return false; + + if constexpr (std::is_signed_v) { + if (*value < 0) + throw std::runtime_error( + "cuda_device_id must be a non-negative integer (got " + + std::to_string(*value) + ")"); + } + + using unsigned_integer = std::make_unsigned_t; + const auto magnitude = + static_cast(static_cast(*value)); + if (magnitude > static_cast(std::numeric_limits::max())) + throw std::runtime_error("cuda_device_id is too large (got " + + std::to_string(*value) + + "); maximum supported value is " + + std::to_string(std::numeric_limits::max())); + + result = static_cast(magnitude); + return true; +} + +static int read_cuda_device_id(const cudaqx::heterogeneous_map ¶ms) { + if (!params.contains("cuda_device_id")) + return -1; + + const std::any *raw_value = nullptr; + for (const auto &[key, value] : params) { + if (key == "cuda_device_id") { + raw_value = &value; + break; + } + } + if (!raw_value) + throw std::runtime_error("cuda_device_id is missing from parameter map"); + + int value = -1; + const bool is_integer = + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value) || + try_read_cuda_device_id(*raw_value, value); + if (!is_integer) + throw std::runtime_error("cuda_device_id must be an integer"); + + int count = 0; + const cudaError_t count_status = cudaGetDeviceCount(&count); + if (count_status != cudaSuccess) + throw std::runtime_error( + "cuda_device_id " + std::to_string(value) + + " could not be validated because cudaGetDeviceCount() failed: " + + cudaGetErrorString(count_status)); + if (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; +} + +/// Selects a device for decoder construction and restores the previous device +/// unless commit() is called. This keeps the successful persistent-pin +/// contract while making failed plugin construction transactional. +class ConstructionDevicePin { +public: + explicit ConstructionDevicePin(int target) : target_(target) { + cudaError_t err = cudaGetDevice(&previous_); + if (err != cudaSuccess) + throw std::runtime_error( + "cuda_device_id " + std::to_string(target_) + + " could not be selected because cudaGetDevice() failed: " + + cudaGetErrorString(err)); + + err = cudaSetDevice(target_); + if (err != cudaSuccess) + throw std::runtime_error( + "cudaSetDevice(" + std::to_string(target_) + + ") failed for cuda_device_id: " + cudaGetErrorString(err)); + selected_ = true; + } + + ~ConstructionDevicePin() { + if (selected_ && !committed_ && previous_ != target_) + (void)cudaSetDevice(previous_); + } + + ConstructionDevicePin(const ConstructionDevicePin &) = delete; + ConstructionDevicePin &operator=(const ConstructionDevicePin &) = delete; + + void commit() { + const cudaError_t err = cudaSetDevice(target_); + if (err != cudaSuccess) + throw std::runtime_error( + "cudaSetDevice(" + std::to_string(target_) + + ") failed while finalizing cuda_device_id placement: " + + cudaGetErrorString(err)); + committed_ = true; + } + +private: + int target_ = -1; + int previous_ = -1; + bool selected_ = false; + bool committed_ = false; +}; + std::unique_ptr decoder::get(const std::string &name, const decoder_init &init, const cudaqx::heterogeneous_map ¶m_map) { @@ -138,7 +276,24 @@ 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. + ConstructionDevicePin device_pin(cuda_device_id); + // 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->pimpl->cuda_device_id = cuda_device_id; + device_pin.commit(); + 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..8c3922467 --- /dev/null +++ b/libs/qec/lib/hardware_guards.h @@ -0,0 +1,61 @@ +/******************************************************************************* + * 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 it can be +/// reused by decoder implementation files without adding a public API. +/// +/// 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; + const cudaError_t count_status = cudaGetDeviceCount(&count); + if (count_status != cudaSuccess) + throw std::runtime_error( + "CudaDeviceGuard: cudaGetDeviceCount() failed for cuda_device_id " + + std::to_string(target) + ": " + cudaGetErrorString(count_status)); + if (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/decoding-server-cqr/DecodingSession.cpp b/libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp index 9e2e7b09c..e8ab4836f 100644 --- a/libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp +++ b/libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp @@ -46,6 +46,10 @@ DecodingSession::create(std::unique_ptr decoder, SyndromeMappingTable mapping_table_arg) { if (!decoder) throw std::invalid_argument("DecodingSession requires a decoder"); + if (decoder->get_cuda_device_id() >= 0) + throw std::invalid_argument( + "DecodingSession does not yet support cuda_device_id; worker-thread " + "affinity and device-owned teardown will be added in a follow-up"); auto s = std::make_unique(); s->dec = std::move(decoder); diff --git a/libs/qec/lib/realtime/qec_realtime_session.cpp b/libs/qec/lib/realtime/qec_realtime_session.cpp index ecfc09daa..b301f31f0 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 @@ -297,6 +298,11 @@ void qec_realtime_session::classify_mode() { if (!decoder) continue; ++non_null; + if (decoder->get_cuda_device_id() >= 0) + throw std::runtime_error( + "qec_realtime_session does not yet support cuda_device_id; " + "construct realtime decoders without device placement until " + "session worker affinity is implemented"); if (decoder->supports_graph_dispatch()) any_graph = true; else diff --git a/libs/qec/python/tests/test_decoder.py b/libs/qec/python/tests/test_decoder.py index 3201de1c6..00cbe37ff 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() + with pytest.raises(RuntimeError, match="cuda_device_id"): + qec.get_decoder("single_error_lut", H, cuda_device_id=-2) + with pytest.raises(RuntimeError, match="cuda_device_id"): + qec.get_decoder("single_error_lut", H, cuda_device_id=1 << 20) + # This value used to narrow to zero on platforms with a 32-bit C++ int. + with pytest.raises(RuntimeError, match="cuda_device_id.*too large"): + qec.get_decoder("single_error_lut", H, cuda_device_id=1 << 32) + with pytest.raises(RuntimeError, match="cuda_device_id.*integer"): + qec.get_decoder("single_error_lut", H, cuda_device_id=True) + + +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) or "cudaGetDeviceCount" 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 42132a4a7..5c0a2ab05 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..16055ce14 100644 --- a/libs/qec/unittests/test_decoders.cpp +++ b/libs/qec/unittests/test_decoders.cpp @@ -10,12 +10,16 @@ #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 +#include namespace { class ScopedEnv { @@ -1190,3 +1194,255 @@ 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, OversizedIntegerThrowsBeforeNarrowing) { + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", std::numeric_limits::max()); + 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); + EXPECT_NE(std::string(e.what()).find("too large"), std::string::npos); + } +} + +TEST(DecoderCudaDeviceId, NonIntegerThrows) { + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", true); + 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); + EXPECT_NE(std::string(e.what()).find("integer"), 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, ConstructionFailureRestoresPreviousDevice) { + if (cuda_device_count() < 2) + GTEST_SKIP() << "needs >= 2 CUDA devices"; + ScopedDeviceRestore restore; + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", 1); + params.insert("bogus_key", 1); // strict_keys_decoder throws in its ctor. + EXPECT_THROW( + cudaq::qec::decoder::get("strict_keys_decoder", make_test_H(), params), + std::runtime_error); + + int current = -1; + ASSERT_EQ(cudaGetDevice(¤t), cudaSuccess); + EXPECT_EQ(current, 0); +} + +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 026754649..73bcff15a 100644 --- a/libs/qec/unittests/test_decoders_yaml.cpp +++ b/libs/qec/unittests/test_decoders_yaml.cpp @@ -84,6 +84,24 @@ unexpected: true misspelled_decoder_argument), std::runtime_error); + // Device placement is deliberately limited to direct decoder construction + // until realtime session workers and teardown own a CUDA device. + const std::string unsupported_cuda_device_id = R"( +decoders: + - id: 0 + type: pymatching + cuda_device_id: 0 + block_size: 1 + syndrome_size: 1 + H_sparse: [0, -1] + O_sparse: [0, -1] + D_sparse: [0, -1] +)"; + EXPECT_THROW( + cudaq::qec::decoding::config::multi_decoder_config::from_yaml_str( + unsupported_cuda_device_id), + std::runtime_error); + EXPECT_THROW( cudaq::qec::decoding::config::multi_decoder_config::from_yaml_str( "decoders: ["),