diff --git a/include/xgboost/cache.h b/include/xgboost/cache.h new file mode 100644 index 000000000000..246a92648de7 --- /dev/null +++ b/include/xgboost/cache.h @@ -0,0 +1,134 @@ +/** + * Copyright 2023 by XGBoost contributors + */ +#ifndef XGBOOST_CACHE_H_ +#define XGBOOST_CACHE_H_ + +#include // CHECK_EQ + +#include // std::size_t +#include // std::weak_ptr,std::shared_ptr,std::make_shared +#include // std:queue +#include // std::unordered_map +#include // std::vector + +namespace xgboost { +class DMatrix; +/** + * \brief FIFO cache for DMatrix related data. + * + * \tparam CacheT The type that needs to be cached. + */ +template +class DMatrixCache { + public: + struct Item { + // A weak pointer for checking whether the DMatrix object has expired. + std::weak_ptr ref; + // The cached item + std::shared_ptr value; + + CacheT const& Value() const { return *value; } + CacheT& Value() { return *value; } + }; + + protected: + std::unordered_map container_; + std::queue queue_; + std::size_t max_size_; + + void CheckConsistent() const { CHECK_EQ(queue_.size(), container_.size()); } + + void ClearExpired() { + // Clear expired entries + this->CheckConsistent(); + std::vector expired; + std::queue 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* p_fmat : expired) { + container_.erase(p_fmat); + } + while (!remained.empty()) { + auto p_fmat = remained.front(); + queue_.push(p_fmat); + remained.pop(); + } + this->CheckConsistent(); + } + + void ClearExcess() { + this->CheckConsistent(); + while (queue_.size() >= max_size_) { + auto p_fmat = queue_.front(); + queue_.pop(); + container_.erase(p_fmat); + } + this->CheckConsistent(); + } + + public: + /** + * \param cache_size Maximum size of the cache. + */ + explicit DMatrixCache(std::size_t cache_size) : max_size_{cache_size} {} + /** + * \brief Cache a new DMatrix if it's no in the cache already. + * + * Passing in a `shared_ptr` is critical here. First to create a `weak_ptr` inside the + * entry this shared pointer is necessary. More importantly, the life time of this + * cache is tied to the shared pointer. + * + * \param m shared pointer to the DMatrix that needs to be cached. + * \param args The arguments for constructing a new cache item, if needed. + * + * \return The cache entry for passed in DMatrix, either an existing cache or newly + * created. + */ + template + std::shared_ptr& CacheItem(std::shared_ptr m, Args const&... args) { + CHECK(m); + this->ClearExpired(); + if (container_.size() >= max_size_) { + this->ClearExcess(); + } + // after clear, cache size < max_size + CHECK_LT(container_.size(), max_size_); + auto it = container_.find(m.get()); + if (it == container_.cend()) { + // after the new DMatrix, cache size is at most max_size + container_[m.get()] = {m, std::make_shared(args...)}; + queue_.push(m.get()); + } + return container_.at(m.get()).value; + } + /** + * \brief Get a const reference to the underlying hash map. Clear expired caches before + * returning. + */ + decltype(container_) const& Container() { + this->ClearExpired(); + return container_; + } + + std::shared_ptr Entry(DMatrix const* m) const { + CHECK(container_.find(m) != container_.cend()); + CHECK(!container_.at(m).ref.expired()); + return container_.at(m).value; + } +}; +} // namespace xgboost +#endif // XGBOOST_CACHE_H_ diff --git a/include/xgboost/gbm.h b/include/xgboost/gbm.h index 3be60a061c46..d00f9ceaf602 100644 --- a/include/xgboost/gbm.h +++ b/include/xgboost/gbm.h @@ -1,5 +1,5 @@ -/*! - * Copyright 2014-2022 by XGBoost Contributors +/** + * Copyright 2014-2023 by XGBoost Contributors * \file gbm.h * \brief Interface of gradient booster, * that learns through gradient statistics. @@ -31,7 +31,6 @@ class ObjFunction; struct Context; struct LearnerModelParam; struct PredictionCacheEntry; -class PredictionContainer; /*! * \brief interface of gradient boosting model. diff --git a/include/xgboost/learner.h b/include/xgboost/learner.h index 35d3cf586d6a..ccb171370048 100644 --- a/include/xgboost/learner.h +++ b/include/xgboost/learner.h @@ -1,5 +1,5 @@ -/*! - * Copyright 2015-2022 by XGBoost Contributors +/** + * Copyright 2015-2023 by XGBoost Contributors * \file learner.h * \brief Learner interface that integrates objective, gbm and evaluation together. * This is the user facing XGBoost training module. @@ -8,12 +8,13 @@ #ifndef XGBOOST_LEARNER_H_ #define XGBOOST_LEARNER_H_ +#include // Serializable #include #include // Context #include #include +#include // Tensor #include -#include #include #include @@ -29,6 +30,7 @@ class GradientBooster; class ObjFunction; class DMatrix; class Json; +struct XGBAPIThreadLocalEntry; enum class PredictionType : std::uint8_t { // NOLINT kValue = 0, @@ -40,26 +42,6 @@ enum class PredictionType : std::uint8_t { // NOLINT kLeaf = 6 }; -/*! \brief entry to to easily hold returning information */ -struct XGBAPIThreadLocalEntry { - /*! \brief result holder for returning string */ - std::string ret_str; - /*! \brief result holder for returning raw buffer */ - std::vector ret_char_vec; - /*! \brief result holder for returning strings */ - std::vector ret_vec_str; - /*! \brief result holder for returning string pointers */ - std::vector ret_vec_charp; - /*! \brief returning float vector. */ - std::vector ret_vec_float; - /*! \brief temp variable of gradient pairs. */ - std::vector tmp_gpair; - /*! \brief Temp variable for returning prediction result. */ - PredictionCacheEntry prediction_entry; - /*! \brief Temp variable for returning prediction shape. */ - std::vector prediction_shape; -}; - /*! * \brief Learner class that does training and prediction. * This is the user facing module of xgboost training. diff --git a/include/xgboost/predictor.h b/include/xgboost/predictor.h index 6625612b402d..438c23465a7b 100644 --- a/include/xgboost/predictor.h +++ b/include/xgboost/predictor.h @@ -1,95 +1,66 @@ -/*! - * Copyright 2017-2022 by Contributors +/** + * Copyright 2017-2023 by Contributors * \file predictor.h * \brief Interface of predictor, * performs predictions for a gradient booster. */ #pragma once #include +#include // DMatrixCache #include #include #include -#include +#include // std::function #include -#include #include -#include -#include #include // Forward declarations namespace xgboost { -class TreeUpdater; namespace gbm { struct GBTreeModel; } // namespace gbm -} +} // namespace xgboost namespace xgboost { /** - * \struct PredictionCacheEntry - * * \brief Contains pointer to input matrix and associated cached predictions. */ struct PredictionCacheEntry { // A storage for caching prediction values HostDeviceVector predictions; // The version of current cache, corresponding number of layers of trees - uint32_t version { 0 }; - // A weak pointer for checking whether the DMatrix object has expired. - std::weak_ptr< DMatrix > ref; + std::uint32_t version{0}; PredictionCacheEntry() = default; - /* \brief Update the cache entry by number of versions. + /** + * \brief Update the cache entry by number of versions. * * \param v Added versions. */ - void Update(uint32_t v) { + void Update(std::uint32_t v) { version += v; } }; -/* \brief A container for managed prediction caches. +/** + * \brief A container for managed prediction caches. */ -class PredictionContainer { - std::unordered_map container_; - void ClearExpiredEntries(); +class PredictionContainer : public DMatrixCache { + // we cache up to 32 DMatrix + std::size_t static constexpr DefaultSize() { return 32; } public: - PredictionContainer() = default; - /* \brief Add a new DMatrix to the cache, at the same time this function will clear out - * all expired caches by checking the `std::weak_ptr`. Caching an existing - * DMatrix won't renew it. - * - * Passing in a `shared_ptr` is critical here. First to create a `weak_ptr` inside the - * entry this shared pointer is necessary. More importantly, the life time of this - * cache is tied to the shared pointer. - * - * Another way to make a safe cache is create a proxy to this entry, with anther shared - * pointer defined inside, and pass this proxy around instead of the real entry. But - * seems to be too messy. In XGBoost, functions like `UpdateOneIter` will have - * (memory) safe access to the DMatrix as long as it's passed in as a `shared_ptr`. - * - * \param m shared pointer to the DMatrix that needs to be cached. - * \param device Which device should the cache be allocated on. Pass - * Context::kCpuId for CPU or positive integer for GPU id. - * - * \return the cache entry for passed in DMatrix, either an existing cache or newly - * created. - */ - PredictionCacheEntry& Cache(std::shared_ptr m, int32_t device); - /* \brief Get a prediction cache entry. This entry must be already allocated by `Cache` - * method. Otherwise a dmlc::Error is thrown. - * - * \param m pointer to the DMatrix. - * \return The prediction cache for passed in DMatrix. - */ - PredictionCacheEntry& Entry(DMatrix* m); - /* \brief Get a const reference to the underlying hash map. Clear expired caches before - * returning. - */ - decltype(container_) const& Container(); + PredictionContainer() : DMatrixCache{DefaultSize()} {} + PredictionCacheEntry& Cache(std::shared_ptr m, int32_t device) { + this->CacheItem(m); + auto p_cache = this->container_.find(m.get()); + if (device != Context::kCpuId) { + p_cache->second.Value().predictions.SetDevice(device); + } + return p_cache->second.Value(); + } }; /** @@ -114,7 +85,7 @@ class Predictor { * * \param cfg The configuration. */ - virtual void Configure(const std::vector>&); + virtual void Configure(Args const&); /** * \brief Initialize output prediction diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc index dcb2eb4c2f26..6069da064850 100644 --- a/src/c_api/c_api.cc +++ b/src/c_api/c_api.cc @@ -12,11 +12,11 @@ #include #include "../collective/communicator-inl.h" +#include "../common/api_entry.h" // XGBAPIThreadLocalEntry #include "../common/charconv.h" #include "../common/io.h" #include "../data/adapter.h" #include "../data/simple_dmatrix.h" -#include "c_api_error.h" #include "c_api_utils.h" #include "xgboost/base.h" #include "xgboost/data.h" diff --git a/src/c_api/c_api.cu b/src/c_api/c_api.cu index b75c70adcab1..e6201b0fdc03 100644 --- a/src/c_api/c_api.cu +++ b/src/c_api/c_api.cu @@ -1,6 +1,7 @@ /** * Copyright 2019-2023 by XGBoost Contributors */ +#include "../common/api_entry.h" // XGBAPIThreadLocalEntry #include "../common/threading_utils.h" #include "../data/device_adapter.cuh" #include "../data/proxy_dmatrix.h" diff --git a/src/common/api_entry.h b/src/common/api_entry.h new file mode 100644 index 000000000000..db3bcfbc3271 --- /dev/null +++ b/src/common/api_entry.h @@ -0,0 +1,35 @@ +/** + * Copyright 2016-2023 by XGBoost contributors + */ +#ifndef XGBOOST_COMMON_API_ENTRY_H_ +#define XGBOOST_COMMON_API_ENTRY_H_ +#include // std::string +#include // std::vector + +#include "xgboost/base.h" // GradientPair,bst_ulong +#include "xgboost/predictor.h" // PredictionCacheEntry + +namespace xgboost { +/** + * \brief entry to to easily hold returning information + */ +struct XGBAPIThreadLocalEntry { + /*! \brief result holder for returning string */ + std::string ret_str; + /*! \brief result holder for returning raw buffer */ + std::vector ret_char_vec; + /*! \brief result holder for returning strings */ + std::vector ret_vec_str; + /*! \brief result holder for returning string pointers */ + std::vector ret_vec_charp; + /*! \brief returning float vector. */ + std::vector ret_vec_float; + /*! \brief temp variable of gradient pairs. */ + std::vector tmp_gpair; + /*! \brief Temp variable for returning prediction result. */ + PredictionCacheEntry prediction_entry; + /*! \brief Temp variable for returning prediction shape. */ + std::vector prediction_shape; +}; +} // namespace xgboost +#endif // XGBOOST_COMMON_API_ENTRY_H_ diff --git a/src/data/data.cc b/src/data/data.cc index 5702262122aa..a935220e5891 100644 --- a/src/data/data.cc +++ b/src/data/data.cc @@ -10,6 +10,7 @@ #include #include "../collective/communicator-inl.h" +#include "../common/api_entry.h" // XGBAPIThreadLocalEntry #include "../common/group_data.h" #include "../common/io.h" #include "../common/linalg_op.h" diff --git a/src/learner.cc b/src/learner.cc index b0010c12615a..669566e45798 100644 --- a/src/learner.cc +++ b/src/learner.cc @@ -25,6 +25,7 @@ #include #include "collective/communicator-inl.h" +#include "common/api_entry.h" // XGBAPIThreadLocalEntry #include "common/charconv.h" #include "common/common.h" #include "common/io.h" @@ -430,7 +431,9 @@ class LearnerConfiguration : public Learner { monitor_.Init("Learner"); auto& local_cache = (*ThreadLocalPredictionCache::Get())[this]; for (std::shared_ptr const& d : cache) { - local_cache.Cache(d, Context::kCpuId); + if (d) { + local_cache.Cache(d, Context::kCpuId); + } } } ~LearnerConfiguration() override { @@ -1296,9 +1299,8 @@ class LearnerImpl : public LearnerIO { this->ValidateDMatrix(train.get(), true); auto local_cache = this->GetPredictionCache(); - local_cache->Cache(train, ctx_.gpu_id); - - gbm_->DoBoost(train.get(), in_gpair, &local_cache->Entry(train.get()), obj_.get()); + auto& predt = local_cache->Cache(train, ctx_.gpu_id); + gbm_->DoBoost(train.get(), in_gpair, &predt, obj_.get()); monitor_.Stop("BoostOneIter"); } diff --git a/src/predictor/predictor.cc b/src/predictor/predictor.cc index 38ac9492f4d5..c6ef7fe51c0e 100644 --- a/src/predictor/predictor.cc +++ b/src/predictor/predictor.cc @@ -1,57 +1,28 @@ -/*! - * Copyright 2017-2021 by Contributors +/** + * Copyright 2017-2023 by Contributors */ #include "xgboost/predictor.h" #include -#include +#include // std::string -#include "../gbm/gbtree.h" -#include "xgboost/context.h" -#include "xgboost/data.h" +#include "../gbm/gbtree.h" // GBTreeModel +#include "xgboost/base.h" // bst_row_t,bst_group_t +#include "xgboost/context.h" // Context +#include "xgboost/data.h" // MetaInfo +#include "xgboost/host_device_vector.h" // HostDeviceVector +#include "xgboost/learner.h" // LearnerModelParam +#include "xgboost/linalg.h" // Tensor +#include "xgboost/logging.h" namespace dmlc { DMLC_REGISTRY_ENABLE(::xgboost::PredictorReg); } // namespace dmlc namespace xgboost { -void PredictionContainer::ClearExpiredEntries() { - std::vector expired; - for (auto& kv : container_) { - if (kv.second.ref.expired()) { - expired.emplace_back(kv.first); - } - } - for (auto const& ptr : expired) { - container_.erase(ptr); - } -} - -PredictionCacheEntry &PredictionContainer::Cache(std::shared_ptr m, int32_t device) { - this->ClearExpiredEntries(); - container_[m.get()].ref = m; - if (device != Context::kCpuId) { - container_[m.get()].predictions.SetDevice(device); - } - return container_[m.get()]; -} +void Predictor::Configure(Args const&) {} -PredictionCacheEntry &PredictionContainer::Entry(DMatrix *m) { - CHECK(container_.find(m) != container_.cend()); - CHECK(container_.at(m).ref.lock()) - << "[Internal error]: DMatrix: " << m << " has expired."; - return container_.at(m); -} - -decltype(PredictionContainer::container_) const& PredictionContainer::Container() { - this->ClearExpiredEntries(); - return container_; -} - -void Predictor::Configure( - const std::vector>&) { -} Predictor* Predictor::Create(std::string const& name, Context const* ctx) { auto* e = ::dmlc::Registry::Get()->Find(name); if (e == nullptr) { diff --git a/tests/cpp/test_cache.cc b/tests/cpp/test_cache.cc new file mode 100644 index 000000000000..4099fa2dead8 --- /dev/null +++ b/tests/cpp/test_cache.cc @@ -0,0 +1,55 @@ +/** + * Copyright 2023 by XGBoost contributors + */ +#include +#include +#include // DMatrix + +#include // std::size_t + +#include "helpers.h" // RandomDataGenerator + +namespace xgboost { +namespace { +struct CacheForTest { + std::size_t i; + + explicit CacheForTest(std::size_t k) : i{k} {} +}; +} // namespace + +TEST(DMatrixCache, Basic) { + std::size_t constexpr kRows = 2, kCols = 1, kCacheSize = 4; + DMatrixCache cache(kCacheSize); + + auto add_cache = [&]() { + // Create a lambda function here, so that p_fmat gets deleted upon the + // end of the lambda. This is to test how the cache handle expired + // cache entries. + auto p_fmat = RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix(); + cache.CacheItem(p_fmat, 3); + DMatrix* m = p_fmat.get(); + return m; + }; + auto m = add_cache(); + ASSERT_EQ(cache.Container().size(), 0); + ASSERT_THROW(cache.Entry(m), dmlc::Error); + + auto p_fmat = RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix(); + + auto item = cache.CacheItem(p_fmat, 1); + ASSERT_EQ(cache.Entry(p_fmat.get())->i, 1); + + std::vector> items; + for (std::size_t i = 0; i < kCacheSize * 2; ++i) { + items.emplace_back(RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix()); + cache.CacheItem(items.back(), i); + ASSERT_EQ(cache.Entry(items.back().get())->i, i); + ASSERT_LE(cache.Container().size(), kCacheSize); + if (i > kCacheSize) { + auto k = i - kCacheSize - 1; + ASSERT_THROW(cache.Entry(items[k].get()), dmlc::Error); + } + } +} +} // namespace xgboost diff --git a/tests/cpp/test_learner.cc b/tests/cpp/test_learner.cc index 292846c0ad92..27bfbf21eaa0 100644 --- a/tests/cpp/test_learner.cc +++ b/tests/cpp/test_learner.cc @@ -1,5 +1,5 @@ /*! - * Copyright 2017-2022 by XGBoost contributors + * Copyright 2017-2023 by XGBoost contributors */ #include #include @@ -10,6 +10,7 @@ #include #include +#include "../../src/common/api_entry.h" // XGBAPIThreadLocalEntry #include "../../src/common/io.h" #include "../../src/common/linalg_op.h" #include "../../src/common/random.h"