cache.h

 
#ifndef XGBOOST_CACHE_H_
#define XGBOOST_CACHE_H_
 
#include <xgboost/logging.h>  // for CHECK_EQ, CHECK
 
#include <cstddef>            // for size_t
#include <memory>             // for weak_ptr, shared_ptr, make_shared
#include <mutex>              // for mutex, lock_guard
#include <queue>              // for queue
#include <thread>             // for thread
#include <unordered_map>      // for unordered_map
#include <utility>            // for move
#include <vector>             // for vector
 
namespace xgboost {
class DMatrix;
template <typename CacheT>
class DMatrixCache {
 public:
  struct Item {
    // A weak pointer for checking whether the DMatrix object has expired.
    std::weak_ptr<DMatrix> ref;
    // The cached item
    std::shared_ptr<CacheT> value;
 
    CacheT const& Value() const { return *value; }
    CacheT& Value() { return *value; }
 
    Item(std::shared_ptr<DMatrix> m, std::shared_ptr<CacheT> v) : ref{m}, value{std::move(v)} {}
  };
 
  static constexpr std::size_t DefaultSize() { return 32; }
 
 private:
  mutable std::mutex lock_;
 
 protected:
  struct Key {
    DMatrix const* ptr;
    std::thread::id const thread_id;
 
    bool operator==(Key const& that) const {
      return ptr == that.ptr && thread_id == that.thread_id;
    }
  };
  struct Hash {
    std::size_t operator()(Key const& key) const noexcept {
      std::size_t f = std::hash<DMatrix const*>()(key.ptr);
      std::size_t s = std::hash<std::thread::id>()(key.thread_id);
      if (f == s) {
        return f;
      }
      return f ^ s;
    }
  };
 
  std::unordered_map<Key, Item, Hash> container_;
  std::queue<Key> queue_;
  std::size_t max_size_;
 
  void CheckConsistent() const { CHECK_EQ(queue_.size(), container_.size()); }
 
  void ClearExpired() {
    // Clear expired entries
    this->CheckConsistent();
    std::vector<Key> expired;
    std::queue<Key> remained;
 
    while (!queue_.empty()) {
      auto p_fmat = queue_.front();
      auto it = container_.find(p_fmat);
      CHECK(it != container_.cend());
      if (it->second.ref.expired()) {
        expired.push_back(it->first);
      } else {
        remained.push(it->first);
      }
      queue_.pop();
    }
    CHECK(queue_.empty());
    CHECK_EQ(remained.size() + expired.size(), container_.size());
 
    for (auto const& key : expired) {
      container_.erase(key);
    }
    while (!remained.empty()) {
      auto p_fmat = remained.front();
      queue_.push(p_fmat);
      remained.pop();
    }
    this->CheckConsistent();
  }
 
  void ClearExcess() {
    this->CheckConsistent();
    // clear half of the entries to prevent repeatingly clearing cache.
    std::size_t half_size = max_size_ / 2;
    while (queue_.size() >= half_size && !queue_.empty()) {
      auto p_fmat = queue_.front();
      queue_.pop();
      container_.erase(p_fmat);
    }
    this->CheckConsistent();
  }
 
 public:
  explicit DMatrixCache(std::size_t cache_size) : max_size_{cache_size} {}
 
  DMatrixCache& operator=(DMatrixCache&& that) {
    CHECK(lock_.try_lock());
    lock_.unlock();
    CHECK(that.lock_.try_lock());
    that.lock_.unlock();
    std::swap(this->container_, that.container_);
    std::swap(this->queue_, that.queue_);
    std::swap(this->max_size_, that.max_size_);
    return *this;
  }
 
  template <typename... Args>
  std::shared_ptr<CacheT> CacheItem(std::shared_ptr<DMatrix> m, Args const&... args) {
    CHECK(m);
    std::lock_guard<std::mutex> guard{lock_};
 
    this->ClearExpired();
    if (container_.size() >= max_size_) {
      this->ClearExcess();
    }
    // after clear, cache size < max_size
    CHECK_LT(container_.size(), max_size_);
    auto key = Key{m.get(), std::this_thread::get_id()};
    auto it = container_.find(key);
    if (it == container_.cend()) {
      // after the new DMatrix, cache size is at most max_size
      container_.emplace(key, Item{m, std::make_shared<CacheT>(args...)});
      queue_.emplace(key);
    }
    return container_.at(key).value;
  }
  template <typename... Args>
  std::shared_ptr<CacheT> ResetItem(std::shared_ptr<DMatrix> m, Args const&... args) {
    std::lock_guard<std::mutex> guard{lock_};
    CheckConsistent();
    auto key = Key{m.get(), std::this_thread::get_id()};
    auto it = container_.find(key);
    CHECK(it != container_.cend());
    it->second = {m, std::make_shared<CacheT>(args...)};
    CheckConsistent();
    return it->second.value;
  }
  decltype(container_) const& Container() {
    std::lock_guard<std::mutex> guard{lock_};
 
    this->ClearExpired();
    return container_;
  }
 
  std::shared_ptr<CacheT> Entry(DMatrix const* m) const {
    std::lock_guard<std::mutex> guard{lock_};
    auto key = Key{m, std::this_thread::get_id()};
    CHECK(container_.find(key) != container_.cend());
    CHECK(!container_.at(key).ref.expired());
    return container_.at(key).value;
  }
};
}  // namespace xgboost
#endif  // XGBOOST_CACHE_H_