From 58eff945e156a6d7a018c43bff89a721b81b6265 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 26 Aug 2026 13:10:11 +0200 Subject: [PATCH] Synchronize the stream on accesses to the device proxy --- cpp/include/raft/core/device_container_policy.hpp | 7 ++++++- cpp/include/raft/core/device_mdarray.hpp | 6 ++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/cpp/include/raft/core/device_container_policy.hpp b/cpp/include/raft/core/device_container_policy.hpp index 30233b69e6..c706c69d6c 100644 --- a/cpp/include/raft/core/device_container_policy.hpp +++ b/cpp/include/raft/core/device_container_policy.hpp @@ -1,6 +1,6 @@ /* * SPDX-FileCopyrightText: Copyright (2019) Sandia Corporation - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause */ /* @@ -30,6 +30,9 @@ namespace RAFT_EXPORT raft { /** * @brief A simplified version of thrust::device_reference with support for CUDA stream. + * + * @note This proxy performs H2D or D2H transfer and a synchronization on the given + * stream on every access. */ template class device_reference { @@ -53,12 +56,14 @@ class device_reference { auto* raw = ptr_.get(); value_type v{}; update_host(&v, raw, 1, stream_); + raft::interruptible::synchronize(stream_); return v; } auto operator=(T const& other) -> device_reference& { auto* raw = ptr_.get(); update_device(raw, &other, 1, stream_); + raft::interruptible::synchronize(stream_); return *this; } }; diff --git a/cpp/include/raft/core/device_mdarray.hpp b/cpp/include/raft/core/device_mdarray.hpp index f7f564283c..65d30c66f8 100644 --- a/cpp/include/raft/core/device_mdarray.hpp +++ b/cpp/include/raft/core/device_mdarray.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -149,7 +149,9 @@ auto make_device_scalar(raft::resources const& handle) /** * @brief Create a device scalar from v - * (async copy in the resource-provided stream). + * + * @note This function performs a H2D copy and implicitly synchronizes the resource's stream, + * which impacts performance when there's unfinished work in the stream. * * @tparam ElementType the data type of the scalar element * @tparam IndexType the index type of the extents