19 #include <rmm/detail/export.hpp>
24 #include <cuda_runtime_api.h>
32 namespace RMM_NAMESPACE {
46 template <
typename Upstream>
58 : upstream_mr_{upstream_resource}
90 int8_t min_size_exponent,
91 int8_t max_size_exponent)
92 : upstream_mr_{upstream_resource}
94 for (
auto i = min_size_exponent; i <= max_size_exponent; i++) {
114 int8_t min_size_exponent,
115 int8_t max_size_exponent)
118 for (
auto i = min_size_exponent; i <= max_size_exponent; i++) {
159 std::optional<device_async_resource_ref> bin_resource = std::nullopt)
163 if (bin_resource.has_value()) {
164 resource_bins_.insert({allocation_size, bin_resource.value()});
165 }
else if (resource_bins_.count(allocation_size) == 0) {
166 owned_bin_resources_.push_back(
168 resource_bins_.insert({allocation_size, owned_bin_resources_.back().get()});
183 auto iter = resource_bins_.lower_bound(bytes);
184 return (iter != resource_bins_.cend()) ? iter->second : get_upstream_resource();
198 if (bytes <= 0) {
return nullptr; }
199 return get_resource_ref(bytes).allocate_async(bytes, stream);
210 void do_deallocate(
void* ptr, std::size_t bytes, cuda_stream_view stream)
override
212 get_resource_ref(bytes).deallocate_async(ptr, bytes, stream);
218 std::vector<std::unique_ptr<fixed_size_memory_resource<Upstream>>> owned_bin_resources_;
220 std::map<std::size_t, device_async_resource_ref> resource_bins_;
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Allocates memory from upstream resources associated with bin sizes.
Definition: binning_memory_resource.hpp:47
binning_memory_resource(Upstream *upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:72
device_async_resource_ref get_upstream_resource() const noexcept
device_async_resource_ref to the upstream resource
Definition: binning_memory_resource.hpp:138
binning_memory_resource(device_async_resource_ref upstream_resource, int8_t min_size_exponent, int8_t max_size_exponent)
Construct a new binning memory resource object with a range of initial bins.
Definition: binning_memory_resource.hpp:89
~binning_memory_resource() override=default
Destroy the binning_memory_resource and free all memory allocated from the upstream resource.
binning_memory_resource(Upstream *upstream_resource, int8_t min_size_exponent, int8_t max_size_exponent)
Construct a new binning memory resource object with a range of initial bins.
Definition: binning_memory_resource.hpp:113
binning_memory_resource(device_async_resource_ref upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:57
void add_bin(std::size_t allocation_size, std::optional< device_async_resource_ref > bin_resource=std::nullopt)
Add a bin allocator to this resource.
Definition: binning_memory_resource.hpp:158
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
A device_memory_resource which allocates memory blocks of a single fixed size.
Definition: fixed_size_memory_resource.hpp:53
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 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