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