/**
* Copyright 2020-2024, XGBoost Contributors
*/
#ifndef HISTOGRAM_CUH_
#define HISTOGRAM_CUH_
#include <memory> // for unique_ptr
#include "../../common/cuda_context.cuh" // for CUDAContext
#include "../../common/device_helpers.cuh" // for LaunchN
#include "../../common/device_vector.cuh" // for device_vector
#include "../../data/ellpack_page.cuh" // for EllpackDeviceAccessor
#include "expand_entry.cuh" // for GPUExpandEntry
#include "feature_groups.cuh" // for FeatureGroupsAccessor
#include "quantiser.cuh" // for GradientQuantiser
#include "xgboost/base.h" // for GradientPair, GradientPairInt64
#include "xgboost/context.h" // for Context
#include "xgboost/span.h" // for Span
namespace xgboost::tree {
/**
* \brief An atomicAdd designed for gradient pair with better performance. For general
* int64_t atomicAdd, one can simply cast it to unsigned long long. Exposed for testing.
*/
XGBOOST_DEV_INLINE void AtomicAdd64As32(int64_t* dst, int64_t src) {
uint32_t* y_low = reinterpret_cast<uint32_t*>(dst);
uint32_t* y_high = y_low + 1;
auto cast_src = reinterpret_cast<uint64_t *>(&src);
uint32_t const x_low = static_cast<uint32_t>(src);
uint32_t const x_high = (*cast_src) >> 32;
auto const old = atomicAdd(y_low, x_low);
uint32_t const carry = old > (std::numeric_limits<uint32_t>::max() - x_low) ? 1 : 0;
uint32_t const sig = x_high + carry;
atomicAdd(y_high, sig);
}
namespace cuda_impl {
// Start with about 16mb
std::size_t constexpr DftReserveSize() { return 1 << 22; }
} // namespace cuda_impl
/**
* @brief Data storage for node histograms on device. Automatically expands.
*
* @author Rory
* @date 28/07/2018
*/
class DeviceHistogramStorage {
private:
using GradientSumT = GradientPairInt64;
std::size_t stop_growing_size_{0};
/** @brief Map nidx to starting index of its histogram. */
std::map<int, size_t> nidx_map_;
// Large buffer of zeroed memory, caches histograms
dh::device_vector<typename GradientSumT::ValueT> data_;
// If we run out of storage allocate one histogram at a time in overflow. Not cached,
// overwritten when a new histogram is requested
dh::device_vector<typename GradientSumT::ValueT> overflow_;
std::map<int, size_t> overflow_nidx_map_;
int n_bins_;
static constexpr std::size_t kNumItemsInGradientSum =
sizeof(GradientSumT) / sizeof(typename GradientSumT::ValueT);
static_assert(kNumItemsInGradientSum == 2, "Number of items in gradient type should be 2.");
public:
explicit DeviceHistogramStorage() { data_.reserve(cuda_impl::DftReserveSize()); }
void Reset(Context const* ctx, bst_bin_t n_total_bins, std::size_t max_cached_nodes) {
this->n_bins_ = n_total_bins;
auto d_data = data_.data().get();
dh::LaunchN(data_.size(), ctx->CUDACtx()->Stream(),
[=] __device__(size_t idx) { d_data[idx] = 0.0f; });
nidx_map_.clear();
overflow_nidx_map_.clear();
auto max_cached_bin_values =
static_cast<std::size_t>(n_total_bins) * max_cached_nodes * kNumItemsInGradientSum;
this->stop_growing_size_ = max_cached_bin_values;
}
[[nodiscard]] bool HistogramExists(bst_node_t nidx) const {
return nidx_map_.find(nidx) != nidx_map_.cend() ||
overflow_nidx_map_.find(nidx) != overflow_nidx_map_.cend();
}
[[nodiscard]] int Bins() const { return n_bins_; }
[[nodiscard]] size_t HistogramSize() const { return n_bins_ * kNumItemsInGradientSum; }
dh::device_vector<typename GradientSumT::ValueT>& Data() { return data_; }
void AllocateHistograms(Context const* ctx, std::vector<bst_node_t> const& new_nidxs) {
for (int nidx : new_nidxs) {
CHECK(!HistogramExists(nidx));
}
// Number of items currently used in data
const size_t used_size = nidx_map_.size() * HistogramSize();
const size_t new_used_size = used_size + HistogramSize() * new_nidxs.size();
CHECK_GE(this->stop_growing_size_, kNumItemsInGradientSum);
if (used_size >= this->stop_growing_size_) {
// Use overflow
// Delete previous entries
overflow_nidx_map_.clear();
overflow_.resize(HistogramSize() * new_nidxs.size());
// Zero memory
auto d_data = overflow_.data().get();
dh::LaunchN(overflow_.size(), ctx->CUDACtx()->Stream(),
[=] __device__(size_t idx) { d_data[idx] = 0.0; });
// Append new histograms
for (int nidx : new_nidxs) {
overflow_nidx_map_[nidx] = overflow_nidx_map_.size() * HistogramSize();
}
} else {
CHECK_GE(data_.size(), used_size);
// Expand if necessary
if (data_.size() < new_used_size) {
data_.resize(std::max(data_.size() * 2, new_used_size));
}
// Append new histograms
for (int nidx : new_nidxs) {
nidx_map_[nidx] = nidx_map_.size() * HistogramSize();
}
}
CHECK_GE(data_.size(), nidx_map_.size() * HistogramSize());
}
/**
* \summary Return pointer to histogram memory for a given node.
* \param nidx Tree node index.
* \return hist pointer.
*/
common::Span<GradientSumT> GetNodeHistogram(int nidx) {
CHECK(this->HistogramExists(nidx));
if (nidx_map_.find(nidx) != nidx_map_.cend()) {
// Fetch from normal cache
auto ptr = data_.data().get() + nidx_map_.at(nidx);
return {reinterpret_cast<GradientSumT*>(ptr), static_cast<std::size_t>(n_bins_)};
} else {
// Fetch from overflow
auto ptr = overflow_.data().get() + overflow_nidx_map_.at(nidx);
return {reinterpret_cast<GradientSumT*>(ptr), static_cast<std::size_t>(n_bins_)};
}
}
};
class DeviceHistogramBuilderImpl;
class DeviceHistogramBuilder {
std::unique_ptr<DeviceHistogramBuilderImpl> p_impl_;
DeviceHistogramStorage hist_;
common::Monitor monitor_;
public:
explicit DeviceHistogramBuilder();
~DeviceHistogramBuilder();
void Reset(Context const* ctx, std::size_t max_cached_hist_nodes,
FeatureGroupsAccessor const& feature_groups, bst_bin_t n_total_bins,
bool force_global_memory);
void BuildHistogram(CUDAContext const* ctx, EllpackDeviceAccessor const& matrix,
FeatureGroupsAccessor const& feature_groups,
common::Span<GradientPair const> gpair,
common::Span<const std::uint32_t> ridx,
common::Span<GradientPairInt64> histogram, GradientQuantiser rounding);
[[nodiscard]] auto GetNodeHistogram(bst_node_t nidx) { return hist_.GetNodeHistogram(nidx); }
// num histograms is the number of contiguous histograms in memory to reduce over
void AllReduceHist(Context const* ctx, MetaInfo const& info, bst_node_t nidx,
std::size_t num_histograms);
// Attempt to do subtraction trick
// return true if succeeded
[[nodiscard]] bool SubtractionTrick(Context const* ctx, bst_node_t nidx_parent,
bst_node_t nidx_histogram, bst_node_t nidx_subtraction) {
if (!hist_.HistogramExists(nidx_histogram) || !hist_.HistogramExists(nidx_parent)) {
return false;
}
auto d_node_hist_parent = hist_.GetNodeHistogram(nidx_parent);
auto d_node_hist_histogram = hist_.GetNodeHistogram(nidx_histogram);
auto d_node_hist_subtraction = hist_.GetNodeHistogram(nidx_subtraction);
dh::LaunchN(d_node_hist_parent.size(), ctx->CUDACtx()->Stream(), [=] __device__(size_t idx) {
d_node_hist_subtraction[idx] = d_node_hist_parent[idx] - d_node_hist_histogram[idx];
});
return true;
}
[[nodiscard]] auto SubtractHist(Context const* ctx, std::vector<GPUExpandEntry> const& candidates,
std::vector<bst_node_t> const& build_nidx,
std::vector<bst_node_t> const& subtraction_nidx) {
this->monitor_.Start(__func__);
std::vector<bst_node_t> need_build;
for (std::size_t i = 0; i < subtraction_nidx.size(); i++) {
auto build_hist_nidx = build_nidx.at(i);
auto subtraction_trick_nidx = subtraction_nidx.at(i);
auto parent_nidx = candidates.at(i).nid;
if (!this->SubtractionTrick(ctx, parent_nidx, build_hist_nidx, subtraction_trick_nidx)) {
need_build.push_back(subtraction_trick_nidx);
}
}
this->monitor_.Stop(__func__);
return need_build;
}
void AllocateHistograms(Context const* ctx, std::vector<bst_node_t> const& nodes_to_build,
std::vector<bst_node_t> const& nodes_to_sub) {
this->monitor_.Start(__func__);
std::vector<bst_node_t> all_new = nodes_to_build;
all_new.insert(all_new.end(), nodes_to_sub.cbegin(), nodes_to_sub.cend());
// Allocate the histograms
// Guaranteed contiguous memory
this->AllocateHistograms(ctx, all_new);
this->monitor_.Stop(__func__);
}
void AllocateHistograms(Context const* ctx, std::vector<int> const& new_nidxs) {
this->hist_.AllocateHistograms(ctx, new_nidxs);
}
};
} // namespace xgboost::tree
#endif // HISTOGRAM_CUH_