binning_memory_resource.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2020-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <rmm/aligned.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/fixed_size_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda_runtime_api.h>
 
#include <cassert>
#include <map>
#include <memory>
#include <optional>
#include <vector>
 
namespace RMM_NAMESPACE {
namespace mr {
template <typename Upstream>
class binning_memory_resource final : public device_memory_resource {
 public:
  explicit binning_memory_resource(device_async_resource_ref upstream_resource)
    : upstream_mr_{upstream_resource}
  {
  }
 
  explicit binning_memory_resource(Upstream* upstream_resource)
    : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
  {
  }
 
  binning_memory_resource(device_async_resource_ref upstream_resource,
                          int8_t min_size_exponent,  // NOLINT(bugprone-easily-swappable-parameters)
                          int8_t max_size_exponent)
    : upstream_mr_{upstream_resource}
  {
    for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
      add_bin(1 << i);
    }
  }
 
  binning_memory_resource(Upstream* upstream_resource,
                          int8_t min_size_exponent,  // NOLINT(bugprone-easily-swappable-parameters)
                          int8_t max_size_exponent)
    : upstream_mr_{to_device_async_resource_ref_checked(upstream_resource)}
  {
    for (auto i = min_size_exponent; i <= max_size_exponent; i++) {
      add_bin(1 << i);
    }
  }
 
  ~binning_memory_resource() override = default;
 
  binning_memory_resource()                                          = delete;
  binning_memory_resource(binning_memory_resource const&)            = delete;
  binning_memory_resource(binning_memory_resource&&)                 = delete;
  binning_memory_resource& operator=(binning_memory_resource const&) = delete;
  binning_memory_resource& operator=(binning_memory_resource&&)      = delete;
 
  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
  {
    return upstream_mr_;
  }
 
  void add_bin(std::size_t allocation_size,
               std::optional<device_async_resource_ref> bin_resource = std::nullopt)
  {
    allocation_size = align_up(allocation_size, CUDA_ALLOCATION_ALIGNMENT);
 
    if (bin_resource.has_value()) {
      resource_bins_.insert({allocation_size, bin_resource.value()});
    } else if (resource_bins_.count(allocation_size) == 0) {  // do nothing if bin already exists
      owned_bin_resources_.push_back(
        std::make_unique<fixed_size_memory_resource<Upstream>>(upstream_mr_, allocation_size));
      resource_bins_.insert({allocation_size, owned_bin_resources_.back().get()});
    }
  }
 
 private:
  device_async_resource_ref get_resource_ref(std::size_t bytes)
  {
    auto iter = resource_bins_.lower_bound(bytes);
    return (iter != resource_bins_.cend()) ? iter->second : get_upstream_resource();
  }
 
  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
  {
    if (bytes <= 0) { return nullptr; }
    return get_resource_ref(bytes).allocate_async(bytes, stream);
  }
 
  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
  {
    get_resource_ref(bytes).deallocate_async(ptr, bytes, stream);
  }
 
  device_async_resource_ref
    upstream_mr_;  // The upstream memory_resource from which to allocate blocks.
 
  std::vector<std::unique_ptr<fixed_size_memory_resource<Upstream>>> owned_bin_resources_;
 
  std::map<std::size_t, device_async_resource_ref> resource_bins_;
};
  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE