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/lib/realtime/decoding-server-cqr/SessionRegistry.cpp b/libs/qec/lib/realtime/decoding-server-cqr/SessionRegistry.cpp index 718929fb1..6afa6b365 100644 --- a/libs/qec/lib/realtime/decoding-server-cqr/SessionRegistry.cpp +++ b/libs/qec/lib/realtime/decoding-server-cqr/SessionRegistry.cpp @@ -7,6 +7,7 @@ ******************************************************************************/ #include "SessionRegistry.h" +#include "../../hardware_guards.h" #include "../realtime_decoding.h" #include "cudaq/qec/logger.h" #include "cudaq/qec/realtime/decoding_config.h" @@ -71,6 +72,11 @@ void SessionRegistry::load_from_config(const multi_decoder_config &config, CUDA_QEC_INFO("SessionRegistry: creating decoder id={} type={}", dc.id, dc.type); + // Keep each construction transaction on its configured device, then + // restore the registry thread. Runtime ownership transfers to the + // session's dedicated worker in start_worker(). + cudaq::qec::detail_affinity::CudaDeviceGuard construction_device( + dc.cuda_device_id.value_or(-1)); auto decoder = cudaq::qec::decoding::host::create_realtime_decoder(dc); auto session = DecodingSession::create(std::move(decoder), make_default_mapping_table()); diff --git a/libs/qec/lib/realtime/qec_realtime_session.h b/libs/qec/lib/realtime/qec_realtime_session.h index fef32a6bc..1d0383d00 100644 --- a/libs/qec/lib/realtime/qec_realtime_session.h +++ b/libs/qec/lib/realtime/qec_realtime_session.h @@ -187,6 +187,10 @@ class __attribute__((visibility("default"))) qec_realtime_session { // ---- Lifetime / mode ---- bool initialized_ = false; bool device_mode_ = false; + // Every DEVICE-mode decoder must share one CUDA owner, so each scheduler + // allocation, graph launch, and cleanup operation belongs to this device. + // -1 means none of the decoders requested explicit placement. + int device_mode_cuda_device_id_ = -1; // ---- Ring buffer (raw pointers; _dev aliases _host in HOST mode) ---- static constexpr std::size_t kDefaultNumSlots = 8; diff --git a/libs/qec/python/tests/test_decoders_yaml.py b/libs/qec/python/tests/test_decoders_yaml.py index a6e017278..acc8d0464 100644 --- a/libs/qec/python/tests/test_decoders_yaml.py +++ b/libs/qec/python/tests/test_decoders_yaml.py @@ -128,6 +128,19 @@ def create_test_decoder_config_nv_qldpc(decoder_id): return config +def test_cuda_device_id_yaml_roundtrip(): + multi_config = qec.multi_decoder_config() + config = create_test_empty_decoder_config(0) + assert config.cuda_device_id is None + config.cuda_device_id = 2 + multi_config.decoders = [config] + + check_decoder_yaml_roundtrip(multi_config) + parsed = qec.multi_decoder_config.from_yaml_str( + multi_config.to_yaml_str(200)) + assert parsed.decoders[0].cuda_device_id == 2 + + def test_single_decoder(): """ Test YAML serialization/deserialization and creation of a single NV-QLDPC decoder. diff --git a/libs/qec/unittests/CMakeLists.txt b/libs/qec/unittests/CMakeLists.txt index 359258a51..cd18b40cc 100644 --- a/libs/qec/unittests/CMakeLists.txt +++ b/libs/qec/unittests/CMakeLists.txt @@ -48,7 +48,8 @@ target_link_libraries(test_decoders_yaml PRIVATE cudaq-qec-decoding-server cudaq-qec-realtime-decoding cudaq-qec-realtime-decoding-simulation - cudaq::cudaq) + cudaq::cudaq + CUDA::cudart) add_dependencies(CUDAQXQECUnitTests test_decoders_yaml) gtest_discover_tests(test_decoders_yaml) diff --git a/libs/qec/unittests/decoders/pymatching/CMakeLists.txt b/libs/qec/unittests/decoders/pymatching/CMakeLists.txt index cc8d5a0f8..f482469b8 100644 --- a/libs/qec/unittests/decoders/pymatching/CMakeLists.txt +++ b/libs/qec/unittests/decoders/pymatching/CMakeLists.txt @@ -42,6 +42,8 @@ if(CUDAQ_REALTIME_ROOT AND CUDAQ_REALTIME_INCLUDE_DIR) GTest::gtest_main cudaq-qec-decoders cudaq-qec-realtime-decoding + cudaq::cudaq + CUDA::cudart ) set_target_properties(test_pymatching_realtime PROPERTIES BUILD_RPATH "${CMAKE_BINARY_DIR}/lib;${CMAKE_BINARY_DIR}/lib/decoder-plugins" diff --git a/libs/qec/unittests/decoders/pymatching/test_pymatching_realtime.cpp b/libs/qec/unittests/decoders/pymatching/test_pymatching_realtime.cpp index c345fc95a..7700c21ba 100644 --- a/libs/qec/unittests/decoders/pymatching/test_pymatching_realtime.cpp +++ b/libs/qec/unittests/decoders/pymatching/test_pymatching_realtime.cpp @@ -12,8 +12,10 @@ #include "cudaq/qec/decoder.h" #include "cudaq/qec/realtime/decoding_config.h" +#include #include +#include #include #include #include @@ -24,6 +26,53 @@ namespace { using DecoderVec = std::vector>; +class RealtimeSessionDeviceRecordingDecoder final : public cudaq::qec::decoder { +public: + std::atomic last_decode_device{-2}; + + RealtimeSessionDeviceRecordingDecoder( + const cudaq::qec::sparse_binary_matrix &H, + const cudaqx::heterogeneous_map &) + : decoder(H) {} + + cudaq::qec::decoder_result + decode(const std::vector &syndrome) override { + int device = -1; + if (cudaGetDevice(&device) != cudaSuccess) + device = -1; + last_decode_device.store(device, std::memory_order_release); + cudaq::qec::decoder_result result; + result.converged = true; + result.result = {syndrome.at(0)}; + return result; + } + + CUDAQ_EXTENSION_CUSTOM_CREATOR_FUNCTION( + RealtimeSessionDeviceRecordingDecoder, + static std::unique_ptr create( + const cudaq::qec::decoder_init &init, + const cudaqx::heterogeneous_map ¶ms) { + return cudaq::qec::make_pcm_decoder< + RealtimeSessionDeviceRecordingDecoder>(init, params); + }) +}; +CUDAQ_EXT_PT_REGISTER_TYPE(RealtimeSessionDeviceRecordingDecoder) + +class ScopedDeviceRestore { +public: + ScopedDeviceRestore() { + if (cudaGetDevice(&previous) != cudaSuccess) + previous = -1; + } + ~ScopedDeviceRestore() { + if (previous >= 0) + (void)cudaSetDevice(previous); + } + +private: + int previous = -1; +}; + DecoderVec make_pymatching_decoders(const std::vector &h_vec, std::size_t syndrome_size, std::size_t block_size) { @@ -208,3 +257,38 @@ TEST(PyMatchingRealtime, ConfiguresViaRealtimeDecoderConfig) { config::finalize_decoders(); } + +TEST(RealtimeSessionCudaDeviceId, HostDispatcherUsesConfiguredDevice) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count < 2) + GTEST_SKIP() << "needs >= 2 CUDA devices"; + + ScopedDeviceRestore restore; + cudaqx::tensor H({std::size_t{1}, std::size_t{1}}); + H.at({0, 0}) = 1; + cudaqx::heterogeneous_map params; + params.insert("cuda_device_id", 1); + auto decoder = cudaq::qec::decoder::get( + "RealtimeSessionDeviceRecordingDecoder", H, params); + auto *recording = + dynamic_cast(decoder.get()); + ASSERT_NE(recording, nullptr); + decoder->set_decoder_id(0); + decoder->set_D_sparse(std::vector>{{0}}); + decoder->set_O_sparse(std::vector>{{0}}); + + DecoderVec decoders; + decoders.push_back(std::move(decoder)); + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + cudaq::qec::realtime::qec_realtime_session session(decoders); + session.initialize(); + + expect_corrections(session, {1}, std::vector{1}, + /*counter=*/1); + EXPECT_EQ(recording->last_decode_device.load(std::memory_order_acquire), 1); + + int current = -1; + ASSERT_EQ(cudaGetDevice(¤t), cudaSuccess); + EXPECT_EQ(current, 0); + session.finalize(); +} diff --git a/libs/qec/unittests/realtime/qec_graph_decode_test/qldpc_config_loader.cpp b/libs/qec/unittests/realtime/qec_graph_decode_test/qldpc_config_loader.cpp index 921f812ed..489b63a6c 100644 --- a/libs/qec/unittests/realtime/qec_graph_decode_test/qldpc_config_loader.cpp +++ b/libs/qec/unittests/realtime/qec_graph_decode_test/qldpc_config_loader.cpp @@ -7,6 +7,7 @@ ******************************************************************************/ #include "qldpc_config_loader.h" +#include "realtime_decoding.h" #include "cudaq/qec/decoder.h" #include "cudaq/qec/realtime/decoding_config.h" @@ -64,7 +65,7 @@ LoadedDecoder load_decoder_from_yaml(const std::string &yaml_path) { for (std::uint32_t j = h_row_ptr[r]; j < h_row_ptr[r + 1]; ++j) H_tensor.at({r, static_cast(h_col_idx[j])}) = 1; - auto params = dec.decoder_custom_args_to_heterogeneous_map(); + auto params = decoding::host::prepare_decoder_params(dec); auto plugin = decoder::get("nv-qldpc-decoder", H_tensor, params); if (!plugin) throw std::runtime_error( diff --git a/libs/qec/unittests/realtime/qec_graph_decode_test/test_realtime_qldpc_graph_decoding.cpp b/libs/qec/unittests/realtime/qec_graph_decode_test/test_realtime_qldpc_graph_decoding.cpp index 13c1d0b26..978254f47 100644 --- a/libs/qec/unittests/realtime/qec_graph_decode_test/test_realtime_qldpc_graph_decoding.cpp +++ b/libs/qec/unittests/realtime/qec_graph_decode_test/test_realtime_qldpc_graph_decoding.cpp @@ -254,6 +254,7 @@ class GraphDecodeTest : public ::testing::Test { num_measurements_ = loaded.num_measurements; num_observables_ = loaded.num_observables; ASSERT_NE(loaded.decoder, nullptr); + EXPECT_EQ(loaded.decoder->get_cuda_device_id(), 0); printf("Config: num_measurements=%zu, num_observables=%zu\n", num_measurements_, num_observables_); diff --git a/libs/qec/unittests/realtime/qec_roce_decode_test/data/config_nv_qldpc_relay.yml b/libs/qec/unittests/realtime/qec_roce_decode_test/data/config_nv_qldpc_relay.yml index c9c489a4a..024b2bc6b 100644 --- a/libs/qec/unittests/realtime/qec_roce_decode_test/data/config_nv_qldpc_relay.yml +++ b/libs/qec/unittests/realtime/qec_roce_decode_test/data/config_nv_qldpc_relay.yml @@ -2,6 +2,7 @@ decoders: - id: 0 type: nv-qldpc-decoder + cuda_device_id: 0 block_size: 69 syndrome_size: 24 H_sparse: [ 0, 1, 2, 3, 4, 5, 6, -1, 1, 4, 7, 8, -1, 2, 5, 9, 10, 11, 12, 13, -1, 10, 11, 14, 15, -1, 3, 6, 16, 17, -1, 4, 5, 8, 13, 18, 19, 20, -1, 5, 6, 11, 15, 17, 19, 21, -1, 12, 13, 20,