Use ptr from mmap for GHistIndexMatrix and ColumnMatrix. (#9315)
* Use ptr from mmap for `GHistIndexMatrix` and `ColumnMatrix`. - Define a resource for holding various types of memory pointers. - Define ref vector for holding resources. - Swap the underlying resources for GHist and ColumnM. - Add documentation for current status. - s390x support is removed. It should work if you can compile XGBoost, all the old workaround code does is to get GCC to compile.
This commit is contained in:
Jiaming Yuan
GitHub
@@ -1,16 +1,27 @@
|
||||
/*!
|
||||
* Copyright 2017-2022 by XGBoost Contributors
|
||||
/**
|
||||
* Copyright 2017-2023, XGBoost Contributors
|
||||
* \brief Utility for fast column-wise access
|
||||
*/
|
||||
#include "column_matrix.h"
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
#include <algorithm> // for transform
|
||||
#include <cstddef> // for size_t
|
||||
#include <cstdint> // for uint64_t, uint8_t
|
||||
#include <limits> // for numeric_limits
|
||||
#include <type_traits> // for remove_reference_t
|
||||
#include <vector> // for vector
|
||||
|
||||
#include "../data/gradient_index.h" // for GHistIndexMatrix
|
||||
#include "io.h" // for AlignedResourceReadStream, AlignedFileWriteStream
|
||||
#include "xgboost/base.h" // for bst_feaature_t
|
||||
#include "xgboost/span.h" // for Span
|
||||
|
||||
namespace xgboost::common {
|
||||
void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold) {
|
||||
auto const nfeature = gmat.Features();
|
||||
const size_t nrow = gmat.Size();
|
||||
// identify type of each column
|
||||
type_.resize(nfeature);
|
||||
type_ = common::MakeFixedVecWithMalloc(nfeature, ColumnType{});
|
||||
|
||||
uint32_t max_val = std::numeric_limits<uint32_t>::max();
|
||||
for (bst_feature_t fid = 0; fid < nfeature; ++fid) {
|
||||
@@ -34,7 +45,7 @@ void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_thres
|
||||
|
||||
// want to compute storage boundary for each feature
|
||||
// using variants of prefix sum scan
|
||||
feature_offsets_.resize(nfeature + 1);
|
||||
feature_offsets_ = common::MakeFixedVecWithMalloc(nfeature + 1, std::size_t{0});
|
||||
size_t accum_index = 0;
|
||||
feature_offsets_[0] = accum_index;
|
||||
for (bst_feature_t fid = 1; fid < nfeature + 1; ++fid) {
|
||||
@@ -49,9 +60,11 @@ void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_thres
|
||||
SetTypeSize(gmat.MaxNumBinPerFeat());
|
||||
auto storage_size =
|
||||
feature_offsets_.back() * static_cast<std::underlying_type_t<BinTypeSize>>(bins_type_size_);
|
||||
index_.resize(storage_size, 0);
|
||||
|
||||
index_ = common::MakeFixedVecWithMalloc(storage_size, std::uint8_t{0});
|
||||
|
||||
if (!all_dense_column) {
|
||||
row_ind_.resize(feature_offsets_[nfeature]);
|
||||
row_ind_ = common::MakeFixedVecWithMalloc(feature_offsets_[nfeature], std::size_t{0});
|
||||
}
|
||||
|
||||
// store least bin id for each feature
|
||||
@@ -59,7 +72,51 @@ void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_thres
|
||||
|
||||
any_missing_ = !gmat.IsDense();
|
||||
|
||||
missing_flags_.clear();
|
||||
missing_ = MissingIndicator{0, false};
|
||||
}
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
|
||||
// IO procedures for external memory.
|
||||
bool ColumnMatrix::Read(AlignedResourceReadStream* fi, uint32_t const* index_base) {
|
||||
if (!common::ReadVec(fi, &index_)) {
|
||||
return false;
|
||||
}
|
||||
if (!common::ReadVec(fi, &type_)) {
|
||||
return false;
|
||||
}
|
||||
if (!common::ReadVec(fi, &row_ind_)) {
|
||||
return false;
|
||||
}
|
||||
if (!common::ReadVec(fi, &feature_offsets_)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!common::ReadVec(fi, &missing_.storage)) {
|
||||
return false;
|
||||
}
|
||||
missing_.InitView();
|
||||
|
||||
index_base_ = index_base;
|
||||
if (!fi->Read(&bins_type_size_)) {
|
||||
return false;
|
||||
}
|
||||
if (!fi->Read(&any_missing_)) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::size_t ColumnMatrix::Write(AlignedFileWriteStream* fo) const {
|
||||
std::size_t bytes{0};
|
||||
|
||||
bytes += common::WriteVec(fo, index_);
|
||||
bytes += common::WriteVec(fo, type_);
|
||||
bytes += common::WriteVec(fo, row_ind_);
|
||||
bytes += common::WriteVec(fo, feature_offsets_);
|
||||
bytes += common::WriteVec(fo, missing_.storage);
|
||||
|
||||
bytes += fo->Write(bins_type_size_);
|
||||
bytes += fo->Write(any_missing_);
|
||||
|
||||
return bytes;
|
||||
}
|
||||
} // namespace xgboost::common
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/*!
|
||||
* Copyright 2017-2022 by Contributors
|
||||
/**
|
||||
* Copyright 2017-2023, XGBoost Contributors
|
||||
* \file column_matrix.h
|
||||
* \brief Utility for fast column-wise access
|
||||
* \author Philip Cho
|
||||
@@ -8,25 +8,30 @@
|
||||
#ifndef XGBOOST_COMMON_COLUMN_MATRIX_H_
|
||||
#define XGBOOST_COMMON_COLUMN_MATRIX_H_
|
||||
|
||||
#include <dmlc/endian.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstddef> // for size_t
|
||||
#include <cstdint> // for uint8_t
|
||||
#include <limits>
|
||||
#include <memory>
|
||||
#include <utility> // std::move
|
||||
#include <utility> // for move
|
||||
#include <vector>
|
||||
|
||||
#include "../data/adapter.h"
|
||||
#include "../data/gradient_index.h"
|
||||
#include "algorithm.h"
|
||||
#include "bitfield.h" // for RBitField8
|
||||
#include "hist_util.h"
|
||||
#include "ref_resource_view.h" // for RefResourceView
|
||||
#include "xgboost/base.h" // for bst_bin_t
|
||||
#include "xgboost/span.h" // for Span
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
|
||||
namespace xgboost::common {
|
||||
class ColumnMatrix;
|
||||
class AlignedFileWriteStream;
|
||||
class AlignedResourceReadStream;
|
||||
|
||||
/*! \brief column type */
|
||||
enum ColumnType : uint8_t { kDenseColumn, kSparseColumn };
|
||||
enum ColumnType : std::uint8_t { kDenseColumn, kSparseColumn };
|
||||
|
||||
/*! \brief a column storage, to be used with ApplySplit. Note that each
|
||||
bin id is stored as index[i] + index_base.
|
||||
@@ -41,12 +46,12 @@ class Column {
|
||||
: index_(index), index_base_(least_bin_idx) {}
|
||||
virtual ~Column() = default;
|
||||
|
||||
bst_bin_t GetGlobalBinIdx(size_t idx) const {
|
||||
[[nodiscard]] bst_bin_t GetGlobalBinIdx(size_t idx) const {
|
||||
return index_base_ + static_cast<bst_bin_t>(index_[idx]);
|
||||
}
|
||||
|
||||
/* returns number of elements in column */
|
||||
size_t Size() const { return index_.size(); }
|
||||
[[nodiscard]] size_t Size() const { return index_.size(); }
|
||||
|
||||
private:
|
||||
/* bin indexes in range [0, max_bins - 1] */
|
||||
@@ -63,7 +68,7 @@ class SparseColumnIter : public Column<BinIdxT> {
|
||||
common::Span<const size_t> row_ind_;
|
||||
size_t idx_;
|
||||
|
||||
size_t const* RowIndices() const { return row_ind_.data(); }
|
||||
[[nodiscard]] size_t const* RowIndices() const { return row_ind_.data(); }
|
||||
|
||||
public:
|
||||
SparseColumnIter(common::Span<const BinIdxT> index, bst_bin_t least_bin_idx,
|
||||
@@ -81,7 +86,7 @@ class SparseColumnIter : public Column<BinIdxT> {
|
||||
SparseColumnIter(SparseColumnIter const&) = delete;
|
||||
SparseColumnIter(SparseColumnIter&&) = default;
|
||||
|
||||
size_t GetRowIdx(size_t idx) const { return RowIndices()[idx]; }
|
||||
[[nodiscard]] size_t GetRowIdx(size_t idx) const { return RowIndices()[idx]; }
|
||||
bst_bin_t operator[](size_t rid) {
|
||||
const size_t column_size = this->Size();
|
||||
if (!((idx_) < column_size)) {
|
||||
@@ -101,6 +106,10 @@ class SparseColumnIter : public Column<BinIdxT> {
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Column stored as a dense vector. It might still contain missing values as
|
||||
* indicated by the missing flags.
|
||||
*/
|
||||
template <typename BinIdxT, bool any_missing>
|
||||
class DenseColumnIter : public Column<BinIdxT> {
|
||||
public:
|
||||
@@ -109,17 +118,19 @@ class DenseColumnIter : public Column<BinIdxT> {
|
||||
private:
|
||||
using Base = Column<BinIdxT>;
|
||||
/* flags for missing values in dense columns */
|
||||
std::vector<ByteType> const& missing_flags_;
|
||||
LBitField32 missing_flags_;
|
||||
size_t feature_offset_;
|
||||
|
||||
public:
|
||||
explicit DenseColumnIter(common::Span<const BinIdxT> index, bst_bin_t index_base,
|
||||
std::vector<ByteType> const& missing_flags, size_t feature_offset)
|
||||
LBitField32 missing_flags, size_t feature_offset)
|
||||
: Base{index, index_base}, missing_flags_{missing_flags}, feature_offset_{feature_offset} {}
|
||||
DenseColumnIter(DenseColumnIter const&) = delete;
|
||||
DenseColumnIter(DenseColumnIter&&) = default;
|
||||
|
||||
bool IsMissing(size_t ridx) const { return missing_flags_[feature_offset_ + ridx]; }
|
||||
[[nodiscard]] bool IsMissing(size_t ridx) const {
|
||||
return missing_flags_.Check(feature_offset_ + ridx);
|
||||
}
|
||||
|
||||
bst_bin_t operator[](size_t ridx) const {
|
||||
if (any_missing) {
|
||||
@@ -131,12 +142,54 @@ class DenseColumnIter : public Column<BinIdxT> {
|
||||
};
|
||||
|
||||
/**
|
||||
* \brief Column major matrix for gradient index. This matrix contains both dense column
|
||||
* and sparse column, the type of the column is controlled by sparse threshold. When the
|
||||
* number of missing values in a column is below the threshold it's classified as dense
|
||||
* column.
|
||||
* @brief Column major matrix for gradient index on CPU.
|
||||
*
|
||||
* This matrix contains both dense columns and sparse columns, the type of the column
|
||||
* is controlled by the sparse threshold parameter. When the number of missing values
|
||||
* in a column is below the threshold it's classified as dense column.
|
||||
*/
|
||||
class ColumnMatrix {
|
||||
/**
|
||||
* @brief A bit set for indicating whether an element in a dense column is missing.
|
||||
*/
|
||||
struct MissingIndicator {
|
||||
LBitField32 missing;
|
||||
RefResourceView<std::uint32_t> storage;
|
||||
|
||||
MissingIndicator() = default;
|
||||
/**
|
||||
* @param n_elements Size of the bit set
|
||||
* @param init Initialize the indicator to true or false.
|
||||
*/
|
||||
MissingIndicator(std::size_t n_elements, bool init) {
|
||||
auto m_size = missing.ComputeStorageSize(n_elements);
|
||||
storage = common::MakeFixedVecWithMalloc(m_size, init ? ~std::uint32_t{0} : std::uint32_t{0});
|
||||
this->InitView();
|
||||
}
|
||||
/** @brief Set the i^th element to be a valid element (instead of missing). */
|
||||
void SetValid(typename LBitField32::index_type i) { missing.Clear(i); }
|
||||
/** @brief assign the storage to the view. */
|
||||
void InitView() {
|
||||
missing = LBitField32{Span{storage.data(), storage.size()}};
|
||||
}
|
||||
|
||||
void GrowTo(std::size_t n_elements, bool init) {
|
||||
CHECK(storage.Resource()->Type() == ResourceHandler::kMalloc)
|
||||
<< "[Internal Error]: Cannot grow the vector when external memory is used.";
|
||||
auto m_size = missing.ComputeStorageSize(n_elements);
|
||||
CHECK_GE(m_size, storage.size());
|
||||
if (m_size == storage.size()) {
|
||||
return;
|
||||
}
|
||||
|
||||
auto new_storage =
|
||||
common::MakeFixedVecWithMalloc(m_size, init ? ~std::uint32_t{0} : std::uint32_t{0});
|
||||
std::copy_n(storage.cbegin(), storage.size(), new_storage.begin());
|
||||
storage = std::move(new_storage);
|
||||
this->InitView();
|
||||
}
|
||||
};
|
||||
|
||||
void InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold);
|
||||
|
||||
template <typename ColumnBinT, typename BinT, typename RIdx>
|
||||
@@ -144,9 +197,10 @@ class ColumnMatrix {
|
||||
if (type_[fid] == kDenseColumn) {
|
||||
ColumnBinT* begin = &local_index[feature_offsets_[fid]];
|
||||
begin[rid] = bin_id - index_base_[fid];
|
||||
// not thread-safe with bool vector. FIXME(jiamingy): We can directly assign
|
||||
// kMissingId to the index to avoid missing flags.
|
||||
missing_flags_[feature_offsets_[fid] + rid] = false;
|
||||
// not thread-safe with bit field.
|
||||
// FIXME(jiamingy): We can directly assign kMissingId to the index to avoid missing
|
||||
// flags.
|
||||
missing_.SetValid(feature_offsets_[fid] + rid);
|
||||
} else {
|
||||
ColumnBinT* begin = &local_index[feature_offsets_[fid]];
|
||||
begin[num_nonzeros_[fid]] = bin_id - index_base_[fid];
|
||||
@@ -158,7 +212,9 @@ class ColumnMatrix {
|
||||
public:
|
||||
using ByteType = bool;
|
||||
// get number of features
|
||||
bst_feature_t GetNumFeature() const { return static_cast<bst_feature_t>(type_.size()); }
|
||||
[[nodiscard]] bst_feature_t GetNumFeature() const {
|
||||
return static_cast<bst_feature_t>(type_.size());
|
||||
}
|
||||
|
||||
ColumnMatrix() = default;
|
||||
ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold) {
|
||||
@@ -166,7 +222,7 @@ class ColumnMatrix {
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Initialize ColumnMatrix from GHistIndexMatrix with reference to the original
|
||||
* @brief Initialize ColumnMatrix from GHistIndexMatrix with reference to the original
|
||||
* SparsePage.
|
||||
*/
|
||||
void InitFromSparse(SparsePage const& page, const GHistIndexMatrix& gmat, double sparse_threshold,
|
||||
@@ -178,8 +234,8 @@ class ColumnMatrix {
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Initialize ColumnMatrix from GHistIndexMatrix without reference to actual
|
||||
* data.
|
||||
* @brief Initialize ColumnMatrix from GHistIndexMatrix without reference to actual
|
||||
* data.
|
||||
*
|
||||
* This function requires a binary search for each bin to get back the feature index
|
||||
* for those bins.
|
||||
@@ -199,7 +255,7 @@ class ColumnMatrix {
|
||||
}
|
||||
}
|
||||
|
||||
bool IsInitialized() const { return !type_.empty(); }
|
||||
[[nodiscard]] bool IsInitialized() const { return !type_.empty(); }
|
||||
|
||||
/**
|
||||
* \brief Push batch of data for Quantile DMatrix support.
|
||||
@@ -257,7 +313,7 @@ class ColumnMatrix {
|
||||
reinterpret_cast<const BinIdxType*>(&index_[feature_offset * bins_type_size_]),
|
||||
column_size};
|
||||
return std::move(DenseColumnIter<BinIdxType, any_missing>{
|
||||
bin_index, static_cast<bst_bin_t>(index_base_[fidx]), missing_flags_, feature_offset});
|
||||
bin_index, static_cast<bst_bin_t>(index_base_[fidx]), missing_.missing, feature_offset});
|
||||
}
|
||||
|
||||
// all columns are dense column and has no missing value
|
||||
@@ -265,7 +321,8 @@ class ColumnMatrix {
|
||||
template <typename RowBinIdxT>
|
||||
void SetIndexNoMissing(bst_row_t base_rowid, RowBinIdxT const* row_index, const size_t n_samples,
|
||||
const size_t n_features, int32_t n_threads) {
|
||||
missing_flags_.resize(feature_offsets_[n_features], false);
|
||||
missing_.GrowTo(feature_offsets_[n_features], false);
|
||||
|
||||
DispatchBinType(bins_type_size_, [&](auto t) {
|
||||
using ColumnBinT = decltype(t);
|
||||
auto column_index = Span<ColumnBinT>{reinterpret_cast<ColumnBinT*>(index_.data()),
|
||||
@@ -290,9 +347,15 @@ class ColumnMatrix {
|
||||
void SetIndexMixedColumns(size_t base_rowid, Batch const& batch, const GHistIndexMatrix& gmat,
|
||||
float missing) {
|
||||
auto n_features = gmat.Features();
|
||||
missing_flags_.resize(feature_offsets_[n_features], true);
|
||||
auto const* row_index = gmat.index.data<uint32_t>() + gmat.row_ptr[base_rowid];
|
||||
num_nonzeros_.resize(n_features, 0);
|
||||
|
||||
missing_.GrowTo(feature_offsets_[n_features], true);
|
||||
auto const* row_index = gmat.index.data<std::uint32_t>() + gmat.row_ptr[base_rowid];
|
||||
if (num_nonzeros_.empty()) {
|
||||
num_nonzeros_ = common::MakeFixedVecWithMalloc(n_features, std::size_t{0});
|
||||
} else {
|
||||
CHECK_EQ(num_nonzeros_.size(), n_features);
|
||||
}
|
||||
|
||||
auto is_valid = data::IsValidFunctor{missing};
|
||||
|
||||
DispatchBinType(bins_type_size_, [&](auto t) {
|
||||
@@ -321,8 +384,9 @@ class ColumnMatrix {
|
||||
*/
|
||||
void SetIndexMixedColumns(const GHistIndexMatrix& gmat) {
|
||||
auto n_features = gmat.Features();
|
||||
missing_flags_.resize(feature_offsets_[n_features], true);
|
||||
num_nonzeros_.resize(n_features, 0);
|
||||
|
||||
missing_ = MissingIndicator{feature_offsets_[n_features], true};
|
||||
num_nonzeros_ = common::MakeFixedVecWithMalloc(n_features, std::size_t{0});
|
||||
|
||||
DispatchBinType(bins_type_size_, [&](auto t) {
|
||||
using ColumnBinT = decltype(t);
|
||||
@@ -335,106 +399,34 @@ class ColumnMatrix {
|
||||
});
|
||||
}
|
||||
|
||||
BinTypeSize GetTypeSize() const { return bins_type_size_; }
|
||||
auto GetColumnType(bst_feature_t fidx) const { return type_[fidx]; }
|
||||
[[nodiscard]] BinTypeSize GetTypeSize() const { return bins_type_size_; }
|
||||
[[nodiscard]] auto GetColumnType(bst_feature_t fidx) const { return type_[fidx]; }
|
||||
|
||||
// And this returns part of state
|
||||
bool AnyMissing() const { return any_missing_; }
|
||||
[[nodiscard]] bool AnyMissing() const { return any_missing_; }
|
||||
|
||||
// IO procedures for external memory.
|
||||
bool Read(dmlc::SeekStream* fi, uint32_t const* index_base) {
|
||||
fi->Read(&index_);
|
||||
#if !DMLC_LITTLE_ENDIAN
|
||||
// s390x
|
||||
std::vector<std::underlying_type<ColumnType>::type> int_types;
|
||||
fi->Read(&int_types);
|
||||
type_.resize(int_types.size());
|
||||
std::transform(
|
||||
int_types.begin(), int_types.end(), type_.begin(),
|
||||
[](std::underlying_type<ColumnType>::type i) { return static_cast<ColumnType>(i); });
|
||||
#else
|
||||
fi->Read(&type_);
|
||||
#endif // !DMLC_LITTLE_ENDIAN
|
||||
|
||||
fi->Read(&row_ind_);
|
||||
fi->Read(&feature_offsets_);
|
||||
|
||||
std::vector<std::uint8_t> missing;
|
||||
fi->Read(&missing);
|
||||
missing_flags_.resize(missing.size());
|
||||
std::transform(missing.cbegin(), missing.cend(), missing_flags_.begin(),
|
||||
[](std::uint8_t flag) { return !!flag; });
|
||||
|
||||
index_base_ = index_base;
|
||||
#if !DMLC_LITTLE_ENDIAN
|
||||
std::underlying_type<BinTypeSize>::type v;
|
||||
fi->Read(&v);
|
||||
bins_type_size_ = static_cast<BinTypeSize>(v);
|
||||
#else
|
||||
fi->Read(&bins_type_size_);
|
||||
#endif
|
||||
|
||||
fi->Read(&any_missing_);
|
||||
return true;
|
||||
}
|
||||
|
||||
size_t Write(dmlc::Stream* fo) const {
|
||||
size_t bytes{0};
|
||||
|
||||
auto write_vec = [&](auto const& vec) {
|
||||
fo->Write(vec);
|
||||
bytes += vec.size() * sizeof(typename std::remove_reference_t<decltype(vec)>::value_type) +
|
||||
sizeof(uint64_t);
|
||||
};
|
||||
write_vec(index_);
|
||||
#if !DMLC_LITTLE_ENDIAN
|
||||
// s390x
|
||||
std::vector<std::underlying_type<ColumnType>::type> int_types(type_.size());
|
||||
std::transform(type_.begin(), type_.end(), int_types.begin(), [](ColumnType t) {
|
||||
return static_cast<std::underlying_type<ColumnType>::type>(t);
|
||||
});
|
||||
write_vec(int_types);
|
||||
#else
|
||||
write_vec(type_);
|
||||
#endif // !DMLC_LITTLE_ENDIAN
|
||||
write_vec(row_ind_);
|
||||
write_vec(feature_offsets_);
|
||||
// dmlc can not handle bool vector
|
||||
std::vector<std::uint8_t> missing(missing_flags_.size());
|
||||
std::transform(missing_flags_.cbegin(), missing_flags_.cend(), missing.begin(),
|
||||
[](bool flag) { return static_cast<std::uint8_t>(flag); });
|
||||
write_vec(missing);
|
||||
|
||||
#if !DMLC_LITTLE_ENDIAN
|
||||
auto v = static_cast<std::underlying_type<BinTypeSize>::type>(bins_type_size_);
|
||||
fo->Write(v);
|
||||
#else
|
||||
fo->Write(bins_type_size_);
|
||||
#endif // DMLC_LITTLE_ENDIAN
|
||||
bytes += sizeof(bins_type_size_);
|
||||
fo->Write(any_missing_);
|
||||
bytes += sizeof(any_missing_);
|
||||
|
||||
return bytes;
|
||||
}
|
||||
[[nodiscard]] bool Read(AlignedResourceReadStream* fi, uint32_t const* index_base);
|
||||
[[nodiscard]] std::size_t Write(AlignedFileWriteStream* fo) const;
|
||||
|
||||
private:
|
||||
std::vector<uint8_t> index_;
|
||||
RefResourceView<std::uint8_t> index_;
|
||||
|
||||
std::vector<ColumnType> type_;
|
||||
/* indptr of a CSC matrix. */
|
||||
std::vector<size_t> row_ind_;
|
||||
/* indicate where each column's index and row_ind is stored. */
|
||||
std::vector<size_t> feature_offsets_;
|
||||
/* The number of nnz of each column. */
|
||||
std::vector<size_t> num_nonzeros_;
|
||||
RefResourceView<ColumnType> type_;
|
||||
/** @brief indptr of a CSC matrix. */
|
||||
RefResourceView<std::size_t> row_ind_;
|
||||
/** @brief indicate where each column's index and row_ind is stored. */
|
||||
RefResourceView<std::size_t> feature_offsets_;
|
||||
/** @brief The number of nnz of each column. */
|
||||
RefResourceView<std::size_t> num_nonzeros_;
|
||||
|
||||
// index_base_[fid]: least bin id for feature fid
|
||||
uint32_t const* index_base_;
|
||||
std::vector<ByteType> missing_flags_;
|
||||
std::uint32_t const* index_base_;
|
||||
|
||||
MissingIndicator missing_;
|
||||
|
||||
BinTypeSize bins_type_size_;
|
||||
bool any_missing_;
|
||||
};
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
} // namespace xgboost::common
|
||||
#endif // XGBOOST_COMMON_COLUMN_MATRIX_H_
|
||||
|
||||
@@ -203,13 +203,33 @@ auto DispatchBinType(BinTypeSize type, Fn&& fn) {
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Optionally compressed gradient index. The compression works only with dense
|
||||
* @brief Optionally compressed gradient index. The compression works only with dense
|
||||
* data.
|
||||
*
|
||||
* The main body of construction code is in gradient_index.cc, this struct is only a
|
||||
* storage class.
|
||||
* view class.
|
||||
*/
|
||||
struct Index {
|
||||
class Index {
|
||||
private:
|
||||
void SetBinTypeSize(BinTypeSize binTypeSize) {
|
||||
binTypeSize_ = binTypeSize;
|
||||
switch (binTypeSize) {
|
||||
case kUint8BinsTypeSize:
|
||||
func_ = &GetValueFromUint8;
|
||||
break;
|
||||
case kUint16BinsTypeSize:
|
||||
func_ = &GetValueFromUint16;
|
||||
break;
|
||||
case kUint32BinsTypeSize:
|
||||
func_ = &GetValueFromUint32;
|
||||
break;
|
||||
default:
|
||||
CHECK(binTypeSize == kUint8BinsTypeSize || binTypeSize == kUint16BinsTypeSize ||
|
||||
binTypeSize == kUint32BinsTypeSize);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
// Inside the compressor, bin_idx is the index for cut value across all features. By
|
||||
// subtracting it with starting pointer of each feature, we can reduce it to smaller
|
||||
// value and store it with smaller types. Usable only with dense data.
|
||||
@@ -233,10 +253,24 @@ struct Index {
|
||||
}
|
||||
|
||||
Index() { SetBinTypeSize(binTypeSize_); }
|
||||
Index(const Index& i) = delete;
|
||||
Index& operator=(Index i) = delete;
|
||||
|
||||
Index(Index const& i) = delete;
|
||||
Index& operator=(Index const& i) = delete;
|
||||
Index(Index&& i) = delete;
|
||||
Index& operator=(Index&& i) = delete;
|
||||
|
||||
/** @brief Move assignment for lazy initialization. */
|
||||
Index& operator=(Index&& i) = default;
|
||||
|
||||
/**
|
||||
* @brief Construct the index from data.
|
||||
*
|
||||
* @param data Storage for compressed histogram bin.
|
||||
* @param bin_size Number of bytes for each bin.
|
||||
*/
|
||||
Index(Span<std::uint8_t> data, BinTypeSize bin_size) : data_{data} {
|
||||
this->SetBinTypeSize(bin_size);
|
||||
}
|
||||
|
||||
uint32_t operator[](size_t i) const {
|
||||
if (!bin_offset_.empty()) {
|
||||
// dense, compressed
|
||||
@@ -247,26 +281,7 @@ struct Index {
|
||||
return func_(data_.data(), i);
|
||||
}
|
||||
}
|
||||
void SetBinTypeSize(BinTypeSize binTypeSize) {
|
||||
binTypeSize_ = binTypeSize;
|
||||
switch (binTypeSize) {
|
||||
case kUint8BinsTypeSize:
|
||||
func_ = &GetValueFromUint8;
|
||||
break;
|
||||
case kUint16BinsTypeSize:
|
||||
func_ = &GetValueFromUint16;
|
||||
break;
|
||||
case kUint32BinsTypeSize:
|
||||
func_ = &GetValueFromUint32;
|
||||
break;
|
||||
default:
|
||||
CHECK(binTypeSize == kUint8BinsTypeSize || binTypeSize == kUint16BinsTypeSize ||
|
||||
binTypeSize == kUint32BinsTypeSize);
|
||||
}
|
||||
}
|
||||
BinTypeSize GetBinTypeSize() const {
|
||||
return binTypeSize_;
|
||||
}
|
||||
[[nodiscard]] BinTypeSize GetBinTypeSize() const { return binTypeSize_; }
|
||||
template <typename T>
|
||||
T const* data() const { // NOLINT
|
||||
return reinterpret_cast<T const*>(data_.data());
|
||||
@@ -275,30 +290,27 @@ struct Index {
|
||||
T* data() { // NOLINT
|
||||
return reinterpret_cast<T*>(data_.data());
|
||||
}
|
||||
uint32_t const* Offset() const { return bin_offset_.data(); }
|
||||
size_t OffsetSize() const { return bin_offset_.size(); }
|
||||
size_t Size() const { return data_.size() / (binTypeSize_); }
|
||||
[[nodiscard]] std::uint32_t const* Offset() const { return bin_offset_.data(); }
|
||||
[[nodiscard]] std::size_t OffsetSize() const { return bin_offset_.size(); }
|
||||
[[nodiscard]] std::size_t Size() const { return data_.size() / (binTypeSize_); }
|
||||
|
||||
void Resize(const size_t n_bytes) {
|
||||
data_.resize(n_bytes);
|
||||
}
|
||||
// set the offset used in compression, cut_ptrs is the CSC indptr in HistogramCuts
|
||||
void SetBinOffset(std::vector<uint32_t> const& cut_ptrs) {
|
||||
bin_offset_.resize(cut_ptrs.size() - 1); // resize to number of features.
|
||||
std::copy_n(cut_ptrs.begin(), bin_offset_.size(), bin_offset_.begin());
|
||||
}
|
||||
std::vector<uint8_t>::const_iterator begin() const { // NOLINT
|
||||
return data_.begin();
|
||||
auto begin() const { // NOLINT
|
||||
return data_.data();
|
||||
}
|
||||
std::vector<uint8_t>::const_iterator end() const { // NOLINT
|
||||
return data_.end();
|
||||
auto end() const { // NOLINT
|
||||
return data_.data() + data_.size();
|
||||
}
|
||||
|
||||
std::vector<uint8_t>::iterator begin() { // NOLINT
|
||||
return data_.begin();
|
||||
auto begin() { // NOLINT
|
||||
return data_.data();
|
||||
}
|
||||
std::vector<uint8_t>::iterator end() { // NOLINT
|
||||
return data_.end();
|
||||
auto end() { // NOLINT
|
||||
return data_.data() + data_.size();
|
||||
}
|
||||
|
||||
private:
|
||||
@@ -313,12 +325,12 @@ struct Index {
|
||||
|
||||
using Func = uint32_t (*)(uint8_t const*, size_t);
|
||||
|
||||
std::vector<uint8_t> data_;
|
||||
Span<std::uint8_t> data_;
|
||||
// starting position of each feature inside the cut values (the indptr of the CSC cut matrix
|
||||
// HistogramCuts without the last entry.) Used for bin compression.
|
||||
std::vector<uint32_t> bin_offset_;
|
||||
|
||||
BinTypeSize binTypeSize_ {kUint8BinsTypeSize};
|
||||
BinTypeSize binTypeSize_{kUint8BinsTypeSize};
|
||||
Func func_;
|
||||
};
|
||||
|
||||
|
||||
104
src/common/io.cc
104
src/common/io.cc
@@ -200,21 +200,43 @@ std::string FileExtension(std::string fname, bool lower) {
|
||||
}
|
||||
}
|
||||
|
||||
struct PrivateMmapConstStream::MMAPFile {
|
||||
// For some reason, NVCC 12.1 marks the function deleted if we expose it in the header.
|
||||
// NVCC 11.8 doesn't allow `noexcept(false) = default` altogether.
|
||||
ResourceHandler::~ResourceHandler() noexcept(false) {} // NOLINT
|
||||
|
||||
struct MMAPFile {
|
||||
#if defined(xgboost_IS_WIN)
|
||||
HANDLE fd{INVALID_HANDLE_VALUE};
|
||||
HANDLE file_map{INVALID_HANDLE_VALUE};
|
||||
#else
|
||||
std::int32_t fd{0};
|
||||
#endif
|
||||
char* base_ptr{nullptr};
|
||||
std::byte* base_ptr{nullptr};
|
||||
std::size_t base_size{0};
|
||||
std::size_t delta{0};
|
||||
std::string path;
|
||||
|
||||
MMAPFile() = default;
|
||||
|
||||
#if defined(xgboost_IS_WIN)
|
||||
MMAPFile(HANDLE fd, HANDLE fm, std::byte* base_ptr, std::size_t base_size, std::size_t delta,
|
||||
std::string path)
|
||||
: fd{fd},
|
||||
file_map{fm},
|
||||
base_ptr{base_ptr},
|
||||
base_size{base_size},
|
||||
delta{delta},
|
||||
path{std::move(path)} {}
|
||||
#else
|
||||
MMAPFile(std::int32_t fd, std::byte* base_ptr, std::size_t base_size, std::size_t delta,
|
||||
std::string path)
|
||||
: fd{fd}, base_ptr{base_ptr}, base_size{base_size}, delta{delta}, path{std::move(path)} {}
|
||||
#endif
|
||||
};
|
||||
|
||||
char* PrivateMmapConstStream::Open(std::string path, std::size_t offset, std::size_t length) {
|
||||
std::unique_ptr<MMAPFile> Open(std::string path, std::size_t offset, std::size_t length) {
|
||||
if (length == 0) {
|
||||
return nullptr;
|
||||
return std::make_unique<MMAPFile>();
|
||||
}
|
||||
|
||||
#if defined(xgboost_IS_WIN)
|
||||
@@ -226,16 +248,18 @@ char* PrivateMmapConstStream::Open(std::string path, std::size_t offset, std::si
|
||||
CHECK_GE(fd, 0) << "Failed to open:" << path << ". " << SystemErrorMsg();
|
||||
#endif
|
||||
|
||||
char* ptr{nullptr};
|
||||
std::byte* ptr{nullptr};
|
||||
// Round down for alignment.
|
||||
auto view_start = offset / GetMmapAlignment() * GetMmapAlignment();
|
||||
auto view_size = length + (offset - view_start);
|
||||
|
||||
#if defined(__linux__) || defined(__GLIBC__)
|
||||
int prot{PROT_READ};
|
||||
ptr = reinterpret_cast<char*>(mmap64(nullptr, view_size, prot, MAP_PRIVATE, fd, view_start));
|
||||
ptr = reinterpret_cast<std::byte*>(mmap64(nullptr, view_size, prot, MAP_PRIVATE, fd, view_start));
|
||||
madvise(ptr, view_size, MADV_WILLNEED);
|
||||
CHECK_NE(ptr, MAP_FAILED) << "Failed to map: " << path << ". " << SystemErrorMsg();
|
||||
handle_.reset(new MMAPFile{fd, ptr, view_size, std::move(path)});
|
||||
auto handle =
|
||||
std::make_unique<MMAPFile>(fd, ptr, view_size, offset - view_start, std::move(path));
|
||||
#elif defined(xgboost_IS_WIN)
|
||||
auto file_size = GetFileSize(fd, nullptr);
|
||||
DWORD access = PAGE_READONLY;
|
||||
@@ -244,33 +268,32 @@ char* PrivateMmapConstStream::Open(std::string path, std::size_t offset, std::si
|
||||
std::uint32_t loff = static_cast<std::uint32_t>(view_start);
|
||||
std::uint32_t hoff = view_start >> 32;
|
||||
CHECK(map_file) << "Failed to map: " << path << ". " << SystemErrorMsg();
|
||||
ptr = reinterpret_cast<char*>(MapViewOfFile(map_file, access, hoff, loff, view_size));
|
||||
ptr = reinterpret_cast<std::byte*>(MapViewOfFile(map_file, access, hoff, loff, view_size));
|
||||
CHECK_NE(ptr, nullptr) << "Failed to map: " << path << ". " << SystemErrorMsg();
|
||||
handle_.reset(new MMAPFile{fd, map_file, ptr, view_size, std::move(path)});
|
||||
auto handle = std::make_unique<MMAPFile>(fd, map_file, ptr, view_size, offset - view_start,
|
||||
std::move(path));
|
||||
#else
|
||||
CHECK_LE(offset, std::numeric_limits<off_t>::max())
|
||||
<< "File size has exceeded the limit on the current system.";
|
||||
int prot{PROT_READ};
|
||||
ptr = reinterpret_cast<char*>(mmap(nullptr, view_size, prot, MAP_PRIVATE, fd, view_start));
|
||||
ptr = reinterpret_cast<std::byte*>(mmap(nullptr, view_size, prot, MAP_PRIVATE, fd, view_start));
|
||||
CHECK_NE(ptr, MAP_FAILED) << "Failed to map: " << path << ". " << SystemErrorMsg();
|
||||
handle_.reset(new MMAPFile{fd, ptr, view_size, std::move(path)});
|
||||
auto handle =
|
||||
std::make_unique<MMAPFile>(fd, ptr, view_size, offset - view_start, std::move(path));
|
||||
#endif // defined(__linux__)
|
||||
|
||||
ptr += (offset - view_start);
|
||||
return ptr;
|
||||
return handle;
|
||||
}
|
||||
|
||||
PrivateMmapConstStream::PrivateMmapConstStream(std::string path, std::size_t offset,
|
||||
std::size_t length)
|
||||
: MemoryFixSizeBuffer{}, handle_{nullptr} {
|
||||
this->p_buffer_ = Open(std::move(path), offset, length);
|
||||
this->buffer_size_ = length;
|
||||
}
|
||||
MmapResource::MmapResource(std::string path, std::size_t offset, std::size_t length)
|
||||
: ResourceHandler{kMmap}, handle_{Open(std::move(path), offset, length)}, n_{length} {}
|
||||
|
||||
PrivateMmapConstStream::~PrivateMmapConstStream() {
|
||||
CHECK(handle_);
|
||||
MmapResource::~MmapResource() noexcept(false) {
|
||||
if (!handle_) {
|
||||
return;
|
||||
}
|
||||
#if defined(xgboost_IS_WIN)
|
||||
if (p_buffer_) {
|
||||
if (handle_->base_ptr) {
|
||||
CHECK(UnmapViewOfFile(handle_->base_ptr)) "Faled to call munmap: " << SystemErrorMsg();
|
||||
}
|
||||
if (handle_->fd != INVALID_HANDLE_VALUE) {
|
||||
@@ -290,6 +313,43 @@ PrivateMmapConstStream::~PrivateMmapConstStream() {
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
[[nodiscard]] void* MmapResource::Data() {
|
||||
if (!handle_) {
|
||||
return nullptr;
|
||||
}
|
||||
return handle_->base_ptr + handle_->delta;
|
||||
}
|
||||
|
||||
[[nodiscard]] std::size_t MmapResource::Size() const { return n_; }
|
||||
|
||||
// For some reason, NVCC 12.1 marks the function deleted if we expose it in the header.
|
||||
// NVCC 11.8 doesn't allow `noexcept(false) = default` altogether.
|
||||
AlignedResourceReadStream::~AlignedResourceReadStream() noexcept(false) {} // NOLINT
|
||||
PrivateMmapConstStream::~PrivateMmapConstStream() noexcept(false) {} // NOLINT
|
||||
|
||||
AlignedFileWriteStream::AlignedFileWriteStream(StringView path, StringView flags)
|
||||
: pimpl_{dmlc::Stream::Create(path.c_str(), flags.c_str())} {}
|
||||
|
||||
[[nodiscard]] std::size_t AlignedFileWriteStream::DoWrite(const void* ptr,
|
||||
std::size_t n_bytes) noexcept(true) {
|
||||
pimpl_->Write(ptr, n_bytes);
|
||||
return n_bytes;
|
||||
}
|
||||
|
||||
AlignedMemWriteStream::AlignedMemWriteStream(std::string* p_buf)
|
||||
: pimpl_{std::make_unique<MemoryBufferStream>(p_buf)} {}
|
||||
AlignedMemWriteStream::~AlignedMemWriteStream() = default;
|
||||
|
||||
[[nodiscard]] std::size_t AlignedMemWriteStream::DoWrite(const void* ptr,
|
||||
std::size_t n_bytes) noexcept(true) {
|
||||
this->pimpl_->Write(ptr, n_bytes);
|
||||
return n_bytes;
|
||||
}
|
||||
|
||||
[[nodiscard]] std::size_t AlignedMemWriteStream::Tell() const noexcept(true) {
|
||||
return this->pimpl_->Tell();
|
||||
}
|
||||
} // namespace xgboost::common
|
||||
|
||||
#if defined(xgboost_IS_WIN)
|
||||
|
||||
336
src/common/io.h
336
src/common/io.h
@@ -4,22 +4,29 @@
|
||||
* \brief general stream interface for serialization, I/O
|
||||
* \author Tianqi Chen
|
||||
*/
|
||||
|
||||
#ifndef XGBOOST_COMMON_IO_H_
|
||||
#define XGBOOST_COMMON_IO_H_
|
||||
|
||||
#include <dmlc/io.h>
|
||||
#include <rabit/rabit.h>
|
||||
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <memory> // for unique_ptr
|
||||
#include <string> // for string
|
||||
#include <algorithm> // for min
|
||||
#include <array> // for array
|
||||
#include <cstddef> // for byte, size_t
|
||||
#include <cstdlib> // for malloc, realloc, free
|
||||
#include <cstring> // for memcpy
|
||||
#include <fstream> // for ifstream
|
||||
#include <limits> // for numeric_limits
|
||||
#include <memory> // for unique_ptr
|
||||
#include <string> // for string
|
||||
#include <type_traits> // for alignment_of_v, enable_if_t
|
||||
#include <utility> // for move
|
||||
#include <vector> // for vector
|
||||
|
||||
#include "common.h"
|
||||
#include "xgboost/string_view.h" // for StringView
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
namespace xgboost::common {
|
||||
using MemoryFixSizeBuffer = rabit::utils::MemoryFixSizeBuffer;
|
||||
using MemoryBufferStream = rabit::utils::MemoryBufferStream;
|
||||
|
||||
@@ -58,8 +65,8 @@ class FixedSizeStream : public PeekableInStream {
|
||||
|
||||
size_t Read(void* dptr, size_t size) override;
|
||||
size_t PeekRead(void* dptr, size_t size) override;
|
||||
size_t Size() const { return buffer_.size(); }
|
||||
size_t Tell() const { return pointer_; }
|
||||
[[nodiscard]] std::size_t Size() const { return buffer_.size(); }
|
||||
[[nodiscard]] std::size_t Tell() const { return pointer_; }
|
||||
void Seek(size_t pos);
|
||||
|
||||
void Write(const void*, size_t) override {
|
||||
@@ -129,18 +136,245 @@ inline std::string ReadAll(std::string const &path) {
|
||||
return content;
|
||||
}
|
||||
|
||||
struct MMAPFile;
|
||||
|
||||
/**
|
||||
* @brief Handler for one-shot resource. Unlike `std::pmr::*`, the resource handler is
|
||||
* fixed once it's constructed. Users cannot use mutable operations like resize
|
||||
* without acquiring the specific resource first.
|
||||
*/
|
||||
class ResourceHandler {
|
||||
public:
|
||||
// RTTI
|
||||
enum Kind : std::uint8_t {
|
||||
kMalloc = 0,
|
||||
kMmap = 1,
|
||||
};
|
||||
|
||||
private:
|
||||
Kind kind_{kMalloc};
|
||||
|
||||
public:
|
||||
virtual void* Data() = 0;
|
||||
template <typename T>
|
||||
[[nodiscard]] T* DataAs() {
|
||||
return reinterpret_cast<T*>(this->Data());
|
||||
}
|
||||
|
||||
[[nodiscard]] virtual std::size_t Size() const = 0;
|
||||
[[nodiscard]] auto Type() const { return kind_; }
|
||||
|
||||
// Allow exceptions for cleaning up resource.
|
||||
virtual ~ResourceHandler() noexcept(false);
|
||||
|
||||
explicit ResourceHandler(Kind kind) : kind_{kind} {}
|
||||
// Use shared_ptr to manage a pool like resource handler. All copy and assignment
|
||||
// operators are disabled.
|
||||
ResourceHandler(ResourceHandler const& that) = delete;
|
||||
ResourceHandler& operator=(ResourceHandler const& that) = delete;
|
||||
ResourceHandler(ResourceHandler&& that) = delete;
|
||||
ResourceHandler& operator=(ResourceHandler&& that) = delete;
|
||||
/**
|
||||
* @brief Wether two resources have the same type. (both malloc or both mmap).
|
||||
*/
|
||||
[[nodiscard]] bool IsSameType(ResourceHandler const& that) const {
|
||||
return this->Type() == that.Type();
|
||||
}
|
||||
};
|
||||
|
||||
class MallocResource : public ResourceHandler {
|
||||
void* ptr_{nullptr};
|
||||
std::size_t n_{0};
|
||||
|
||||
void Clear() noexcept(true) {
|
||||
std::free(ptr_);
|
||||
ptr_ = nullptr;
|
||||
n_ = 0;
|
||||
}
|
||||
|
||||
public:
|
||||
explicit MallocResource(std::size_t n_bytes) : ResourceHandler{kMalloc} { this->Resize(n_bytes); }
|
||||
~MallocResource() noexcept(true) override { this->Clear(); }
|
||||
|
||||
void* Data() override { return ptr_; }
|
||||
[[nodiscard]] std::size_t Size() const override { return n_; }
|
||||
/**
|
||||
* @brief Resize the resource to n_bytes. Unlike std::vector::resize, it prefers realloc
|
||||
* over malloc.
|
||||
*
|
||||
* @tparam force_malloc Force the use of malloc over realloc. Used for testing.
|
||||
*
|
||||
* @param n_bytes The new size.
|
||||
*/
|
||||
template <bool force_malloc = false>
|
||||
void Resize(std::size_t n_bytes) {
|
||||
// realloc(ptr, 0) works, but is deprecated.
|
||||
if (n_bytes == 0) {
|
||||
this->Clear();
|
||||
return;
|
||||
}
|
||||
|
||||
// If realloc fails, we need to copy the data ourselves.
|
||||
bool need_copy{false};
|
||||
void* new_ptr{nullptr};
|
||||
// use realloc first, it can handle nullptr.
|
||||
if constexpr (!force_malloc) {
|
||||
new_ptr = std::realloc(ptr_, n_bytes);
|
||||
}
|
||||
// retry with malloc if realloc fails
|
||||
if (!new_ptr) {
|
||||
// ptr_ is preserved if realloc fails
|
||||
new_ptr = std::malloc(n_bytes);
|
||||
need_copy = true;
|
||||
}
|
||||
if (!new_ptr) {
|
||||
// malloc fails
|
||||
LOG(FATAL) << "bad_malloc: Failed to allocate " << n_bytes << " bytes.";
|
||||
}
|
||||
|
||||
if (need_copy) {
|
||||
std::copy_n(reinterpret_cast<std::byte*>(ptr_), n_, reinterpret_cast<std::byte*>(new_ptr));
|
||||
}
|
||||
// default initialize
|
||||
std::memset(reinterpret_cast<std::byte*>(new_ptr) + n_, '\0', n_bytes - n_);
|
||||
// free the old ptr if malloc is used.
|
||||
if (need_copy) {
|
||||
this->Clear();
|
||||
}
|
||||
|
||||
ptr_ = new_ptr;
|
||||
n_ = n_bytes;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief A class for wrapping mmap as a resource for RAII.
|
||||
*/
|
||||
class MmapResource : public ResourceHandler {
|
||||
std::unique_ptr<MMAPFile> handle_;
|
||||
std::size_t n_;
|
||||
|
||||
public:
|
||||
MmapResource(std::string path, std::size_t offset, std::size_t length);
|
||||
~MmapResource() noexcept(false) override;
|
||||
|
||||
[[nodiscard]] void* Data() override;
|
||||
[[nodiscard]] std::size_t Size() const override;
|
||||
};
|
||||
|
||||
/**
|
||||
* @param Alignment for resource read stream and aligned write stream.
|
||||
*/
|
||||
constexpr std::size_t IOAlignment() {
|
||||
// For most of the pod types in XGBoost, 8 byte is sufficient.
|
||||
return 8;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Wrap resource into a dmlc stream.
|
||||
*
|
||||
* This class is to facilitate the use of mmap. Caller can optionally use the `Read()`
|
||||
* method or the `Consume()` method. The former copies data into output, while the latter
|
||||
* makes copy only if it's a primitive type.
|
||||
*
|
||||
* Input is required to be aligned to IOAlignment().
|
||||
*/
|
||||
class AlignedResourceReadStream {
|
||||
std::shared_ptr<ResourceHandler> resource_;
|
||||
std::size_t curr_ptr_{0};
|
||||
|
||||
// Similar to SEEK_END in libc
|
||||
static std::size_t constexpr kSeekEnd = std::numeric_limits<std::size_t>::max();
|
||||
|
||||
public:
|
||||
explicit AlignedResourceReadStream(std::shared_ptr<ResourceHandler> resource)
|
||||
: resource_{std::move(resource)} {}
|
||||
|
||||
[[nodiscard]] std::shared_ptr<ResourceHandler> Share() noexcept(true) { return resource_; }
|
||||
/**
|
||||
* @brief Consume n_bytes of data, no copying is performed.
|
||||
*
|
||||
* @return A pair with the beginning pointer and the number of available bytes, which
|
||||
* may be smaller than requested.
|
||||
*/
|
||||
[[nodiscard]] auto Consume(std::size_t n_bytes) noexcept(true) {
|
||||
auto res_size = resource_->Size();
|
||||
auto data = reinterpret_cast<std::byte*>(resource_->Data());
|
||||
auto ptr = data + curr_ptr_;
|
||||
|
||||
// Move the cursor
|
||||
auto aligned_n_bytes = DivRoundUp(n_bytes, IOAlignment()) * IOAlignment();
|
||||
auto aligned_forward = std::min(res_size - curr_ptr_, aligned_n_bytes);
|
||||
std::size_t forward = std::min(res_size - curr_ptr_, n_bytes);
|
||||
|
||||
curr_ptr_ += aligned_forward;
|
||||
|
||||
return std::pair{ptr, forward};
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
[[nodiscard]] auto Consume(T* out) noexcept(false) -> std::enable_if_t<std::is_pod_v<T>, bool> {
|
||||
auto [ptr, size] = this->Consume(sizeof(T));
|
||||
if (size != sizeof(T)) {
|
||||
return false;
|
||||
}
|
||||
CHECK_EQ(reinterpret_cast<std::uintptr_t>(ptr) % std::alignment_of_v<T>, 0);
|
||||
*out = *reinterpret_cast<T*>(ptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
[[nodiscard]] virtual std::size_t Tell() noexcept(true) { return curr_ptr_; }
|
||||
/**
|
||||
* @brief Read n_bytes of data, output is copied into ptr.
|
||||
*/
|
||||
[[nodiscard]] std::size_t Read(void* ptr, std::size_t n_bytes) noexcept(true) {
|
||||
auto [res_ptr, forward] = this->Consume(n_bytes);
|
||||
if (forward != 0) {
|
||||
std::memcpy(ptr, res_ptr, forward);
|
||||
}
|
||||
return forward;
|
||||
}
|
||||
/**
|
||||
* @brief Read a primitive type.
|
||||
*
|
||||
* @return Whether the read is successful.
|
||||
*/
|
||||
template <typename T>
|
||||
[[nodiscard]] auto Read(T* out) noexcept(false) -> std::enable_if_t<std::is_pod_v<T>, bool> {
|
||||
return this->Consume(out);
|
||||
}
|
||||
/**
|
||||
* @brief Read a vector.
|
||||
*
|
||||
* @return Whether the read is successful.
|
||||
*/
|
||||
template <typename T>
|
||||
[[nodiscard]] bool Read(std::vector<T>* out) noexcept(true) {
|
||||
std::uint64_t n{0};
|
||||
if (!this->Consume(&n)) {
|
||||
return false;
|
||||
}
|
||||
out->resize(n);
|
||||
|
||||
auto n_bytes = sizeof(T) * n;
|
||||
if (this->Read(out->data(), n_bytes) != n_bytes) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
virtual ~AlignedResourceReadStream() noexcept(false);
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Private mmap file as a read-only stream.
|
||||
*
|
||||
* It can calculate alignment automatically based on system page size (or allocation
|
||||
* granularity on Windows).
|
||||
*
|
||||
* The file is required to be aligned by IOAlignment().
|
||||
*/
|
||||
class PrivateMmapConstStream : public MemoryFixSizeBuffer {
|
||||
struct MMAPFile;
|
||||
std::unique_ptr<MMAPFile> handle_;
|
||||
|
||||
char* Open(std::string path, std::size_t offset, std::size_t length);
|
||||
|
||||
class PrivateMmapConstStream : public AlignedResourceReadStream {
|
||||
public:
|
||||
/**
|
||||
* @brief Construct a private mmap stream.
|
||||
@@ -149,11 +383,71 @@ class PrivateMmapConstStream : public MemoryFixSizeBuffer {
|
||||
* @param offset See the `offset` parameter of `mmap` for details.
|
||||
* @param length See the `length` parameter of `mmap` for details.
|
||||
*/
|
||||
explicit PrivateMmapConstStream(std::string path, std::size_t offset, std::size_t length);
|
||||
void Write(void const*, std::size_t) override { LOG(FATAL) << "Read-only stream."; }
|
||||
|
||||
~PrivateMmapConstStream() override;
|
||||
explicit PrivateMmapConstStream(std::string path, std::size_t offset, std::size_t length)
|
||||
: AlignedResourceReadStream{std::make_shared<MmapResource>(path, offset, length)} {}
|
||||
~PrivateMmapConstStream() noexcept(false) override;
|
||||
};
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
|
||||
/**
|
||||
* @brief Base class for write stream with alignment defined by IOAlignment().
|
||||
*/
|
||||
class AlignedWriteStream {
|
||||
protected:
|
||||
[[nodiscard]] virtual std::size_t DoWrite(const void* ptr,
|
||||
std::size_t n_bytes) noexcept(true) = 0;
|
||||
|
||||
public:
|
||||
virtual ~AlignedWriteStream() = default;
|
||||
|
||||
[[nodiscard]] std::size_t Write(const void* ptr, std::size_t n_bytes) noexcept(false) {
|
||||
auto aligned_n_bytes = DivRoundUp(n_bytes, IOAlignment()) * IOAlignment();
|
||||
auto w_n_bytes = this->DoWrite(ptr, n_bytes);
|
||||
CHECK_EQ(w_n_bytes, n_bytes);
|
||||
auto remaining = aligned_n_bytes - n_bytes;
|
||||
if (remaining > 0) {
|
||||
std::array<std::uint8_t, IOAlignment()> padding;
|
||||
std::memset(padding.data(), '\0', padding.size());
|
||||
w_n_bytes = this->DoWrite(padding.data(), remaining);
|
||||
CHECK_EQ(w_n_bytes, remaining);
|
||||
}
|
||||
return aligned_n_bytes;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
[[nodiscard]] std::enable_if_t<std::is_pod_v<T>, std::size_t> Write(T const& v) {
|
||||
return this->Write(&v, sizeof(T));
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Output stream backed by a file. Aligned to IOAlignment() bytes.
|
||||
*/
|
||||
class AlignedFileWriteStream : public AlignedWriteStream {
|
||||
std::unique_ptr<dmlc::Stream> pimpl_;
|
||||
|
||||
protected:
|
||||
[[nodiscard]] std::size_t DoWrite(const void* ptr, std::size_t n_bytes) noexcept(true) override;
|
||||
|
||||
public:
|
||||
AlignedFileWriteStream() = default;
|
||||
AlignedFileWriteStream(StringView path, StringView flags);
|
||||
~AlignedFileWriteStream() override = default;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Output stream backed by memory buffer. Aligned to IOAlignment() bytes.
|
||||
*/
|
||||
class AlignedMemWriteStream : public AlignedFileWriteStream {
|
||||
std::unique_ptr<MemoryBufferStream> pimpl_;
|
||||
|
||||
protected:
|
||||
[[nodiscard]] std::size_t DoWrite(const void* ptr, std::size_t n_bytes) noexcept(true) override;
|
||||
|
||||
public:
|
||||
explicit AlignedMemWriteStream(std::string* p_buf);
|
||||
~AlignedMemWriteStream() override;
|
||||
|
||||
[[nodiscard]] std::size_t Tell() const noexcept(true);
|
||||
};
|
||||
} // namespace xgboost::common
|
||||
#endif // XGBOOST_COMMON_IO_H_
|
||||
|
||||
158
src/common/ref_resource_view.h
Normal file
158
src/common/ref_resource_view.h
Normal file
@@ -0,0 +1,158 @@
|
||||
/**
|
||||
* Copyright 2023, XGBoost Contributors
|
||||
*/
|
||||
#ifndef XGBOOST_COMMON_REF_RESOURCE_VIEW_H_
|
||||
#define XGBOOST_COMMON_REF_RESOURCE_VIEW_H_
|
||||
|
||||
#include <algorithm> // for fill_n
|
||||
#include <cstdint> // for uint64_t
|
||||
#include <cstring> // for memcpy
|
||||
#include <memory> // for shared_ptr, make_shared
|
||||
#include <type_traits> // for is_reference_v, remove_reference_t, is_same_v
|
||||
#include <utility> // for swap, move
|
||||
|
||||
#include "io.h" // for ResourceHandler, AlignedResourceReadStream, MallocResource
|
||||
#include "xgboost/logging.h"
|
||||
#include "xgboost/span.h" // for Span
|
||||
|
||||
namespace xgboost::common {
|
||||
/**
|
||||
* @brief A vector-like type that holds a reference counted resource.
|
||||
*
|
||||
* The vector size is immutable after construction. This way we can swap the underlying
|
||||
* resource when needed.
|
||||
*/
|
||||
template <typename T>
|
||||
class RefResourceView {
|
||||
static_assert(!std::is_reference_v<T>);
|
||||
|
||||
public:
|
||||
using value_type = T; // NOLINT
|
||||
using size_type = std::uint64_t; // NOLINT
|
||||
|
||||
private:
|
||||
value_type* ptr_{nullptr};
|
||||
size_type size_{0};
|
||||
std::shared_ptr<common::ResourceHandler> mem_{nullptr};
|
||||
|
||||
public:
|
||||
RefResourceView(value_type* ptr, size_type n, std::shared_ptr<common::ResourceHandler> mem)
|
||||
: ptr_{ptr}, size_{n}, mem_{std::move(mem)} {
|
||||
CHECK_GE(mem_->Size(), n);
|
||||
}
|
||||
/**
|
||||
* @brief Construct a view on ptr with length n. The ptr is held by the mem resource.
|
||||
*
|
||||
* @param ptr The pointer to view.
|
||||
* @param n The length of the view.
|
||||
* @param mem The owner of the pointer.
|
||||
* @param init Initialize the view with this value.
|
||||
*/
|
||||
RefResourceView(value_type* ptr, size_type n, std::shared_ptr<common::ResourceHandler> mem,
|
||||
T const& init)
|
||||
: RefResourceView{ptr, n, mem} {
|
||||
if (n != 0) {
|
||||
std::fill_n(ptr_, n, init);
|
||||
}
|
||||
}
|
||||
|
||||
~RefResourceView() = default;
|
||||
|
||||
RefResourceView() = default;
|
||||
RefResourceView(RefResourceView const& that) = delete;
|
||||
RefResourceView(RefResourceView&& that) = delete;
|
||||
RefResourceView& operator=(RefResourceView const& that) = delete;
|
||||
/**
|
||||
* @brief We allow move assignment for lazy initialization.
|
||||
*/
|
||||
RefResourceView& operator=(RefResourceView&& that) = default;
|
||||
|
||||
[[nodiscard]] size_type size() const { return size_; } // NOLINT
|
||||
[[nodiscard]] size_type size_bytes() const { // NOLINT
|
||||
return Span{data(), size()}.size_bytes();
|
||||
}
|
||||
[[nodiscard]] value_type* data() { return ptr_; }; // NOLINT
|
||||
[[nodiscard]] value_type const* data() const { return ptr_; }; // NOLINT
|
||||
[[nodiscard]] bool empty() const { return size() == 0; } // NOLINT
|
||||
|
||||
[[nodiscard]] auto cbegin() const { return data(); } // NOLINT
|
||||
[[nodiscard]] auto begin() { return data(); } // NOLINT
|
||||
[[nodiscard]] auto begin() const { return cbegin(); } // NOLINT
|
||||
[[nodiscard]] auto cend() const { return data() + size(); } // NOLINT
|
||||
[[nodiscard]] auto end() { return data() + size(); } // NOLINT
|
||||
[[nodiscard]] auto end() const { return cend(); } // NOLINT
|
||||
|
||||
[[nodiscard]] auto const& front() const { return data()[0]; } // NOLINT
|
||||
[[nodiscard]] auto& front() { return data()[0]; } // NOLINT
|
||||
[[nodiscard]] auto const& back() const { return data()[size() - 1]; } // NOLINT
|
||||
[[nodiscard]] auto& back() { return data()[size() - 1]; } // NOLINT
|
||||
|
||||
[[nodiscard]] value_type& operator[](size_type i) { return ptr_[i]; }
|
||||
[[nodiscard]] value_type const& operator[](size_type i) const { return ptr_[i]; }
|
||||
|
||||
/**
|
||||
* @brief Get the underlying resource.
|
||||
*/
|
||||
auto Resource() const { return mem_; }
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Read a vector from stream. Accepts both `std::vector` and `RefResourceView`.
|
||||
*
|
||||
* If the output vector is a referenced counted view, no copying occur.
|
||||
*/
|
||||
template <typename Vec>
|
||||
[[nodiscard]] bool ReadVec(common::AlignedResourceReadStream* fi, Vec* vec) {
|
||||
std::uint64_t n{0};
|
||||
if (!fi->Read(&n)) {
|
||||
return false;
|
||||
}
|
||||
if (n == 0) {
|
||||
return true;
|
||||
}
|
||||
|
||||
using T = typename Vec::value_type;
|
||||
auto expected_bytes = sizeof(T) * n;
|
||||
|
||||
auto [ptr, n_bytes] = fi->Consume(expected_bytes);
|
||||
if (n_bytes != expected_bytes) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if constexpr (std::is_same_v<Vec, RefResourceView<T>>) {
|
||||
*vec = RefResourceView<T>{reinterpret_cast<T*>(ptr), n, fi->Share()};
|
||||
} else {
|
||||
vec->resize(n);
|
||||
std::memcpy(vec->data(), ptr, n_bytes);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Write a vector to stream. Accepts both `std::vector` and `RefResourceView`.
|
||||
*/
|
||||
template <typename Vec>
|
||||
[[nodiscard]] std::size_t WriteVec(AlignedFileWriteStream* fo, Vec const& vec) {
|
||||
std::size_t bytes{0};
|
||||
auto n = static_cast<std::uint64_t>(vec.size());
|
||||
bytes += fo->Write(n);
|
||||
if (n == 0) {
|
||||
return sizeof(n);
|
||||
}
|
||||
|
||||
using T = typename std::remove_reference_t<decltype(vec)>::value_type;
|
||||
bytes += fo->Write(vec.data(), vec.size() * sizeof(T));
|
||||
|
||||
return bytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Make a fixed size `RefResourceView` with malloc resource.
|
||||
*/
|
||||
template <typename T>
|
||||
[[nodiscard]] RefResourceView<T> MakeFixedVecWithMalloc(std::size_t n_elements, T const& init) {
|
||||
auto resource = std::make_shared<common::MallocResource>(n_elements * sizeof(T));
|
||||
return RefResourceView{resource->DataAs<T>(), n_elements, resource, init};
|
||||
}
|
||||
} // namespace xgboost::common
|
||||
#endif // XGBOOST_COMMON_REF_RESOURCE_VIEW_H_
|
||||
Reference in New Issue
Block a user