binning_memory_resource.hpp

/*
 * Copyright (c) 2020-2021, 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/detail/aligned.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/fixed_size_memory_resource.hpp>
 
#include <cuda_runtime_api.h>
 
#include <algorithm>
#include <cassert>
#include <map>
#include <memory>
#include <vector>
 
namespace rmm::mr {
 
template <typename Upstream>
class binning_memory_resource final : public device_memory_resource {
 public:
  explicit binning_memory_resource(Upstream* upstream_resource)
    : upstream_mr_{[upstream_resource]() {
        RMM_EXPECTS(nullptr != upstream_resource, "Unexpected null upstream pointer.");
        return upstream_resource;
      }()}
  {
  }
 
  binning_memory_resource(Upstream* upstream_resource,
                          int8_t min_size_exponent,  // NOLINT(bugprone-easily-swappable-parameters)
                          int8_t max_size_exponent)
    : upstream_mr_{[upstream_resource]() {
        RMM_EXPECTS(nullptr != upstream_resource, "Unexpected null upstream pointer.");
        return 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]] bool supports_streams() const noexcept override { return true; }
 
  [[nodiscard]] bool supports_get_mem_info() const noexcept override { return false; }
 
  [[nodiscard]] Upstream* get_upstream() const noexcept { return upstream_mr_; }
 
  void add_bin(std::size_t allocation_size, device_memory_resource* bin_resource = nullptr)
  {
    allocation_size =
      rmm::detail::align_up(allocation_size, rmm::detail::CUDA_ALLOCATION_ALIGNMENT);
 
    if (nullptr != bin_resource) {
      resource_bins_.insert({allocation_size, bin_resource});
    } 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_memory_resource* get_resource(std::size_t bytes)
  {
    auto iter = resource_bins_.lower_bound(bytes);
    return (iter != resource_bins_.cend()) ? iter->second
                                           : static_cast<device_memory_resource*>(get_upstream());
  }
 
  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
  {
    if (bytes <= 0) { return nullptr; }
    return get_resource(bytes)->allocate(bytes, stream);
  }
 
  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
  {
    auto res = get_resource(bytes);
    if (res != nullptr) { res->deallocate(ptr, bytes, stream); }
  }
 
  [[nodiscard]] std::pair<std::size_t, std::size_t> do_get_mem_info(
    cuda_stream_view stream) const override
  {
    return std::make_pair(0, 0);
  }
 
  Upstream* 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_memory_resource*> resource_bins_;
};
 
}  // namespace rmm::mr