coalescing_free_list.hpp

/*
 * Copyright (c) 2019-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 <iterator>
#include <rmm/detail/error.hpp>
#include <rmm/mr/device/detail/free_list.hpp>
 
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <iostream>
#include <list>
 
namespace rmm::mr::detail {
 
struct block : public block_base {
  block() = default;
  block(char* ptr, std::size_t size, bool is_head)
    : block_base{ptr}, size_bytes{size}, head{is_head}
  {
  }
 
  [[nodiscard]] inline char* pointer() const { return static_cast<char*>(ptr); }
 
  [[nodiscard]] inline std::size_t size() const { return size_bytes; }
 
  [[nodiscard]] inline bool is_head() const { return head; }
 
  inline bool operator<(block const& rhs) const noexcept { return pointer() < rhs.pointer(); };
 
  [[nodiscard]] inline block merge(block const& blk) const noexcept
  {
    assert(is_contiguous_before(blk));
    return {pointer(), size() + blk.size(), is_head()};
  }
 
  [[nodiscard]] inline bool is_contiguous_before(block const& blk) const noexcept
  {
    // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
    return (pointer() + size() == blk.ptr) and not(blk.is_head());
  }
 
  [[nodiscard]] inline bool fits(std::size_t bytes) const noexcept { return size() >= bytes; }
 
  [[nodiscard]] inline bool is_better_fit(std::size_t bytes, block const& blk) const noexcept
  {
    return fits(bytes) && (size() < blk.size() || blk.size() < bytes);
  }
 
#ifdef RMM_DEBUG_PRINT
 
  inline void print() const
  {
    std::cout << fmt::format("{} {} B", fmt::ptr(pointer()), size()) << std::endl;
  }
#endif
 
 private:
  std::size_t size_bytes{};  
  bool head{};               
};
 
#ifdef RMM_DEBUG_PRINT
inline std::ostream& operator<<(std::ostream& out, const block& blk)
{
  out << fmt::format("{} {} B\n", fmt::ptr(blk.pointer()), blk.size());
  return out;
}
#endif
 
template <typename block_type>
struct compare_blocks {
  // is_transparent (C++14 feature) allows search key type for set<block_type>::find()
  using is_transparent = void;
 
  bool operator()(block_type const& lhs, block_type const& rhs) const { return lhs < rhs; }
  bool operator()(char const* ptr, block_type const& rhs) const { return ptr < rhs.pointer(); }
  bool operator()(block_type const& lhs, char const* ptr) const { return lhs.pointer() < ptr; };
};
 
struct coalescing_free_list : free_list<block> {
  coalescing_free_list()           = default;
  ~coalescing_free_list() override = default;
 
  coalescing_free_list(coalescing_free_list const&) = delete;
  coalescing_free_list& operator=(coalescing_free_list const&) = delete;
  coalescing_free_list(coalescing_free_list&&)                 = delete;
  coalescing_free_list& operator=(coalescing_free_list&&) = delete;
 
  void insert(block_type const& block)
  {
    if (is_empty()) {
      free_list::insert(cend(), block);
      return;
    }
 
    // Find the right place (in ascending ptr order) to insert the block
    // Can't use binary_search because it's a linked list and will be quadratic
    auto const next =
      std::find_if(begin(), end(), [block](block_type const& blk) { return block < blk; });
    auto const previous = (next == cbegin()) ? next : std::prev(next);
 
    // Coalesce with neighboring blocks or insert the new block if it can't be coalesced
    bool const merge_prev = previous->is_contiguous_before(block);
    bool const merge_next = (next != cend()) && block.is_contiguous_before(*next);
 
    if (merge_prev && merge_next) {
      *previous = previous->merge(block).merge(*next);
      erase(next);
    } else if (merge_prev) {
      *previous = previous->merge(block);
    } else if (merge_next) {
      *next = block.merge(*next);
    } else {
      free_list::insert(next, block);  // cannot be coalesced, just insert
    }
  }
 
  void insert(free_list&& other)
  {
    using std::make_move_iterator;
    auto inserter = [this](block_type&& block) { this->insert(block); };
    std::for_each(make_move_iterator(other.begin()), make_move_iterator(other.end()), inserter);
  }
 
  block_type get_block(std::size_t size)
  {
    // find best fit block
    auto finder = [size](block_type const& lhs, block_type const& rhs) {
      return lhs.is_better_fit(size, rhs);
    };
    auto const iter = std::min_element(cbegin(), cend(), finder);
 
    if (iter != cend() && iter->fits(size)) {
      // Remove the block from the free_list and return it.
      block_type const found = *iter;
      erase(iter);
      return found;
    }
 
    return block_type{};  // not found
  }
 
#ifdef RMM_DEBUG_PRINT
 
  void print() const
  {
    std::cout << size() << '\n';
    std::for_each(cbegin(), cend(), [](auto const iter) { iter.print(); });
  }
#endif
};  // coalescing_free_list
 
}  // namespace rmm::mr::detail