hist_cache.h

 
#ifndef XGBOOST_TREE_HIST_HIST_CACHE_H_
#define XGBOOST_TREE_HIST_HIST_CACHE_H_
#include <cstddef>  // for size_t
#include <map>      // for map
#include <memory>   // for unique_ptr
#include <vector>   // for vector
 
#include "../../common/hist_util.h"          // for GHistRow, ConstGHistRow
#include "../../common/ref_resource_view.h"  // for ReallocVector
#include "xgboost/base.h"                    // for bst_node_t, bst_bin_t
#include "xgboost/logging.h"                 // for CHECK_GT
#include "xgboost/span.h"                    // for Span
 
namespace xgboost::tree {
class BoundedHistCollection {
  // maps node index to offset in `data_`.
  std::map<bst_node_t, std::size_t> node_map_;
  // currently allocated bins, used for tracking consistentcy.
  std::size_t current_size_{0};
 
  // stores the histograms in a contiguous buffer
  using Vec = common::ReallocVector<GradientPairPrecise>;
  std::unique_ptr<Vec> data_{new Vec{}};  // nvcc 12.1 trips over std::make_unique
 
  // number of histogram bins across all features
  bst_bin_t n_total_bins_{0};
  // limits the number of nodes that can be in the cache for each tree
  std::size_t n_cached_nodes_{0};
  // whether the tree has grown beyond the cache limit
  bool has_exceeded_{false};
 
 public:
  BoundedHistCollection() = default;
  common::GHistRow operator[](std::size_t idx) {
    auto offset = node_map_.at(idx);
    return common::Span{data_->data(), data_->size()}.subspan(offset, n_total_bins_);
  }
  common::ConstGHistRow operator[](std::size_t idx) const {
    auto offset = node_map_.at(idx);
    return common::Span{data_->data(), data_->size()}.subspan(offset, n_total_bins_);
  }
  void Reset(bst_bin_t n_total_bins, std::size_t n_cached_nodes) {
    n_total_bins_ = n_total_bins;
    n_cached_nodes_ = n_cached_nodes;
    this->Clear(false);
  }
  void Clear(bool exceeded) {
    node_map_.clear();
    current_size_ = 0;
    has_exceeded_ = exceeded;
  }
 
  [[nodiscard]] bool CanHost(common::Span<bst_node_t const> nodes_to_build,
                             common::Span<bst_node_t const> nodes_to_sub) const {
    auto n_new_nodes = nodes_to_build.size() + nodes_to_sub.size();
    return n_new_nodes + node_map_.size() <= n_cached_nodes_;
  }
 
  void AllocateHistograms(common::Span<bst_node_t const> nodes_to_build,
                          common::Span<bst_node_t const> nodes_to_sub) {
    auto n_new_nodes = nodes_to_build.size() + nodes_to_sub.size();
    auto alloc_size = n_new_nodes * n_total_bins_;
    auto new_size = alloc_size + current_size_;
    if (new_size > data_->size()) {
      data_->Resize(new_size);
    }
    for (auto nidx : nodes_to_build) {
      node_map_[nidx] = current_size_;
      current_size_ += n_total_bins_;
    }
    for (auto nidx : nodes_to_sub) {
      node_map_[nidx] = current_size_;
      current_size_ += n_total_bins_;
    }
    CHECK_EQ(current_size_, new_size);
  }
  void AllocateHistograms(std::vector<bst_node_t> const& nodes) {
    this->AllocateHistograms(common::Span<bst_node_t const>{nodes},
                             common::Span<bst_node_t const>{});
  }
 
  [[nodiscard]] bool HasExceeded() const { return has_exceeded_; }
  [[nodiscard]] bool HistogramExists(bst_node_t nidx) const {
    return node_map_.find(nidx) != node_map_.cend();
  }
  [[nodiscard]] std::size_t Size() const { return current_size_; }
};
}  // namespace xgboost::tree
#endif  // XGBOOST_TREE_HIST_HIST_CACHE_H_