diff --git a/cpp/bench/prims/common/benchmark.hpp b/cpp/bench/prims/common/benchmark.hpp index 913729fd03..5889a2c3bf 100644 --- a/cpp/bench/prims/common/benchmark.hpp +++ b/cpp/bench/prims/common/benchmark.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -20,6 +20,7 @@ #include #include #include +#include #include @@ -33,26 +34,26 @@ namespace raft::bench { */ struct using_pool_memory_res { private: - rmm::mr::device_memory_resource* orig_res_; + rmm::device_async_resource_ref orig_res_; rmm::mr::cuda_memory_resource cuda_res_{}; rmm::mr::pool_memory_resource pool_res_; public: using_pool_memory_res(size_t initial_size, size_t max_size) - : orig_res_(rmm::mr::get_current_device_resource()), + : orig_res_(rmm::mr::get_current_device_resource_ref()), pool_res_(&cuda_res_, initial_size, max_size) { - rmm::mr::set_current_device_resource(&pool_res_); + rmm::mr::set_current_device_resource_ref(&pool_res_); } using_pool_memory_res() - : orig_res_(rmm::mr::get_current_device_resource()), + : orig_res_(rmm::mr::get_current_device_resource_ref()), pool_res_(&cuda_res_, rmm::percent_of_free_device_memory(50)) { - rmm::mr::set_current_device_resource(&pool_res_); + rmm::mr::set_current_device_resource_ref(&pool_res_); } - ~using_pool_memory_res() { rmm::mr::set_current_device_resource(orig_res_); } + ~using_pool_memory_res() { rmm::mr::set_current_device_resource_ref(orig_res_); } }; /** diff --git a/cpp/bench/prims/matrix/gather.cu b/cpp/bench/prims/matrix/gather.cu index 0f9fdf4c34..ce465cc15f 100644 --- a/cpp/bench/prims/matrix/gather.cu +++ b/cpp/bench/prims/matrix/gather.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,7 +14,9 @@ #include #include +#include #include +#include namespace raft::bench::matrix { @@ -35,7 +37,7 @@ template struct Gather : public fixture { Gather(const GatherParams& p) : params(p), - old_mr(rmm::mr::get_current_device_resource()), + old_mr(rmm::mr::get_current_device_resource_ref()), pool_mr(rmm::mr::get_current_device_resource(), 2 * (1ULL << 30)), matrix(this->handle), map(this->handle), @@ -43,10 +45,10 @@ struct Gather : public fixture { stencil(this->handle), matrix_h(this->handle) { - rmm::mr::set_current_device_resource(&pool_mr); + rmm::mr::set_current_device_resource_ref(&pool_mr); } - ~Gather() { rmm::mr::set_current_device_resource(old_mr); } + ~Gather() { rmm::mr::set_current_device_resource_ref(old_mr); } void allocate_data(const ::benchmark::State& state) override { @@ -107,7 +109,7 @@ struct Gather : public fixture { private: GatherParams params; - rmm::mr::device_memory_resource* old_mr; + rmm::device_async_resource_ref old_mr; rmm::mr::pool_memory_resource pool_mr; raft::device_matrix matrix, out; raft::host_matrix matrix_h; diff --git a/cpp/bench/prims/random/subsample.cu b/cpp/bench/prims/random/subsample.cu index 004b940f7f..63560d4832 100644 --- a/cpp/bench/prims/random/subsample.cu +++ b/cpp/bench/prims/random/subsample.cu @@ -17,6 +17,7 @@ #include #include #include +#include namespace raft::bench::random { @@ -50,16 +51,16 @@ template struct sample : public fixture { sample(const sample_inputs& p) : params(p), - old_mr(rmm::mr::get_current_device_resource()), + old_mr(rmm::mr::get_current_device_resource_ref()), pool_mr(rmm::mr::get_current_device_resource(), 2 * GiB), in(make_device_vector(res, p.n_samples)), out(make_device_vector(res, p.n_train)) { - rmm::mr::set_current_device_resource(&pool_mr); + rmm::mr::set_current_device_resource_ref(&pool_mr); raft::random::RngState r(123456ULL); } - ~sample() { rmm::mr::set_current_device_resource(old_mr); } + ~sample() { rmm::mr::set_current_device_resource_ref(old_mr); } void run_benchmark(::benchmark::State& state) override { std::ostringstream label_stream; @@ -81,7 +82,7 @@ struct sample : public fixture { private: float GiB = 1073741824.0f; raft::device_resources res; - rmm::mr::device_memory_resource* old_mr; + rmm::device_async_resource_ref old_mr; rmm::mr::pool_memory_resource pool_mr; sample_inputs params; raft::device_vector out, in; diff --git a/cpp/include/raft/core/device_resources_manager.hpp b/cpp/include/raft/core/device_resources_manager.hpp index 7d615cec5f..7acf3bd744 100644 --- a/cpp/include/raft/core/device_resources_manager.hpp +++ b/cpp/include/raft/core/device_resources_manager.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -165,7 +165,7 @@ struct device_resources_manager { upstream, params.init_mem_pool_size.value_or(rmm::percent_of_free_device_memory(50)), params.max_mem_pool_size); - rmm::mr::set_current_device_resource(result.get()); + rmm::mr::set_current_device_resource_ref(result.get()); } else { RAFT_LOG_WARN( "Pool allocation requested, but other memory resource has already been set and " diff --git a/cpp/include/raft/random/detail/multi_variable_gaussian.cuh b/cpp/include/raft/random/detail/multi_variable_gaussian.cuh index 0cec08e6e3..1081b5462a 100644 --- a/cpp/include/raft/random/detail/multi_variable_gaussian.cuh +++ b/cpp/include/raft/random/detail/multi_variable_gaussian.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -22,6 +22,8 @@ #include #include +#include + #include #include #include @@ -365,7 +367,7 @@ class multi_variable_gaussian_setup_token { private: std::unique_ptr> impl_; raft::resources const& handle_; - rmm::device_async_resource_ref mem_resource_; + mutable cuda::mr::any_resource mem_resource_; int dim_ = 0; auto allocate_workspace() const diff --git a/cpp/tests/mr/device/buffer.cpp b/cpp/tests/mr/device/buffer.cpp index 77106b9f04..09e7b76043 100644 --- a/cpp/tests/mr/device/buffer.cpp +++ b/cpp/tests/mr/device/buffer.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -53,7 +53,7 @@ TEST(Raft, DeviceBufferZeroResize) std::make_shared>(curr_mr, 1000); - rmm::mr::set_current_device_resource(limit_mr.get()); + rmm::mr::set_current_device_resource_ref(limit_mr.get()); cudaStream_t stream; RAFT_CUDA_TRY(cudaStreamCreate(&stream)); @@ -73,7 +73,7 @@ TEST(Raft, DeviceBufferZeroResize) // Now check that there is no memory left. (Used to not be true) ASSERT_EQ(0, limit_mr->get_allocated_bytes()); - rmm::mr::set_current_device_resource(curr_mr); + rmm::mr::set_current_device_resource_ref(curr_mr); RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); RAFT_CUDA_TRY(cudaStreamDestroy(stream));