Skip to content

Store any_resource in device_buffer and device_uvector #2201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 6 commits into from
Jan 12, 2026
Merged

Store any_resource in device_buffer and device_uvector #2201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fc0f3c0
Store any_resource in device-resource global mapping
bdice Jan 7, 2026
e791a1e
Use const&
bdice Jan 8, 2026
1c3dad9
Revert "Use const&"
bdice Jan 8, 2026
92193f8
Store any_resource in device_buffer
bdice Jan 8, 2026
6e7e098
Add test
bdice Jan 8, 2026
71ba0aa
Merge branch 'main' into containers-any-resource
bdice Jan 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions cpp/include/rmm/device_buffer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/
#pragma once
Expand All @@ -11,6 +11,7 @@
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/memory_resource>
#include <cuda_runtime_api.h>

#include <cassert>
Expand Down Expand Up @@ -320,17 +321,16 @@ class device_buffer {
/**
* @briefreturn{The resource used to allocate and deallocate}
*/
[[nodiscard]] rmm::device_async_resource_ref memory_resource() const noexcept { return _mr; }
[[nodiscard]] rmm::device_async_resource_ref memory_resource() noexcept { return _mr; }
Copy link
Copy Markdown
Collaborator Author

@bdice bdice Jan 10, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a breaking change. Should we make the resource mutable or make this method non-const? I discussed a bit with the CCCL team.

The CCCL buffer returns any_resource const& which requires a copy to be made before the resource (or ref) is usable (constructing a ref and making an allocation are not possible from a const resource). I would consider that approach but I don’t want to alter the API in that way just yet. It might be better to handle that as we transition to the CCCL buffer type.

Copy link
Copy Markdown
Collaborator Author

@bdice bdice Jan 13, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After merging this, I decided to go the other way and use mutable in PR #2208 to keep this const. That seems more aligned with the spirit of how we want to use memory resources. See commit here: 47d7bf4.


private:
void* _data{nullptr}; ///< Pointer to device memory allocation
std::size_t _size{}; ///< Requested size of the device memory allocation
std::size_t _capacity{}; ///< The actual size of the device memory allocation
cuda_stream_view _stream{}; ///< Stream to use for device memory deallocation

rmm::device_async_resource_ref _mr{
rmm::mr::get_current_device_resource_ref()}; ///< The memory resource used to
///< allocate/deallocate device memory
cuda::mr::any_resource<cuda::mr::device_accessible> _mr; ///< The memory resource used to
///< allocate/deallocate device memory
cuda_device_id _device{get_current_cuda_device()};

/**
Expand Down
4 changes: 2 additions & 2 deletions cpp/include/rmm/device_uvector.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -602,7 +602,7 @@ class device_uvector {
* @briefreturn{The resource used to allocate and deallocate the device
* storage}
*/
[[nodiscard]] rmm::device_async_resource_ref memory_resource() const noexcept
[[nodiscard]] rmm::device_async_resource_ref memory_resource() noexcept
{
return _storage.memory_resource();
}
Expand Down
6 changes: 3 additions & 3 deletions cpp/src/device_buffer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -44,7 +44,7 @@ device_buffer::device_buffer(device_buffer&& other) noexcept
_size{other._size},
_capacity{other._capacity},
_stream{other.stream()},
_mr{other._mr},
_mr{std::move(other._mr)},
_device{other._device}
{
other._data = nullptr;
Expand All @@ -64,7 +64,7 @@ device_buffer& device_buffer::operator=(device_buffer&& other) noexcept
_size = other._size;
_capacity = other._capacity;
set_stream(other.stream());
_mr = other._mr;
_mr = std::move(other._mr);
_device = other._device;

other._data = nullptr;
Expand Down
30 changes: 29 additions & 1 deletion cpp/tests/mr/mr_ref_default_tests.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -63,6 +63,34 @@ TEST(DefaultTest, GetCurrentDeviceResourceRef)
EXPECT_EQ(mr, rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()});
}

TEST(DefaultTest, SetCurrentDeviceResourceRefFromPointer)
Copy link
Copy Markdown
Collaborator Author

@bdice bdice Jan 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is an extra test I wanted to add after #2200. Not buffer related.

{
// Construct a cuda_memory_resource
rmm::mr::cuda_memory_resource cuda_mr{};

// Get a pointer to it (device_memory_resource*)
rmm::mr::device_memory_resource* mr_ptr = &cuda_mr;

// Set with set_current_device_resource_ref using the pointer
rmm::mr::set_current_device_resource_ref(mr_ptr);

// Get the ref with get_current_device_resource_ref
auto ref = rmm::mr::get_current_device_resource_ref();

// Use that ref to allocate
constexpr std::size_t size{1024};
void* ptr = ref.allocate_sync(size);
EXPECT_NE(ptr, nullptr);
EXPECT_TRUE(is_properly_aligned(ptr));
EXPECT_TRUE(is_device_accessible_memory(ptr));

// Deallocate
ref.deallocate_sync(ptr, size);

// Reset to initial resource
rmm::mr::reset_current_device_resource_ref();
}

// Multi-threaded default resource tests

TEST(DefaultTest, UseCurrentDeviceResource_mt) { spawn(test_get_current_device_resource); }
Expand Down