20 #include <rmm/detail/error.hpp>
21 #include <rmm/detail/export.hpp>
28 #include <unordered_map>
30 namespace RMM_NAMESPACE {
56 template <
typename Upstream>
71 std::size_t alignment_threshold = default_alignment_threshold)
72 : upstream_{upstream}, alignment_{alignment}, alignment_threshold_{alignment_threshold}
75 "Allocation alignment is not a power of 2.");
91 std::size_t alignment_threshold = default_alignment_threshold)
93 alignment_{alignment},
94 alignment_threshold_{alignment_threshold}
97 "Allocation alignment is not a power of 2.");
118 static constexpr std::size_t default_alignment_threshold = 0;
121 using lock_guard = std::lock_guard<std::mutex>;
137 return get_upstream_resource().allocate_async(bytes, 1, stream);
139 auto const size = upstream_allocation_size(bytes);
140 void* pointer = get_upstream_resource().allocate_async(size, 1, stream);
142 auto const address =
reinterpret_cast<std::size_t
>(pointer);
143 auto const aligned_address =
rmm::align_up(address, alignment_);
145 void* aligned_pointer =
reinterpret_cast<void*
>(aligned_address);
146 if (pointer != aligned_pointer) {
147 lock_guard lock(mtx_);
148 pointers_.emplace(aligned_pointer, pointer);
150 return aligned_pointer;
160 void do_deallocate(
void* ptr, std::size_t bytes,
cuda_stream_view stream)
override
163 get_upstream_resource().deallocate_async(ptr, bytes, 1, stream);
166 lock_guard lock(mtx_);
167 auto const iter = pointers_.find(ptr);
168 if (iter != pointers_.end()) {
170 pointers_.erase(iter);
173 get_upstream_resource().deallocate_async(ptr, upstream_allocation_size(bytes), 1, stream);
184 [[nodiscard]]
bool do_is_equal(device_memory_resource
const& other)
const noexcept
override
186 if (
this == &other) {
return true; }
187 auto cast =
dynamic_cast<aligned_resource_adaptor<Upstream> const*
>(&other);
188 if (cast ==
nullptr) {
return false; }
189 return get_upstream_resource() == cast->get_upstream_resource() &&
190 alignment_ == cast->alignment_ && alignment_threshold_ == cast->alignment_threshold_;
200 std::size_t upstream_allocation_size(std::size_t bytes)
const
208 std::unordered_map<void*, void*> pointers_;
209 std::size_t alignment_;
210 std::size_t alignment_threshold_;
211 mutable std::mutex mtx_;
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Resource that adapts Upstream memory resource to allocate memory in a specified alignment size.
Definition: aligned_resource_adaptor.hpp:57
aligned_resource_adaptor(device_async_resource_ref upstream, std::size_t alignment=rmm::CUDA_ALLOCATION_ALIGNMENT, std::size_t alignment_threshold=default_alignment_threshold)
Construct an aligned resource adaptor using upstream to satisfy allocation requests.
Definition: aligned_resource_adaptor.hpp:69
aligned_resource_adaptor(Upstream *upstream, std::size_t alignment=rmm::CUDA_ALLOCATION_ALIGNMENT, std::size_t alignment_threshold=default_alignment_threshold)
Construct an aligned resource adaptor using upstream to satisfy allocation requests.
Definition: aligned_resource_adaptor.hpp:89
rmm::device_async_resource_ref get_upstream_resource() const noexcept
rmm::device_async_resource_ref to the upstream resource
Definition: aligned_resource_adaptor.hpp:110
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
cuda::mr::async_resource_ref< cuda::mr::device_accessible > device_async_resource_ref
Alias for a cuda::mr::async_resource_ref with the property cuda::mr::device_accessible.
Definition: resource_ref.hpp:41
device_async_resource_ref to_device_async_resource_ref_checked(Resource *res)
Convert pointer to memory resource into device_async_resource_ref, checking for nullptr
Definition: resource_ref.hpp:79
static constexpr std::size_t CUDA_ALLOCATION_ALIGNMENT
Default alignment used for CUDA memory allocation.
Definition: aligned.hpp:43
constexpr bool is_supported_alignment(std::size_t alignment) noexcept
Returns whether or not alignment is a valid memory alignment.
Definition: aligned.hpp:64
constexpr std::size_t align_up(std::size_t value, std::size_t alignment) noexcept
Align up to nearest multiple of specified power of 2.
Definition: aligned.hpp:77
Management of per-device device_memory_resources.