19 #include <rmm/detail/export.hpp>
24 #include <cuda_runtime_api.h>
33 namespace RMM_NAMESPACE {
47 template <
typename Upstream>
59 : upstream_mr_{upstream_resource}
91 int8_t min_size_exponent,
92 int8_t max_size_exponent)
93 : upstream_mr_{upstream_resource}
95 for (
auto i = min_size_exponent; i <= max_size_exponent; i++) {
115 int8_t min_size_exponent,
116 int8_t max_size_exponent)
119 for (
auto i = min_size_exponent; i <= max_size_exponent; i++) {
160 std::optional<device_async_resource_ref> bin_resource = std::nullopt)
164 if (bin_resource.has_value()) {
165 resource_bins_.insert({allocation_size, bin_resource.value()});
166 }
else if (resource_bins_.count(allocation_size) == 0) {
167 owned_bin_resources_.push_back(
169 resource_bins_.insert({allocation_size, owned_bin_resources_.back().get()});
184 auto iter = resource_bins_.lower_bound(bytes);
185 return (iter != resource_bins_.cend()) ? iter->second : get_upstream_resource();
199 if (bytes <= 0) {
return nullptr; }
200 return get_resource_ref(bytes).allocate_async(bytes, stream);
211 void do_deallocate(
void* ptr, std::size_t bytes, cuda_stream_view stream)
override
213 get_resource_ref(bytes).deallocate_async(ptr, bytes, stream);
219 std::vector<std::unique_ptr<fixed_size_memory_resource<Upstream>>> owned_bin_resources_;
221 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:41
Allocates memory from upstream resources associated with bin sizes.
Definition: binning_memory_resource.hpp:48
binning_memory_resource(Upstream *upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:73
device_async_resource_ref get_upstream_resource() const noexcept
device_async_resource_ref to the upstream resource
Definition: binning_memory_resource.hpp:139
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:90
~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:114
binning_memory_resource(device_async_resource_ref upstream_resource)
Construct a new binning memory resource object.
Definition: binning_memory_resource.hpp:58
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:159
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:94
A device_memory_resource which allocates memory blocks of a single fixed size.
Definition: fixed_size_memory_resource.hpp:55
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