Skip to content

Store any_resource members in polymorphic_allocator, thrust_allocator, and device_check_resource_adaptor #2340

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
bdice merged 1 commit into from
Mar 31, 2026
Merged

Store any_resource members in polymorphic_allocator, thrust_allocator, and device_check_resource_adaptor #2340

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 8 additions & 6 deletions cpp/include/rmm/mr/polymorphic_allocator.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/memory_resource>

#include <cstddef>
#include <memory>

Expand Down Expand Up @@ -52,7 +54,7 @@ class polymorphic_allocator {
*
* @param mr The upstream memory resource to use for allocation.
*/
polymorphic_allocator(device_async_resource_ref mr) : mr_{mr} {}
polymorphic_allocator(device_async_resource_ref mr) : mr_(mr) {}

/**
* @brief Construct a `polymorphic_allocator` using the underlying memory resource of `other`.
Expand All @@ -62,7 +64,7 @@ class polymorphic_allocator {
*/
template <typename U>
polymorphic_allocator(polymorphic_allocator<U> const& other) noexcept
: mr_{other.get_upstream_resource()}
: mr_(other.get_upstream_resource())
{
}

Expand All @@ -75,7 +77,7 @@ class polymorphic_allocator {
*/
value_type* allocate(std::size_t num, cuda_stream_view stream)
{
return static_cast<value_type*>(get_upstream_resource().allocate(stream, num * sizeof(T)));
return static_cast<value_type*>(mr_.allocate(stream, num * sizeof(T)));
}

/**
Expand All @@ -90,19 +92,19 @@ class polymorphic_allocator {
*/
void deallocate(value_type* ptr, std::size_t num, cuda_stream_view stream) noexcept
{
get_upstream_resource().deallocate(stream, ptr, num * sizeof(T));
mr_.deallocate(stream, ptr, num * sizeof(T));
}

/**
* @briefreturn{rmm::device_async_resource_ref to the upstream resource}
*/
[[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
{
return mr_;
return rmm::device_async_resource_ref{mr_};
}

private:
rmm::device_async_resource_ref mr_{
mutable cuda::mr::any_resource<cuda::mr::device_accessible> mr_{
get_current_device_resource_ref()}; ///< Underlying resource used for (de)allocation
};

Expand Down
6 changes: 4 additions & 2 deletions cpp/include/rmm/mr/thrust_allocator_adaptor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cuda/memory_resource>
#include <thrust/device_malloc_allocator.h>
#include <thrust/device_ptr.h>
#include <thrust/memory.h>
Expand Down Expand Up @@ -120,7 +121,7 @@ class thrust_allocator : public thrust::device_malloc_allocator<T> {
*/
[[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
{
return _mr;
return rmm::device_async_resource_ref{_mr};
}

/**
Expand All @@ -140,7 +141,8 @@ class thrust_allocator : public thrust::device_malloc_allocator<T> {

private:
cuda_stream_view _stream{};
rmm::device_async_resource_ref _mr{rmm::mr::get_current_device_resource_ref()};
mutable cuda::mr::any_resource<cuda::mr::device_accessible> _mr{
rmm::mr::get_current_device_resource_ref()};
cuda_device_id _device{get_current_cuda_device()};
};
/** @} */ // end of group
Expand Down
5 changes: 3 additions & 2 deletions cpp/tests/device_check_resource_adaptor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#include <gtest/gtest.h>

#include <cstddef>
#include <utility>

class device_check_resource_adaptor final {
public:
Expand All @@ -31,7 +32,7 @@ class device_check_resource_adaptor final {
*/
[[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
{
return upstream_;
return rmm::device_async_resource_ref{upstream_};
}

void* allocate(cuda::stream_ref stream,
Expand Down Expand Up @@ -88,7 +89,7 @@ class device_check_resource_adaptor final {
[[nodiscard]] bool check_device_id() const { return device_id == rmm::get_current_cuda_device(); }

rmm::cuda_device_id device_id;
rmm::device_async_resource_ref upstream_;
mutable cuda::mr::any_resource<cuda::mr::device_accessible> upstream_;
};

static_assert(cuda::mr::resource_with<device_check_resource_adaptor, cuda::mr::device_accessible>);
Loading