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
2023-06-27 19:05:46 +08:00
committed by GitHub
parent 96c3071a8a
commit bc267dd729
29 changed files with 1448 additions and 509 deletions

View File

@@ -1,60 +1,59 @@
/*!
* Copyright 2019-2021 XGBoost contributors
/**
* Copyright 2019-2023, XGBoost contributors
*/
#include <xgboost/data.h>
#include <dmlc/registry.h>
#include <cstddef> // for size_t
#include "../common/io.h" // for AlignedResourceReadStream, AlignedFileWriteStream
#include "../common/ref_resource_view.h" // for ReadVec, WriteVec
#include "ellpack_page.cuh"
#include "sparse_page_writer.h"
#include "histogram_cut_format.h"
namespace xgboost {
namespace data {
#include "histogram_cut_format.h" // for ReadHistogramCuts, WriteHistogramCuts
#include "sparse_page_writer.h" // for SparsePageFormat
namespace xgboost::data {
DMLC_REGISTRY_FILE_TAG(ellpack_page_raw_format);
class EllpackPageRawFormat : public SparsePageFormat<EllpackPage> {
public:
bool Read(EllpackPage* page, dmlc::SeekStream* fi) override {
bool Read(EllpackPage* page, common::AlignedResourceReadStream* fi) override {
auto* impl = page->Impl();
if (!ReadHistogramCuts(&impl->Cuts(), fi)) {
return false;
}
fi->Read(&impl->n_rows);
fi->Read(&impl->is_dense);
fi->Read(&impl->row_stride);
fi->Read(&impl->gidx_buffer.HostVector());
if (!fi->Read(&impl->n_rows)) {
return false;
}
if (!fi->Read(&impl->is_dense)) {
return false;
}
if (!fi->Read(&impl->row_stride)) {
return false;
}
if (!common::ReadVec(fi, &impl->gidx_buffer.HostVector())) {
return false;
}
if (!fi->Read(&impl->base_rowid)) {
return false;
}
return true;
}
size_t Write(const EllpackPage& page, dmlc::Stream* fo) override {
size_t bytes = 0;
size_t Write(const EllpackPage& page, common::AlignedFileWriteStream* fo) override {
std::size_t bytes{0};
auto* impl = page.Impl();
bytes += WriteHistogramCuts(impl->Cuts(), fo);
fo->Write(impl->n_rows);
bytes += sizeof(impl->n_rows);
fo->Write(impl->is_dense);
bytes += sizeof(impl->is_dense);
fo->Write(impl->row_stride);
bytes += sizeof(impl->row_stride);
bytes += fo->Write(impl->n_rows);
bytes += fo->Write(impl->is_dense);
bytes += fo->Write(impl->row_stride);
CHECK(!impl->gidx_buffer.ConstHostVector().empty());
fo->Write(impl->gidx_buffer.HostVector());
bytes += impl->gidx_buffer.ConstHostSpan().size_bytes() + sizeof(uint64_t);
fo->Write(impl->base_rowid);
bytes += sizeof(impl->base_rowid);
bytes += common::WriteVec(fo, impl->gidx_buffer.HostVector());
bytes += fo->Write(impl->base_rowid);
return bytes;
}
};
XGBOOST_REGISTER_ELLPACK_PAGE_FORMAT(raw)
.describe("Raw ELLPACK binary data format.")
.set_body([]() {
return new EllpackPageRawFormat();
});
} // namespace data
} // namespace xgboost
.set_body([]() { return new EllpackPageRawFormat(); });
} // namespace xgboost::data

View File

@@ -29,7 +29,7 @@ GHistIndexMatrix::GHistIndexMatrix(Context const *ctx, DMatrix *p_fmat, bst_bin_
cut = common::SketchOnDMatrix(ctx, p_fmat, max_bins_per_feat, sorted_sketch, hess);
const uint32_t nbins = cut.Ptrs().back();
hit_count.resize(nbins, 0);
hit_count = common::MakeFixedVecWithMalloc(nbins, std::size_t{0});
hit_count_tloc_.resize(ctx->Threads() * nbins, 0);
size_t new_size = 1;
@@ -37,8 +37,7 @@ GHistIndexMatrix::GHistIndexMatrix(Context const *ctx, DMatrix *p_fmat, bst_bin_
new_size += batch.Size();
}
row_ptr.resize(new_size);
row_ptr[0] = 0;
row_ptr = common::MakeFixedVecWithMalloc(new_size, std::size_t{0});
const bool isDense = p_fmat->IsDense();
this->isDense_ = isDense;
@@ -61,8 +60,8 @@ GHistIndexMatrix::GHistIndexMatrix(Context const *ctx, DMatrix *p_fmat, bst_bin_
GHistIndexMatrix::GHistIndexMatrix(MetaInfo const &info, common::HistogramCuts &&cuts,
bst_bin_t max_bin_per_feat)
: row_ptr(info.num_row_ + 1, 0),
hit_count(cuts.TotalBins(), 0),
: row_ptr{common::MakeFixedVecWithMalloc(info.num_row_ + 1, std::size_t{0})},
hit_count{common::MakeFixedVecWithMalloc(cuts.TotalBins(), std::size_t{0})},
cut{std::forward<common::HistogramCuts>(cuts)},
max_numeric_bins_per_feat(max_bin_per_feat),
isDense_{info.num_col_ * info.num_row_ == info.num_nonzero_} {}
@@ -95,12 +94,10 @@ GHistIndexMatrix::GHistIndexMatrix(SparsePage const &batch, common::Span<Feature
isDense_{isDense} {
CHECK_GE(n_threads, 1);
CHECK_EQ(row_ptr.size(), 0);
// The number of threads is pegged to the batch size. If the OMP
// block is parallelized on anything other than the batch/block size,
// it should be reassigned
row_ptr.resize(batch.Size() + 1, 0);
row_ptr = common::MakeFixedVecWithMalloc(batch.Size() + 1, std::size_t{0});
const uint32_t nbins = cut.Ptrs().back();
hit_count.resize(nbins, 0);
hit_count = common::MakeFixedVecWithMalloc(nbins, std::size_t{0});
hit_count_tloc_.resize(n_threads * nbins, 0);
this->PushBatch(batch, ft, n_threads);
@@ -128,20 +125,45 @@ INSTANTIATION_PUSH(data::SparsePageAdapterBatch)
#undef INSTANTIATION_PUSH
void GHistIndexMatrix::ResizeIndex(const size_t n_index, const bool isDense) {
auto make_index = [this, n_index](auto t, common::BinTypeSize t_size) {
// Must resize instead of allocating a new one. This function is called everytime a
// new batch is pushed, and we grow the size accordingly without loosing the data the
// previous batches.
using T = decltype(t);
std::size_t n_bytes = sizeof(T) * n_index;
CHECK_GE(n_bytes, this->data.size());
auto resource = this->data.Resource();
decltype(this->data) new_vec;
if (!resource) {
CHECK(this->data.empty());
new_vec = common::MakeFixedVecWithMalloc(n_bytes, std::uint8_t{0});
} else {
CHECK(resource->Type() == common::ResourceHandler::kMalloc);
auto malloc_resource = std::dynamic_pointer_cast<common::MallocResource>(resource);
CHECK(malloc_resource);
malloc_resource->Resize(n_bytes);
// gcc-11.3 doesn't work if DataAs is used.
std::uint8_t *new_ptr = reinterpret_cast<std::uint8_t *>(malloc_resource->Data());
new_vec = {new_ptr, n_bytes / sizeof(std::uint8_t), malloc_resource};
}
this->data = std::move(new_vec);
this->index = common::Index{common::Span{data.data(), data.size()}, t_size};
};
if ((MaxNumBinPerFeat() - 1 <= static_cast<int>(std::numeric_limits<uint8_t>::max())) &&
isDense) {
// compress dense index to uint8
index.SetBinTypeSize(common::kUint8BinsTypeSize);
index.Resize((sizeof(uint8_t)) * n_index);
make_index(std::uint8_t{}, common::kUint8BinsTypeSize);
} else if ((MaxNumBinPerFeat() - 1 > static_cast<int>(std::numeric_limits<uint8_t>::max()) &&
MaxNumBinPerFeat() - 1 <= static_cast<int>(std::numeric_limits<uint16_t>::max())) &&
isDense) {
// compress dense index to uint16
index.SetBinTypeSize(common::kUint16BinsTypeSize);
index.Resize((sizeof(uint16_t)) * n_index);
make_index(std::uint16_t{}, common::kUint16BinsTypeSize);
} else {
index.SetBinTypeSize(common::kUint32BinsTypeSize);
index.Resize((sizeof(uint32_t)) * n_index);
// no compression
make_index(std::uint32_t{}, common::kUint32BinsTypeSize);
}
}
@@ -214,11 +236,11 @@ float GHistIndexMatrix::GetFvalue(std::vector<std::uint32_t> const &ptrs,
return std::numeric_limits<float>::quiet_NaN();
}
bool GHistIndexMatrix::ReadColumnPage(dmlc::SeekStream *fi) {
bool GHistIndexMatrix::ReadColumnPage(common::AlignedResourceReadStream *fi) {
return this->columns_->Read(fi, this->cut.Ptrs().data());
}
size_t GHistIndexMatrix::WriteColumnPage(dmlc::Stream *fo) const {
std::size_t GHistIndexMatrix::WriteColumnPage(common::AlignedFileWriteStream *fo) const {
return this->columns_->Write(fo);
}
} // namespace xgboost

View File

@@ -1,5 +1,5 @@
/*!
* Copyright 2022 by XGBoost Contributors
/**
* Copyright 2022-2023, XGBoost Contributors
*/
#include <memory> // std::unique_ptr
@@ -41,9 +41,9 @@ void SetIndexData(Context const* ctx, EllpackPageImpl const* page,
}
void GetRowPtrFromEllpack(Context const* ctx, EllpackPageImpl const* page,
std::vector<size_t>* p_out) {
common::RefResourceView<std::size_t>* p_out) {
auto& row_ptr = *p_out;
row_ptr.resize(page->Size() + 1, 0);
row_ptr = common::MakeFixedVecWithMalloc(page->Size() + 1, std::size_t{0});
if (page->is_dense) {
std::fill(row_ptr.begin() + 1, row_ptr.end(), page->row_stride);
} else {
@@ -95,7 +95,7 @@ GHistIndexMatrix::GHistIndexMatrix(Context const* ctx, MetaInfo const& info,
ctx, page, &hit_count_tloc_, [&](auto bin_idx, auto) { return bin_idx; }, this);
}
this->hit_count.resize(n_bins_total, 0);
this->hit_count = common::MakeFixedVecWithMalloc(n_bins_total, std::size_t{0});
this->GatherHitCount(ctx->Threads(), n_bins_total);
// sanity checks

View File

@@ -9,13 +9,14 @@
#include <atomic> // for atomic
#include <cinttypes> // for uint32_t
#include <cstddef> // for size_t
#include <memory>
#include <memory> // for make_unique
#include <vector>
#include "../common/categorical.h"
#include "../common/error_msg.h" // for InfInData
#include "../common/hist_util.h"
#include "../common/numeric.h"
#include "../common/ref_resource_view.h" // for RefResourceView
#include "../common/threading_utils.h"
#include "../common/transform_iterator.h" // for MakeIndexTransformIter
#include "adapter.h"
@@ -25,9 +26,11 @@
namespace xgboost {
namespace common {
class ColumnMatrix;
class AlignedFileWriteStream;
} // namespace common
/*!
* \brief preprocessed global index matrix, in CSR format
/**
* @brief preprocessed global index matrix, in CSR format.
*
* Transform floating values to integer index in histogram This is a global histogram
* index for CPU histogram. On GPU ellpack page is used.
@@ -133,20 +136,22 @@ class GHistIndexMatrix {
}
public:
/*! \brief row pointer to rows by element position */
std::vector<size_t> row_ptr;
/*! \brief The index data */
/** @brief row pointer to rows by element position */
common::RefResourceView<std::size_t> row_ptr;
/** @brief data storage for index. */
common::RefResourceView<std::uint8_t> data;
/** @brief The histogram index. */
common::Index index;
/*! \brief hit count of each index, used for constructing the ColumnMatrix */
std::vector<size_t> hit_count;
/*! \brief The corresponding cuts */
/** @brief hit count of each index, used for constructing the ColumnMatrix */
common::RefResourceView<std::size_t> hit_count;
/** @brief The corresponding cuts */
common::HistogramCuts cut;
/** \brief max_bin for each feature. */
/** @brief max_bin for each feature. */
bst_bin_t max_numeric_bins_per_feat;
/*! \brief base row index for current page (used by external memory) */
size_t base_rowid{0};
/** @brief base row index for current page (used by external memory) */
bst_row_t base_rowid{0};
bst_bin_t MaxNumBinPerFeat() const {
[[nodiscard]] bst_bin_t MaxNumBinPerFeat() const {
return std::max(static_cast<bst_bin_t>(cut.MaxCategory() + 1), max_numeric_bins_per_feat);
}
@@ -218,29 +223,27 @@ class GHistIndexMatrix {
}
}
bool IsDense() const {
return isDense_;
}
[[nodiscard]] bool IsDense() const { return isDense_; }
void SetDense(bool is_dense) { isDense_ = is_dense; }
/**
* \brief Get the local row index.
* @brief Get the local row index.
*/
size_t RowIdx(size_t ridx) const { return row_ptr[ridx - base_rowid]; }
[[nodiscard]] std::size_t RowIdx(size_t ridx) const { return row_ptr[ridx - base_rowid]; }
bst_row_t Size() const { return row_ptr.empty() ? 0 : row_ptr.size() - 1; }
bst_feature_t Features() const { return cut.Ptrs().size() - 1; }
[[nodiscard]] bst_row_t Size() const { return row_ptr.empty() ? 0 : row_ptr.size() - 1; }
[[nodiscard]] bst_feature_t Features() const { return cut.Ptrs().size() - 1; }
bool ReadColumnPage(dmlc::SeekStream* fi);
size_t WriteColumnPage(dmlc::Stream* fo) const;
[[nodiscard]] bool ReadColumnPage(common::AlignedResourceReadStream* fi);
[[nodiscard]] std::size_t WriteColumnPage(common::AlignedFileWriteStream* fo) const;
common::ColumnMatrix const& Transpose() const;
[[nodiscard]] common::ColumnMatrix const& Transpose() const;
bst_bin_t GetGindex(size_t ridx, size_t fidx) const;
[[nodiscard]] bst_bin_t GetGindex(size_t ridx, size_t fidx) const;
float GetFvalue(size_t ridx, size_t fidx, bool is_cat) const;
float GetFvalue(std::vector<std::uint32_t> const& ptrs, std::vector<float> const& values,
std::vector<float> const& mins, bst_row_t ridx, bst_feature_t fidx,
bool is_cat) const;
[[nodiscard]] float GetFvalue(size_t ridx, size_t fidx, bool is_cat) const;
[[nodiscard]] float GetFvalue(std::vector<std::uint32_t> const& ptrs,
std::vector<float> const& values, std::vector<float> const& mins,
bst_row_t ridx, bst_feature_t fidx, bool is_cat) const;
private:
std::unique_ptr<common::ColumnMatrix> columns_;
@@ -294,5 +297,5 @@ void AssignColumnBinIndex(GHistIndexMatrix const& page, Fn&& assign) {
}
});
}
} // namespace xgboost
} // namespace xgboost
#endif // XGBOOST_DATA_GRADIENT_INDEX_H_

View File

@@ -1,38 +1,49 @@
/*!
* Copyright 2021-2022 XGBoost contributors
/**
* Copyright 2021-2023 XGBoost contributors
*/
#include "sparse_page_writer.h"
#include "gradient_index.h"
#include "histogram_cut_format.h"
#include <cstddef> // for size_t
#include <cstdint> // for uint8_t
#include <type_traits> // for underlying_type_t
#include <vector> // for vector
namespace xgboost {
namespace data {
#include "../common/io.h" // for AlignedResourceReadStream
#include "../common/ref_resource_view.h" // for ReadVec, WriteVec
#include "gradient_index.h" // for GHistIndexMatrix
#include "histogram_cut_format.h" // for ReadHistogramCuts
#include "sparse_page_writer.h" // for SparsePageFormat
namespace xgboost::data {
class GHistIndexRawFormat : public SparsePageFormat<GHistIndexMatrix> {
public:
bool Read(GHistIndexMatrix* page, dmlc::SeekStream* fi) override {
bool Read(GHistIndexMatrix* page, common::AlignedResourceReadStream* fi) override {
CHECK(fi);
if (!ReadHistogramCuts(&page->cut, fi)) {
return false;
}
// indptr
fi->Read(&page->row_ptr);
// data
std::vector<uint8_t> data;
if (!fi->Read(&data)) {
if (!common::ReadVec(fi, &page->row_ptr)) {
return false;
}
page->index.Resize(data.size());
std::copy(data.cbegin(), data.cend(), page->index.begin());
// bin type
// data
// - bin type
// Old gcc doesn't support reading from enum.
std::underlying_type_t<common::BinTypeSize> uint_bin_type{0};
if (!fi->Read(&uint_bin_type)) {
return false;
}
common::BinTypeSize size_type =
static_cast<common::BinTypeSize>(uint_bin_type);
page->index.SetBinTypeSize(size_type);
common::BinTypeSize size_type = static_cast<common::BinTypeSize>(uint_bin_type);
// - index buffer
if (!common::ReadVec(fi, &page->data)) {
return false;
}
// - index
page->index = common::Index{common::Span{page->data.data(), page->data.size()}, size_type};
// hit count
if (!fi->Read(&page->hit_count)) {
if (!common::ReadVec(fi, &page->hit_count)) {
return false;
}
if (!fi->Read(&page->max_numeric_bins_per_feat)) {
@@ -50,38 +61,33 @@ class GHistIndexRawFormat : public SparsePageFormat<GHistIndexMatrix> {
page->index.SetBinOffset(page->cut.Ptrs());
}
page->ReadColumnPage(fi);
if (!page->ReadColumnPage(fi)) {
return false;
}
return true;
}
size_t Write(GHistIndexMatrix const &page, dmlc::Stream *fo) override {
size_t bytes = 0;
std::size_t Write(GHistIndexMatrix const& page, common::AlignedFileWriteStream* fo) override {
std::size_t bytes = 0;
bytes += WriteHistogramCuts(page.cut, fo);
// indptr
fo->Write(page.row_ptr);
bytes += page.row_ptr.size() * sizeof(decltype(page.row_ptr)::value_type) +
sizeof(uint64_t);
bytes += common::WriteVec(fo, page.row_ptr);
// data
std::vector<uint8_t> data(page.index.begin(), page.index.end());
fo->Write(data);
bytes += data.size() * sizeof(decltype(data)::value_type) + sizeof(uint64_t);
// bin type
std::underlying_type_t<common::BinTypeSize> uint_bin_type =
page.index.GetBinTypeSize();
fo->Write(uint_bin_type);
bytes += sizeof(page.index.GetBinTypeSize());
// - bin type
std::underlying_type_t<common::BinTypeSize> uint_bin_type = page.index.GetBinTypeSize();
bytes += fo->Write(uint_bin_type);
// - index buffer
std::vector<std::uint8_t> data(page.index.begin(), page.index.end());
bytes += fo->Write(static_cast<std::uint64_t>(data.size()));
bytes += fo->Write(data.data(), data.size());
// hit count
fo->Write(page.hit_count);
bytes +=
page.hit_count.size() * sizeof(decltype(page.hit_count)::value_type) +
sizeof(uint64_t);
bytes += common::WriteVec(fo, page.hit_count);
// max_bins, base row, is_dense
fo->Write(page.max_numeric_bins_per_feat);
bytes += sizeof(page.max_numeric_bins_per_feat);
fo->Write(page.base_rowid);
bytes += sizeof(page.base_rowid);
fo->Write(page.IsDense());
bytes += sizeof(page.IsDense());
bytes += fo->Write(page.max_numeric_bins_per_feat);
bytes += fo->Write(page.base_rowid);
bytes += fo->Write(page.IsDense());
bytes += page.WriteColumnPage(fo);
return bytes;
@@ -93,6 +99,4 @@ DMLC_REGISTRY_FILE_TAG(gradient_index_format);
XGBOOST_REGISTER_GHIST_INDEX_PAGE_FORMAT(raw)
.describe("Raw GHistIndex binary data format.")
.set_body([]() { return new GHistIndexRawFormat(); });
} // namespace data
} // namespace xgboost
} // namespace xgboost::data

View File

@@ -1,36 +1,38 @@
/*!
* Copyright 2021 XGBoost contributors
/**
* Copyright 2021-2023, XGBoost contributors
*/
#ifndef XGBOOST_DATA_HISTOGRAM_CUT_FORMAT_H_
#define XGBOOST_DATA_HISTOGRAM_CUT_FORMAT_H_
#include "../common/hist_util.h"
#include <dmlc/io.h> // for Stream
namespace xgboost {
namespace data {
inline bool ReadHistogramCuts(common::HistogramCuts *cuts, dmlc::SeekStream *fi) {
if (!fi->Read(&cuts->cut_values_.HostVector())) {
#include <cstddef> // for size_t
#include "../common/hist_util.h" // for HistogramCuts
#include "../common/io.h" // for AlignedResourceReadStream, AlignedFileWriteStream
#include "../common/ref_resource_view.h" // for WriteVec, ReadVec
namespace xgboost::data {
inline bool ReadHistogramCuts(common::HistogramCuts *cuts, common::AlignedResourceReadStream *fi) {
if (!common::ReadVec(fi, &cuts->cut_values_.HostVector())) {
return false;
}
if (!fi->Read(&cuts->cut_ptrs_.HostVector())) {
if (!common::ReadVec(fi, &cuts->cut_ptrs_.HostVector())) {
return false;
}
if (!fi->Read(&cuts->min_vals_.HostVector())) {
if (!common::ReadVec(fi, &cuts->min_vals_.HostVector())) {
return false;
}
return true;
}
inline size_t WriteHistogramCuts(common::HistogramCuts const &cuts, dmlc::Stream *fo) {
size_t bytes = 0;
fo->Write(cuts.cut_values_.ConstHostVector());
bytes += cuts.cut_values_.ConstHostSpan().size_bytes() + sizeof(uint64_t);
fo->Write(cuts.cut_ptrs_.ConstHostVector());
bytes += cuts.cut_ptrs_.ConstHostSpan().size_bytes() + sizeof(uint64_t);
fo->Write(cuts.min_vals_.ConstHostVector());
bytes += cuts.min_vals_.ConstHostSpan().size_bytes() + sizeof(uint64_t);
inline std::size_t WriteHistogramCuts(common::HistogramCuts const &cuts,
common::AlignedFileWriteStream *fo) {
std::size_t bytes = 0;
bytes += common::WriteVec(fo, cuts.Values());
bytes += common::WriteVec(fo, cuts.Ptrs());
bytes += common::WriteVec(fo, cuts.MinValues());
return bytes;
}
} // namespace data
} // namespace xgboost
} // namespace xgboost::data
#endif // XGBOOST_DATA_HISTOGRAM_CUT_FORMAT_H_

View File

@@ -240,9 +240,9 @@ void IterativeDMatrix::InitFromCPU(Context const* ctx, BatchParam const& p,
* Generate gradient index.
*/
this->ghist_ = std::make_unique<GHistIndexMatrix>(Info(), std::move(cuts), p.max_bin);
size_t rbegin = 0;
size_t prev_sum = 0;
size_t i = 0;
std::size_t rbegin = 0;
std::size_t prev_sum = 0;
std::size_t i = 0;
while (iter.Next()) {
HostAdapterDispatch(proxy, [&](auto const& batch) {
proxy->Info().num_nonzero_ = batch_nnz[i];

View File

@@ -1,59 +1,57 @@
/*!
* Copyright (c) 2015-2021 by Contributors
/**
* Copyright 2015-2023, XGBoost Contributors
* \file sparse_page_raw_format.cc
* Raw binary format of sparse page.
*/
#include <xgboost/data.h>
#include <dmlc/registry.h>
#include "xgboost/logging.h"
#include "../common/io.h" // for AlignedResourceReadStream, AlignedFileWriteStream
#include "../common/ref_resource_view.h" // for WriteVec
#include "./sparse_page_writer.h"
#include "xgboost/data.h"
#include "xgboost/logging.h"
namespace xgboost {
namespace data {
namespace xgboost::data {
DMLC_REGISTRY_FILE_TAG(sparse_page_raw_format);
template<typename T>
template <typename T>
class SparsePageRawFormat : public SparsePageFormat<T> {
public:
bool Read(T* page, dmlc::SeekStream* fi) override {
bool Read(T* page, common::AlignedResourceReadStream* fi) override {
auto& offset_vec = page->offset.HostVector();
if (!fi->Read(&offset_vec)) {
if (!common::ReadVec(fi, &offset_vec)) {
return false;
}
auto& data_vec = page->data.HostVector();
CHECK_NE(page->offset.Size(), 0U) << "Invalid SparsePage file";
data_vec.resize(offset_vec.back());
if (page->data.Size() != 0) {
size_t n_bytes = fi->Read(dmlc::BeginPtr(data_vec),
(page->data).Size() * sizeof(Entry));
CHECK_EQ(n_bytes, (page->data).Size() * sizeof(Entry))
<< "Invalid SparsePage file";
if (!common::ReadVec(fi, &data_vec)) {
return false;
}
}
if (!fi->Read(&page->base_rowid, sizeof(page->base_rowid))) {
return false;
}
fi->Read(&page->base_rowid, sizeof(page->base_rowid));
return true;
}
size_t Write(const T& page, dmlc::Stream* fo) override {
std::size_t Write(const T& page, common::AlignedFileWriteStream* fo) override {
const auto& offset_vec = page.offset.HostVector();
const auto& data_vec = page.data.HostVector();
CHECK(page.offset.Size() != 0 && offset_vec[0] == 0);
CHECK_EQ(offset_vec.back(), page.data.Size());
fo->Write(offset_vec);
auto bytes = page.MemCostBytes();
bytes += sizeof(uint64_t);
std::size_t bytes{0};
bytes += common::WriteVec(fo, offset_vec);
if (page.data.Size() != 0) {
fo->Write(dmlc::BeginPtr(data_vec), page.data.Size() * sizeof(Entry));
bytes += common::WriteVec(fo, data_vec);
}
fo->Write(&page.base_rowid, sizeof(page.base_rowid));
bytes += sizeof(page.base_rowid);
bytes += fo->Write(&page.base_rowid, sizeof(page.base_rowid));
return bytes;
}
private:
/*! \brief external memory column offset */
std::vector<size_t> disk_offset_;
};
XGBOOST_REGISTER_SPARSE_PAGE_FORMAT(raw)
@@ -74,5 +72,4 @@ XGBOOST_REGISTER_SORTED_CSC_PAGE_FORMAT(raw)
return new SparsePageRawFormat<SortedCSCPage>();
});
} // namespace data
} // namespace xgboost
} // namespace xgboost::data

View File

@@ -6,9 +6,11 @@
#define XGBOOST_DATA_SPARSE_PAGE_SOURCE_H_
#include <algorithm> // for min
#include <atomic> // for atomic
#include <future> // for async
#include <map>
#include <memory>
#include <mutex> // for mutex
#include <string>
#include <thread>
#include <utility> // for pair, move
@@ -18,7 +20,6 @@
#include "../common/io.h" // for PrivateMmapConstStream
#include "../common/timer.h" // for Monitor, Timer
#include "adapter.h"
#include "dmlc/common.h" // for OMPException
#include "proxy_dmatrix.h" // for DMatrixProxy
#include "sparse_page_writer.h" // for SparsePageFormat
#include "xgboost/base.h"
@@ -93,6 +94,47 @@ class TryLockGuard {
}
};
// Similar to `dmlc::OMPException`, but doesn't need the threads to be joined before rethrow
class ExceHandler {
std::mutex mutex_;
std::atomic<bool> flag_{false};
std::exception_ptr curr_exce_{nullptr};
public:
template <typename Fn>
decltype(auto) Run(Fn&& fn) noexcept(true) {
try {
return fn();
} catch (dmlc::Error const& e) {
std::lock_guard<std::mutex> guard{mutex_};
if (!curr_exce_) {
curr_exce_ = std::current_exception();
}
flag_ = true;
} catch (std::exception const& e) {
std::lock_guard<std::mutex> guard{mutex_};
if (!curr_exce_) {
curr_exce_ = std::current_exception();
}
flag_ = true;
} catch (...) {
std::lock_guard<std::mutex> guard{mutex_};
if (!curr_exce_) {
curr_exce_ = std::current_exception();
}
flag_ = true;
}
return std::invoke_result_t<Fn>();
}
void Rethrow() noexcept(false) {
if (flag_) {
CHECK(curr_exce_);
std::rethrow_exception(curr_exce_);
}
}
};
/**
* @brief Base class for all page sources. Handles fetching, writing, and iteration.
*/
@@ -122,7 +164,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
// Catching exception in pre-fetch threads to prevent segfault. Not always work though,
// OOM error can be delayed due to lazy commit. On the bright side, if mmap is used then
// OOM error should be rare.
dmlc::OMPException exec_;
ExceHandler exce_;
common::Monitor monitor_;
bool ReadCache() {
@@ -141,7 +183,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
CHECK_GT(n_prefetch_batches, 0) << "total batches:" << n_batches_;
std::size_t fetch_it = count_;
exec_.Rethrow();
exce_.Rethrow();
for (std::size_t i = 0; i < n_prefetch_batches; ++i, ++fetch_it) {
fetch_it %= n_batches_; // ring
@@ -152,7 +194,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
CHECK_LT(fetch_it, cache_info_->offset.size());
ring_->at(fetch_it) = std::async(std::launch::async, [fetch_it, self, this]() {
auto page = std::make_shared<S>();
this->exec_.Run([&] {
this->exce_.Run([&] {
std::unique_ptr<SparsePageFormat<S>> fmt{CreatePageFormat<S>("raw")};
auto name = self->cache_info_->ShardName();
auto [offset, length] = self->cache_info_->View(fetch_it);
@@ -172,7 +214,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
CHECK(!(*ring_)[count_].valid());
monitor_.Stop("Wait");
exec_.Rethrow();
exce_.Rethrow();
return true;
}
@@ -184,11 +226,11 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
std::unique_ptr<SparsePageFormat<S>> fmt{CreatePageFormat<S>("raw")};
auto name = cache_info_->ShardName();
std::unique_ptr<dmlc::Stream> fo;
std::unique_ptr<common::AlignedFileWriteStream> fo;
if (this->Iter() == 0) {
fo.reset(dmlc::Stream::Create(name.c_str(), "wb"));
fo = std::make_unique<common::AlignedFileWriteStream>(StringView{name}, "wb");
} else {
fo.reset(dmlc::Stream::Create(name.c_str(), "ab"));
fo = std::make_unique<common::AlignedFileWriteStream>(StringView{name}, "ab");
}
auto bytes = fmt->Write(*page_, fo.get());

View File

@@ -1,52 +1,44 @@
/*!
* Copyright (c) 2014-2019 by Contributors
/**
* Copyright 2014-2023, XGBoost Contributors
* \file sparse_page_writer.h
* \author Tianqi Chen
*/
#ifndef XGBOOST_DATA_SPARSE_PAGE_WRITER_H_
#define XGBOOST_DATA_SPARSE_PAGE_WRITER_H_
#include <xgboost/data.h>
#include <dmlc/io.h>
#include <vector>
#include <algorithm>
#include <cstring>
#include <string>
#include <utility>
#include <memory>
#include <functional>
#include <functional> // for function
#include <string> // for string
#if DMLC_ENABLE_STD_THREAD
#include <dmlc/concurrency.h>
#include <thread>
#endif // DMLC_ENABLE_STD_THREAD
namespace xgboost {
namespace data {
#include "../common/io.h" // for AlignedResourceReadStream, AlignedFileWriteStream
#include "dmlc/io.h" // for Stream
#include "dmlc/registry.h" // for Registry, FunctionRegEntryBase
#include "xgboost/data.h" // for SparsePage,CSCPage,SortedCSCPage,EllpackPage ...
namespace xgboost::data {
template<typename T>
struct SparsePageFormatReg;
/*!
* \brief Format specification of SparsePage.
/**
* @brief Format specification of various data formats like SparsePage.
*/
template<typename T>
template <typename T>
class SparsePageFormat {
public:
/*! \brief virtual destructor */
virtual ~SparsePageFormat() = default;
/*!
* \brief Load all the segments into page, advance fi to end of the block.
* \param page The data to read page into.
* \param fi the input stream of the file
* \return true of the loading as successful, false if end of file was reached
/**
* @brief Load all the segments into page, advance fi to end of the block.
*
* @param page The data to read page into.
* @param fi the input stream of the file
* @return true of the loading as successful, false if end of file was reached
*/
virtual bool Read(T* page, dmlc::SeekStream* fi) = 0;
/*!
* \brief save the data to fo, when a page was written.
* \param fo output stream
virtual bool Read(T* page, common::AlignedResourceReadStream* fi) = 0;
/**
* @brief save the data to fo, when a page was written.
*
* @param fo output stream
*/
virtual size_t Write(const T& page, dmlc::Stream* fo) = 0;
virtual size_t Write(const T& page, common::AlignedFileWriteStream* fo) = 0;
};
/*!
@@ -105,6 +97,5 @@ struct SparsePageFormatReg
DMLC_REGISTRY_REGISTER(SparsePageFormatReg<GHistIndexMatrix>, \
GHistIndexPageFmt, Name)
} // namespace data
} // namespace xgboost
} // namespace xgboost::data
#endif // XGBOOST_DATA_SPARSE_PAGE_WRITER_H_