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 {
 // 用于检查 DMatrix 对象是否已过期的弱指针。
 std::weak_ptr<DMatrix> ref;
 // 缓存项
 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() {
 // 清除过期的条目
 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();
 // 清除一半条目以防止重复清除缓存。
 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();
    }
 // 清理后，缓存大小 < 最大大小
 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()) {
 // 新 DMatrix 后，缓存大小最多为 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()};
} // namespace xgboost
#endif // XGBOOST_CACHE_H_
 CHECK(container_.find(key) != container_.cend());
  }
};
 CHECK(!container_.at(key).ref.expired());
 return container_.at(key).value;