From 11d4b6368c23cf49d9bb228fa00e410c558fc86e Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Fri, 4 Apr 2025 02:49:10 -0700 Subject: [PATCH 01/35] optimise data initialisation --- src/common/column_matrix.cc | 6 +-- src/common/column_matrix.h | 94 ++++++++++++++++++++++++++++------ src/common/ref_resource_view.h | 24 +++++++++ src/data/gradient_index.cc | 14 ++--- src/data/gradient_index.h | 8 +-- 5 files changed, 116 insertions(+), 30 deletions(-) diff --git a/src/common/column_matrix.cc b/src/common/column_matrix.cc index 1d44f184033f..c703f938170f 100644 --- a/src/common/column_matrix.cc +++ b/src/common/column_matrix.cc @@ -17,7 +17,7 @@ #include "xgboost/span.h" // for Span namespace xgboost::common { -void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold) { +void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads) { auto const nfeature = gmat.Features(); const size_t nrow = gmat.Size(); // identify type of each column @@ -61,10 +61,10 @@ void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_thres auto storage_size = feature_offsets_.back() * static_cast>(bins_type_size_); - index_ = common::MakeFixedVecWithMalloc(storage_size, std::uint8_t{0}); + index_ = common::MakeFixedVecWithMalloc(storage_size, std::uint8_t{0}, n_threads); if (!all_dense_column) { - row_ind_ = common::MakeFixedVecWithMalloc(feature_offsets_[nfeature], std::size_t{0}); + row_ind_ = common::MakeFixedVecWithMalloc(feature_offsets_[nfeature], std::size_t{0}, n_threads); } // store least bin id for each feature diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 17f3ed4c6824..27bc03ebe176 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -195,7 +195,7 @@ class ColumnMatrix { } }; - void InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold); + void InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads); template void SetBinSparse(BinT bin_id, RIdx rid, bst_feature_t fid, ColumnBinT* local_index) { @@ -214,6 +214,22 @@ class ColumnMatrix { } } + template + void SetBinSparse(BinT bin_id, RIdx rid, bst_feature_t fid, ColumnBinT* local_index, size_t nnz) { + if (type_[fid] == kDenseColumn) { + ColumnBinT* begin = &local_index[feature_offsets_[fid]]; + begin[rid] = bin_id - index_base_[fid]; + // 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[nnz] = bin_id - index_base_[fid]; + row_ind_[feature_offsets_[fid] + nnz] = rid; + } + } + public: // get number of features [[nodiscard]] bst_feature_t GetNumFeature() const { @@ -221,8 +237,8 @@ class ColumnMatrix { } ColumnMatrix() = default; - ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold) { - this->InitStorage(gmat, sparse_threshold); + ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads) { + this->InitStorage(gmat, sparse_threshold, n_threads); } /** @@ -232,7 +248,7 @@ class ColumnMatrix { void InitFromSparse(SparsePage const& page, const GHistIndexMatrix& gmat, double sparse_threshold, int32_t n_threads) { auto batch = data::SparsePageAdapterBatch{page.GetView()}; - this->InitStorage(gmat, sparse_threshold); + this->InitStorage(gmat, sparse_threshold, n_threads); // ignore base row id here as we always has one column matrix for each sparse page. this->PushBatch(n_threads, batch, std::numeric_limits::quiet_NaN(), gmat, 0); } @@ -283,7 +299,7 @@ class ColumnMatrix { SetIndexNoMissing(base_rowid, gmat.index.data(), size, n_features, n_threads); }); } else { - SetIndexMixedColumns(base_rowid, batch, gmat, missing); + SetIndexMixedColumns(base_rowid, batch, gmat, missing, n_threads); } } @@ -349,7 +365,7 @@ class ColumnMatrix { */ template void SetIndexMixedColumns(size_t base_rowid, Batch const& batch, const GHistIndexMatrix& gmat, - float missing) { + float missing, int n_threads) { auto n_features = gmat.Features(); missing_.GrowTo(feature_offsets_[n_features], true); @@ -366,19 +382,65 @@ class ColumnMatrix { using ColumnBinT = decltype(t); ColumnBinT* local_index = reinterpret_cast(index_.data()); size_t const batch_size = batch.Size(); - size_t k{0}; - for (size_t rid = 0; rid < batch_size; ++rid) { - auto line = batch.GetLine(rid); - for (size_t i = 0; i < line.Size(); ++i) { - auto coo = line.GetElement(i); - if (is_valid(coo)) { - auto fid = coo.column_idx; - const uint32_t bin_id = row_index[k]; - SetBinSparse(bin_id, rid + base_rowid, fid, local_index); - ++k; + + dmlc::OMPException exc; + std::vector n_elements((n_threads + 1) * n_features, 0); + std::vector k_offsets(n_threads + 1, 0); + size_t block_size = DivRoundUp(batch_size, n_threads); + #pragma omp parallel num_threads(n_threads) + { + exc.Run([&]() { + int tid = omp_get_thread_num(); + size_t begin = block_size * tid; + size_t end = std::min(begin + block_size, batch_size); + for (size_t rid = begin; rid < end; ++rid) { + const auto& line = batch.GetLine(rid); + for (size_t i = 0; i < line.Size(); ++i) { + auto coo = line.GetElement(i); + if (is_valid(coo)) { + auto fid = coo.column_idx; + n_elements[(tid + 1) * n_features + fid] += 1; + k_offsets[tid + 1] += 1; + } + } } + }); + } + exc.Rethrow(); + + ParallelFor(n_features, n_threads, [&](auto fid) { + for (int tid = 0; tid < n_threads; ++tid) { + n_elements[(tid + 1) * n_features + fid] += + n_elements[tid * n_features + fid]; } + num_nonzeros_[fid] = n_elements[n_threads * n_features + fid]; + }); + std::partial_sum(k_offsets.cbegin(), k_offsets.cend(), k_offsets.begin()); + + #pragma omp parallel num_threads(n_threads) + { + exc.Run([&]() { + int tid = omp_get_thread_num(); + size_t begin = block_size * tid; + size_t end = std::min(begin + block_size, batch_size); + size_t k = 0; + for (size_t rid = begin; rid < end; ++rid) { + const auto& line = batch.GetLine(rid); + for (size_t i = 0; i < line.Size(); ++i) { + auto coo = line.GetElement(i); + if (is_valid(coo)) { + auto fid = coo.column_idx; + const uint32_t bin_id = row_index[k_offsets[tid] + k]; + size_t& nnz = n_elements[tid * n_features + fid]; + SetBinSparse(bin_id, rid + base_rowid, fid, local_index, nnz); + ++k; + ++nnz; + } + } + } + }); } + exc.Rethrow(); }); } diff --git a/src/common/ref_resource_view.h b/src/common/ref_resource_view.h index 3c33a839ab77..6d8a5d91c5ba 100644 --- a/src/common/ref_resource_view.h +++ b/src/common/ref_resource_view.h @@ -164,6 +164,30 @@ template return ref; } +/** + * @brief Make a fixed size `RefResourceView` with malloc resource. + * Use n_threads to initilise the storage + */ +template +[[nodiscard]] RefResourceView MakeFixedVecWithMalloc(std::size_t n_elements, T const& init, + int n_threads) { + if (n_elements < n_threads) return MakeFixedVecWithMalloc(n_elements, init); + + auto resource = std::make_shared(n_elements * sizeof(T)); + auto ref = RefResourceView{resource->DataAs(), n_elements, resource}; + + size_t block_size = n_elements / n_threads + (n_elements % n_threads > 0); + #pragma omp parallel num_threads(n_threads) + { + int tid = omp_get_thread_num(); + auto begin = tid * block_size; + auto end = std::min((tid + 1) * block_size, n_elements); + auto size = end > begin ? end - begin : 0; + std::fill_n(ref.data() + begin, size, init); + } + return ref; +} + template class ReallocVector : public RefResourceView { static_assert(!std::is_reference_v); diff --git a/src/data/gradient_index.cc b/src/data/gradient_index.cc index e9a1a7e329ff..b7844719ef43 100644 --- a/src/data/gradient_index.cc +++ b/src/data/gradient_index.cc @@ -27,7 +27,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 = common::MakeFixedVecWithMalloc(nbins, std::size_t{0}); + hit_count = common::MakeFixedVecWithMalloc(nbins, std::size_t{0}, ctx->Threads()); hit_count_tloc_.resize(ctx->Threads() * nbins, 0); size_t new_size = 1; @@ -35,7 +35,7 @@ GHistIndexMatrix::GHistIndexMatrix(Context const *ctx, DMatrix *p_fmat, bst_bin_ new_size += batch.Size(); } - row_ptr = common::MakeFixedVecWithMalloc(new_size, std::size_t{0}); + row_ptr = common::MakeFixedVecWithMalloc(new_size, std::size_t{0}, ctx->Threads()); const bool isDense = p_fmat->IsDense(); this->isDense_ = isDense; @@ -132,13 +132,13 @@ INSTANTIATION_PUSH(data::SparsePageAdapterBatch) INSTANTIATION_PUSH(data::ColumnarAdapterBatch) #undef INSTANTIATION_PUSH -void GHistIndexMatrix::ResizeColumns(double sparse_thresh) { +void GHistIndexMatrix::ResizeColumns(double sparse_thresh, int n_threads) { CHECK(!std::isnan(sparse_thresh)); - this->columns_ = std::make_unique(*this, sparse_thresh); + this->columns_ = std::make_unique(*this, sparse_thresh, n_threads); } -void GHistIndexMatrix::ResizeIndex(const size_t n_index, const bool isDense) { - auto make_index = [this, n_index](auto t, common::BinTypeSize t_size) { +void GHistIndexMatrix::ResizeIndex(const size_t n_index, const bool isDense, int n_threads) { + auto make_index = [this, n_index, n_threads](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 in // the previous batches. @@ -150,7 +150,7 @@ void GHistIndexMatrix::ResizeIndex(const size_t n_index, const bool isDense) { decltype(this->data) new_vec; if (!resource) { CHECK(this->data.empty()); - new_vec = common::MakeFixedVecWithMalloc(n_bytes, std::uint8_t{0}); + new_vec = common::MakeFixedVecWithMalloc(n_bytes, std::uint8_t{0}, n_threads); } else { CHECK(resource->Type() == common::ResourceHandler::kMalloc); auto malloc_resource = std::dynamic_pointer_cast(resource); diff --git a/src/data/gradient_index.h b/src/data/gradient_index.h index 6560198093c1..26c1496bdbd8 100644 --- a/src/data/gradient_index.h +++ b/src/data/gradient_index.h @@ -121,7 +121,7 @@ class GHistIndexMatrix { auto n_bins_total = cut.TotalBins(); const size_t n_index = row_ptr[rbegin + batch.Size()]; // number of entries in this page - ResizeIndex(n_index, isDense_); + ResizeIndex(n_index, isDense_, n_threads); if (isDense_) { index.SetBinOffset(cut.Ptrs()); } @@ -142,7 +142,7 @@ class GHistIndexMatrix { } // The function is only created to avoid using the column matrix in the header. - void ResizeColumns(double sparse_thresh); + void ResizeColumns(double sparse_thresh, int n_threads); public: /** @brief row pointer to rows by element position */ @@ -224,7 +224,7 @@ class GHistIndexMatrix { if (rbegin + batch.Size() == n_samples_total) { // finished - this->ResizeColumns(sparse_thresh); + this->ResizeColumns(sparse_thresh, ctx->Threads()); } } @@ -233,7 +233,7 @@ class GHistIndexMatrix { void PushAdapterBatchColumns(Context const* ctx, Batch const& batch, float missing, size_t rbegin); - void ResizeIndex(const size_t n_index, const bool isDense); + void ResizeIndex(const size_t n_index, const bool isDense, int n_threads); void GetFeatureCounts(size_t* counts) const { auto nfeature = cut.Ptrs().size() - 1; From 3464f8c996d83208fcce07075e60bfb030e347c2 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Mon, 7 Apr 2025 05:44:08 -0700 Subject: [PATCH 02/35] linting --- dmlc-core | 2 +- src/common/column_matrix.cc | 6 ++++-- src/common/column_matrix.h | 3 ++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/dmlc-core b/dmlc-core index 8986b0598df7..133418575498 160000 --- a/dmlc-core +++ b/dmlc-core @@ -1 +1 @@ -Subproject commit 8986b0598df709117570984571476c3614f55724 +Subproject commit 13341857549852a9a86b1894b5ba84c6276ab381 diff --git a/src/common/column_matrix.cc b/src/common/column_matrix.cc index c703f938170f..21271cbda28a 100644 --- a/src/common/column_matrix.cc +++ b/src/common/column_matrix.cc @@ -17,7 +17,8 @@ #include "xgboost/span.h" // for Span namespace xgboost::common { -void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads) { +void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold, + int n_threads) { auto const nfeature = gmat.Features(); const size_t nrow = gmat.Size(); // identify type of each column @@ -64,7 +65,8 @@ void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_thres index_ = common::MakeFixedVecWithMalloc(storage_size, std::uint8_t{0}, n_threads); if (!all_dense_column) { - row_ind_ = common::MakeFixedVecWithMalloc(feature_offsets_[nfeature], std::size_t{0}, n_threads); + row_ind_ = common::MakeFixedVecWithMalloc(feature_offsets_[nfeature], + std::size_t{0}, n_threads); } // store least bin id for each feature diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 27bc03ebe176..9460c9e7f4fe 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -13,6 +13,7 @@ #include // for uint8_t #include #include +#include #include // for enable_if_t, is_same_v, is_signed_v #include "../data/adapter.h" @@ -423,7 +424,7 @@ class ColumnMatrix { int tid = omp_get_thread_num(); size_t begin = block_size * tid; size_t end = std::min(begin + block_size, batch_size); - size_t k = 0; + size_t k = 0; for (size_t rid = begin; rid < end; ++rid) { const auto& line = batch.GetLine(rid); for (size_t i = 0; i < line.Size(); ++i) { From e211ab99b2a8ec077abd22e65bd1f4096c825fc2 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 00:29:14 -0700 Subject: [PATCH 03/35] changing the capture for inner lambdas --- src/common/column_matrix.h | 4 ++-- src/data/gradient_index.h | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 9460c9e7f4fe..e32505f0387f 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -390,7 +390,7 @@ class ColumnMatrix { size_t block_size = DivRoundUp(batch_size, n_threads); #pragma omp parallel num_threads(n_threads) { - exc.Run([&]() { + exc.Run([&, is_valid]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; size_t end = std::min(begin + block_size, batch_size); @@ -420,7 +420,7 @@ class ColumnMatrix { #pragma omp parallel num_threads(n_threads) { - exc.Run([&]() { + exc.Run([&, is_valid, base_rowid, row_index]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; size_t end = std::min(begin + block_size, batch_size); diff --git a/src/data/gradient_index.h b/src/data/gradient_index.h index 26c1496bdbd8..e859208a4c0f 100644 --- a/src/data/gradient_index.h +++ b/src/data/gradient_index.h @@ -233,7 +233,7 @@ class GHistIndexMatrix { void PushAdapterBatchColumns(Context const* ctx, Batch const& batch, float missing, size_t rbegin); - void ResizeIndex(const size_t n_index, const bool isDense, int n_threads); + void ResizeIndex(const size_t n_index, const bool isDense, int n_threads=1); void GetFeatureCounts(size_t* counts) const { auto nfeature = cut.Ptrs().size() - 1; From 92215736a4292e383d9af3b28c7c22d1b94d27bc Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 00:52:02 -0700 Subject: [PATCH 04/35] fix --- src/data/gradient_index.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/data/gradient_index.h b/src/data/gradient_index.h index e859208a4c0f..fba59f177165 100644 --- a/src/data/gradient_index.h +++ b/src/data/gradient_index.h @@ -233,7 +233,7 @@ class GHistIndexMatrix { void PushAdapterBatchColumns(Context const* ctx, Batch const& batch, float missing, size_t rbegin); - void ResizeIndex(const size_t n_index, const bool isDense, int n_threads=1); + void ResizeIndex(const size_t n_index, const bool isDense, int n_threads = 1); void GetFeatureCounts(size_t* counts) const { auto nfeature = cut.Ptrs().size() - 1; From 0a793e32e83494cda1e3985438b4e87e4968188b Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 00:55:46 -0700 Subject: [PATCH 05/35] set default --- src/common/column_matrix.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index e32505f0387f..30b67ba50f2f 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -238,7 +238,7 @@ class ColumnMatrix { } ColumnMatrix() = default; - ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads) { + ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads=1) { this->InitStorage(gmat, sparse_threshold, n_threads); } From e249a3bed441349389014e71b1e309a36c13f43a Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 05:38:53 -0700 Subject: [PATCH 06/35] linting --- src/common/column_matrix.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 30b67ba50f2f..abeeff06c3c1 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -238,7 +238,7 @@ class ColumnMatrix { } ColumnMatrix() = default; - ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads=1) { + ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads = 1) { this->InitStorage(gmat, sparse_threshold, n_threads); } From 1be6f5d80606cb49bbcf0c67e86937fb3df9f95d Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 08:32:41 -0700 Subject: [PATCH 07/35] fix test --- src/common/column_matrix.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index abeeff06c3c1..93f48a28e9f7 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -420,6 +420,7 @@ class ColumnMatrix { #pragma omp parallel num_threads(n_threads) { + std::vector nnz_offsets(n_features, 0); exc.Run([&, is_valid, base_rowid, row_index]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; @@ -432,10 +433,10 @@ class ColumnMatrix { if (is_valid(coo)) { auto fid = coo.column_idx; const uint32_t bin_id = row_index[k_offsets[tid] + k]; - size_t& nnz = n_elements[tid * n_features + fid]; + size_t nnz = n_elements[tid * n_features + fid] + nnz_offsets[fid]; SetBinSparse(bin_id, rid + base_rowid, fid, local_index, nnz); ++k; - ++nnz; + nnz_offsets[fid]++; } } } From 55a89d7645911f311c1c0472a2e9e84390772726 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 08:42:24 -0700 Subject: [PATCH 08/35] submodule fix --- dmlc-core | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmlc-core b/dmlc-core index 133418575498..8986b0598df7 160000 --- a/dmlc-core +++ b/dmlc-core @@ -1 +1 @@ -Subproject commit 13341857549852a9a86b1894b5ba84c6276ab381 +Subproject commit 8986b0598df709117570984571476c3614f55724 From 8a15c70eefb19ac038a7b6269264ad5a18247a58 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Tue, 8 Apr 2025 08:49:37 -0700 Subject: [PATCH 09/35] fix for i386 --- src/common/ref_resource_view.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/common/ref_resource_view.h b/src/common/ref_resource_view.h index 6d8a5d91c5ba..10fa9b4f13e3 100644 --- a/src/common/ref_resource_view.h +++ b/src/common/ref_resource_view.h @@ -171,8 +171,6 @@ template template [[nodiscard]] RefResourceView MakeFixedVecWithMalloc(std::size_t n_elements, T const& init, int n_threads) { - if (n_elements < n_threads) return MakeFixedVecWithMalloc(n_elements, init); - auto resource = std::make_shared(n_elements * sizeof(T)); auto ref = RefResourceView{resource->DataAs(), n_elements, resource}; From 085627fb25ee52f29841e09114bf730ac2e7e5ee Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Wed, 9 Apr 2025 04:33:23 -0700 Subject: [PATCH 10/35] proteckt thread-unsafe code --- src/common/column_matrix.h | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 93f48a28e9f7..2441fcfbecfd 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -223,7 +223,10 @@ class ColumnMatrix { // 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); + #pragma omp critical + { + missing_.SetValid(feature_offsets_[fid] + rid); + } } else { ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[nnz] = bin_id - index_base_[fid]; From b1e714fa53c60f67031d4e5ea9d409ab122439d9 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Thu, 10 Apr 2025 01:00:01 -0700 Subject: [PATCH 11/35] fix for multi-batch --- src/common/column_matrix.h | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 2441fcfbecfd..cf848ab321f4 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -403,7 +403,9 @@ class ColumnMatrix { auto coo = line.GetElement(i); if (is_valid(coo)) { auto fid = coo.column_idx; - n_elements[(tid + 1) * n_features + fid] += 1; + if ((type_[fid] != kDenseColumn)) { + n_elements[(tid + 1) * n_features + fid] += 1; + } k_offsets[tid + 1] += 1; } } @@ -413,11 +415,16 @@ class ColumnMatrix { exc.Rethrow(); ParallelFor(n_features, n_threads, [&](auto fid) { + for (size_t fid = 0; fid < n_features; ++fid) { + n_elements[fid] += num_nonzeros_[fid]; for (int tid = 0; tid < n_threads; ++tid) { n_elements[(tid + 1) * n_features + fid] += n_elements[tid * n_features + fid]; } - num_nonzeros_[fid] = n_elements[n_threads * n_features + fid]; + if (type_[fid] != kDenseColumn) { + num_nonzeros_[fid] = n_elements[n_threads * n_features + fid]; + } + } }); std::partial_sum(k_offsets.cbegin(), k_offsets.cend(), k_offsets.begin()); @@ -439,7 +446,7 @@ class ColumnMatrix { size_t nnz = n_elements[tid * n_features + fid] + nnz_offsets[fid]; SetBinSparse(bin_id, rid + base_rowid, fid, local_index, nnz); ++k; - nnz_offsets[fid]++; + nnz_offsets[fid] += (type_[fid] != kDenseColumn); } } } From 70fd6bce733f6906d441c7b3068b15bafe0b15c6 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Thu, 10 Apr 2025 02:39:09 -0700 Subject: [PATCH 12/35] fix compilation error --- src/common/column_matrix.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index cf848ab321f4..850dccf6764a 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -415,7 +415,6 @@ class ColumnMatrix { exc.Rethrow(); ParallelFor(n_features, n_threads, [&](auto fid) { - for (size_t fid = 0; fid < n_features; ++fid) { n_elements[fid] += num_nonzeros_[fid]; for (int tid = 0; tid < n_threads; ++tid) { n_elements[(tid + 1) * n_features + fid] += @@ -424,7 +423,6 @@ class ColumnMatrix { if (type_[fid] != kDenseColumn) { num_nonzeros_[fid] = n_elements[n_threads * n_features + fid]; } - } }); std::partial_sum(k_offsets.cbegin(), k_offsets.cend(), k_offsets.begin()); From edef9e717adc224594172c26c452f2725e66f0c7 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Thu, 10 Apr 2025 07:27:32 -0700 Subject: [PATCH 13/35] remove critical section; avoid using of bit filds --- src/common/column_matrix.h | 60 ++++++++++---------------- tests/cpp/common/test_column_matrix.cc | 2 +- 2 files changed, 24 insertions(+), 38 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 850dccf6764a..109116361268 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -114,18 +114,18 @@ class DenseColumnIter : public Column { private: using Base = Column; /* flags for missing values in dense columns */ - LBitField32 missing_flags_; + Span missing_flags_; size_t feature_offset_; public: explicit DenseColumnIter(common::Span index, bst_bin_t index_base, - LBitField32 missing_flags, size_t feature_offset) + Span 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; [[nodiscard]] bool IsMissing(size_t ridx) const { - return missing_flags_.Check(feature_offset_ + ridx); + return missing_flags_[feature_offset_ + ridx]; } bst_bin_t operator[](size_t ridx) const { @@ -149,48 +149,43 @@ class ColumnMatrix { * @brief A bit set for indicating whether an element in a dense column is missing. */ struct MissingIndicator { - using BitFieldT = LBitField32; - using T = typename BitFieldT::value_type; - - BitFieldT missing; - RefResourceView storage; - static_assert(std::is_same_v); + Span missing; + RefResourceView storage; template - [[nodiscard]] std::enable_if_t, U> static InitValue(bool init) { + [[nodiscard]] std::enable_if_t, U> static InitValue(uint8_t init) { return init ? ~U{0} : U{0}; } MissingIndicator() = default; - /** - * @param n_elements Size of the bit set - * @param init Initialize the indicator to true or false. - */ + // /** + // * @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, InitValue(init)); - this->InitView(); + // auto m_size = missing.ComputeStorageSize(n_elements); + storage = common::MakeFixedVecWithMalloc(n_elements, uint8_t(init)); + // 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); } + void SetValid(size_t i) { missing[i] = 0; } /** @brief assign the storage to the view. */ void InitView() { - missing = LBitField32{Span{storage.data(), static_cast(storage.size())}}; + missing = Span{storage.data(), static_cast(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()) { + CHECK_GE(n_elements, storage.size()); + if (n_elements == storage.size()) { return; } // grow the storage auto resource = std::dynamic_pointer_cast(storage.Resource()); CHECK(resource); - resource->Resize(m_size * sizeof(T), InitValue(init)); - storage = RefResourceView{resource->DataAs(), m_size, resource}; + resource->Resize(n_elements * sizeof(uint8_t), InitValue(init)); + storage = RefResourceView{resource->DataAs(), n_elements, resource}; this->InitView(); } @@ -203,9 +198,6 @@ class ColumnMatrix { if (type_[fid] == kDenseColumn) { ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[rid] = bin_id - index_base_[fid]; - // 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]]; @@ -220,13 +212,7 @@ class ColumnMatrix { if (type_[fid] == kDenseColumn) { ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[rid] = bin_id - index_base_[fid]; - // not thread-safe with bit field. - // FIXME(jiamingy): We can directly assign kMissingId to the index to avoid missing - // flags. - #pragma omp critical - { - missing_.SetValid(feature_offsets_[fid] + rid); - } + missing_.SetValid(feature_offsets_[fid] + rid); } else { ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[nnz] = bin_id - index_base_[fid]; @@ -420,9 +406,7 @@ class ColumnMatrix { n_elements[(tid + 1) * n_features + fid] += n_elements[tid * n_features + fid]; } - if (type_[fid] != kDenseColumn) { - num_nonzeros_[fid] = n_elements[n_threads * n_features + fid]; - } + num_nonzeros_[fid] = n_elements[n_threads * n_features + fid]; }); std::partial_sum(k_offsets.cbegin(), k_offsets.cend(), k_offsets.begin()); @@ -496,6 +480,8 @@ class ColumnMatrix { RefResourceView feature_offsets_; /** @brief The number of nnz of each column. */ RefResourceView num_nonzeros_; + /** @brief The number of nnz of each column. */ + RefResourceView missing_flag_; // index_base_[fid]: least bin id for feature fid std::uint32_t const* index_base_; diff --git a/tests/cpp/common/test_column_matrix.cc b/tests/cpp/common/test_column_matrix.cc index 38f64f29798e..645b8c7f058f 100644 --- a/tests/cpp/common/test_column_matrix.cc +++ b/tests/cpp/common/test_column_matrix.cc @@ -136,7 +136,7 @@ TEST(ColumnMatrix, GrowMissing) { auto const& column_matrix = page.Transpose(); auto const& missing = column_matrix.Missing(); auto n = NumpyArrayIterForTest::Rows() * NumpyArrayIterForTest::Cols(); - auto expected = std::remove_reference_t::BitFieldT::ComputeStorageSize(n); + auto expected = n; auto got = missing.storage.size(); ASSERT_EQ(expected, got); DispatchBinType(column_matrix.GetTypeSize(), [&](auto dtype) { From 606c53751fcccdf6a0ca430e447af5250218f353 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Thu, 10 Apr 2025 08:00:57 -0700 Subject: [PATCH 14/35] tildy --- src/common/column_matrix.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 109116361268..a787ac52de5c 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -164,7 +164,7 @@ class ColumnMatrix { // */ MissingIndicator(std::size_t n_elements, bool init) { // auto m_size = missing.ComputeStorageSize(n_elements); - storage = common::MakeFixedVecWithMalloc(n_elements, uint8_t(init)); + storage = common::MakeFixedVecWithMalloc(n_elements, static_cast(init)); // this->InitView(); } /** @brief Set the i^th element to be a valid element (instead of missing). */ From 6f885b0847ced6477678fc78afb584c6f4ffa7c1 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Fri, 11 Apr 2025 00:42:55 -0700 Subject: [PATCH 15/35] return deleted code --- src/common/column_matrix.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index a787ac52de5c..ba4c6d29f463 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -163,9 +163,8 @@ class ColumnMatrix { // * @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(n_elements, static_cast(init)); - // this->InitView(); + this->InitView(); } /** @brief Set the i^th element to be a valid element (instead of missing). */ void SetValid(size_t i) { missing[i] = 0; } From c0dbd7e6388aee64e6e73506756bae5ccd2ab489 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Fri, 11 Apr 2025 01:38:07 -0700 Subject: [PATCH 16/35] remove unactual code --- src/common/column_matrix.h | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index ba4c6d29f463..5225985b343c 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -158,10 +158,10 @@ class ColumnMatrix { } MissingIndicator() = default; - // /** - // * @param n_elements Size of the bit set - // * @param init Initialize the indicator to true or false. - // */ + /** + * @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) { storage = common::MakeFixedVecWithMalloc(n_elements, static_cast(init)); this->InitView(); @@ -479,8 +479,6 @@ class ColumnMatrix { RefResourceView feature_offsets_; /** @brief The number of nnz of each column. */ RefResourceView num_nonzeros_; - /** @brief The number of nnz of each column. */ - RefResourceView missing_flag_; // index_base_[fid]: least bin id for feature fid std::uint32_t const* index_base_; From 1cb3693b3522bd5b565a36d3df21526b90e4c86c Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Mon, 5 May 2025 07:11:04 -0700 Subject: [PATCH 17/35] address comments --- src/common/column_matrix.h | 13 ++++++++----- src/data/gradient_index.cu | 2 +- src/data/gradient_index.h | 2 +- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 5225985b343c..c65c63dadde2 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -194,12 +194,11 @@ class ColumnMatrix { template void SetBinSparse(BinT bin_id, RIdx rid, bst_feature_t fid, ColumnBinT* local_index) { + ColumnBinT* begin = &local_index[feature_offsets_[fid]]; if (type_[fid] == kDenseColumn) { - ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[rid] = bin_id - index_base_[fid]; missing_.SetValid(feature_offsets_[fid] + rid); } else { - ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[num_nonzeros_[fid]] = bin_id - index_base_[fid]; row_ind_[feature_offsets_[fid] + num_nonzeros_[fid]] = rid; ++num_nonzeros_[fid]; @@ -208,12 +207,11 @@ class ColumnMatrix { template void SetBinSparse(BinT bin_id, RIdx rid, bst_feature_t fid, ColumnBinT* local_index, size_t nnz) { + ColumnBinT* begin = &local_index[feature_offsets_[fid]]; if (type_[fid] == kDenseColumn) { - ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[rid] = bin_id - index_base_[fid]; missing_.SetValid(feature_offsets_[fid] + rid); } else { - ColumnBinT* begin = &local_index[feature_offsets_[fid]]; begin[nnz] = bin_id - index_base_[fid]; row_ind_[feature_offsets_[fid] + nnz] = rid; } @@ -226,7 +224,7 @@ class ColumnMatrix { } ColumnMatrix() = default; - ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads = 1) { + ColumnMatrix(GHistIndexMatrix const& gmat, double sparse_threshold, int n_threads) { this->InitStorage(gmat, sparse_threshold, n_threads); } @@ -372,10 +370,13 @@ class ColumnMatrix { ColumnBinT* local_index = reinterpret_cast(index_.data()); size_t const batch_size = batch.Size(); + // Parallel sparse batch processing dmlc::OMPException exc; std::vector n_elements((n_threads + 1) * n_features, 0); std::vector k_offsets(n_threads + 1, 0); size_t block_size = DivRoundUp(batch_size, n_threads); + + // Parallel row processing for thread-local counting. #pragma omp parallel num_threads(n_threads) { exc.Run([&, is_valid]() { @@ -399,6 +400,7 @@ class ColumnMatrix { } exc.Rethrow(); + // Parallel feature processing to aggregate counts & calculate offsets. ParallelFor(n_features, n_threads, [&](auto fid) { n_elements[fid] += num_nonzeros_[fid]; for (int tid = 0; tid < n_threads; ++tid) { @@ -409,6 +411,7 @@ class ColumnMatrix { }); std::partial_sum(k_offsets.cbegin(), k_offsets.cend(), k_offsets.begin()); + // Parallel row processing to place data using offsets into sparse structure. #pragma omp parallel num_threads(n_threads) { std::vector nnz_offsets(n_features, 0); diff --git a/src/data/gradient_index.cu b/src/data/gradient_index.cu index 5e15ff5f0fa2..0d11432e545a 100644 --- a/src/data/gradient_index.cu +++ b/src/data/gradient_index.cu @@ -92,7 +92,7 @@ GHistIndexMatrix::GHistIndexMatrix(Context const* ctx, MetaInfo const& info, this->cut.Values(); this->cut.MinValues(); - this->ResizeIndex(info.num_nonzero_, page->IsDense()); + this->ResizeIndex(info.num_nonzero_, page->IsDense(), ctx->Threads()); if (page->IsDense()) { this->index.SetBinOffset(page->Cuts().Ptrs()); } diff --git a/src/data/gradient_index.h b/src/data/gradient_index.h index fba59f177165..26c1496bdbd8 100644 --- a/src/data/gradient_index.h +++ b/src/data/gradient_index.h @@ -233,7 +233,7 @@ class GHistIndexMatrix { void PushAdapterBatchColumns(Context const* ctx, Batch const& batch, float missing, size_t rbegin); - void ResizeIndex(const size_t n_index, const bool isDense, int n_threads = 1); + void ResizeIndex(const size_t n_index, const bool isDense, int n_threads); void GetFeatureCounts(size_t* counts) const { auto nfeature = cut.Ptrs().size() - 1; From 560a67a2e8ed1534a34b873b93215cf4451ad442 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Mon, 5 May 2025 07:21:24 -0700 Subject: [PATCH 18/35] fix calling ColumnMatrix constructor --- src/data/gradient_index.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/data/gradient_index.cu b/src/data/gradient_index.cu index 0d11432e545a..16e30b40a1e1 100644 --- a/src/data/gradient_index.cu +++ b/src/data/gradient_index.cu @@ -123,7 +123,7 @@ GHistIndexMatrix::GHistIndexMatrix(Context const* ctx, MetaInfo const& info, CHECK(this->cut.cut_values_.HostCanRead()); CHECK(this->cut.min_vals_.HostCanRead()); - this->columns_ = std::make_unique(*this, p.sparse_thresh); + this->columns_ = std::make_unique(*this, p.sparse_thresh, ctx->Threads()); this->columns_->InitFromGHist(ctx, *this); } } // namespace xgboost From 0ac338e2547d48c2a65ef49762f4519f1cc93314 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Mon, 19 May 2025 00:51:55 -0700 Subject: [PATCH 19/35] switch back to bitfield --- src/common/column_matrix.cc | 6 +- src/common/column_matrix.h | 124 +++++++++++++++++++------ tests/cpp/common/test_column_matrix.cc | 2 +- 3 files changed, 103 insertions(+), 29 deletions(-) diff --git a/src/common/column_matrix.cc b/src/common/column_matrix.cc index 21271cbda28a..7a7cfe107d79 100644 --- a/src/common/column_matrix.cc +++ b/src/common/column_matrix.cc @@ -74,7 +74,7 @@ void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_thres any_missing_ = !gmat.IsDense(); - missing_ = MissingIndicator{0, false}; + missing_ = MissingIndicator{feature_offsets_, type_, any_missing_}; } // IO procedures for external memory. @@ -95,6 +95,9 @@ bool ColumnMatrix::Read(AlignedResourceReadStream* fi, uint32_t const* index_bas if (!common::ReadVec(fi, &missing_.storage)) { return false; } + if (!common::ReadVec(fi, &missing_.feature_offsets_expand)) { + return false; + } missing_.InitView(); index_base_ = index_base; @@ -115,6 +118,7 @@ std::size_t ColumnMatrix::Write(AlignedFileWriteStream* fo) const { bytes += common::WriteVec(fo, row_ind_); bytes += common::WriteVec(fo, feature_offsets_); bytes += common::WriteVec(fo, missing_.storage); + bytes += common::WriteVec(fo, missing_.feature_offsets_expand); bytes += fo->Write(bins_type_size_); bytes += fo->Write(any_missing_); diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index c65c63dadde2..e0e6654f2b50 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -114,18 +114,18 @@ class DenseColumnIter : public Column { private: using Base = Column; /* flags for missing values in dense columns */ - Span missing_flags_; + LBitField32 missing_flags_; size_t feature_offset_; public: explicit DenseColumnIter(common::Span index, bst_bin_t index_base, - Span 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; [[nodiscard]] bool IsMissing(size_t ridx) const { - return missing_flags_[feature_offset_ + ridx]; + return missing_flags_.Check(feature_offset_ + ridx); } bst_bin_t operator[](size_t ridx) const { @@ -145,46 +145,91 @@ class DenseColumnIter : public Column { * 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. - */ +/** + * @brief A bit set for indicating whether an element in a dense column is missing. + * Access is carefully managed to ensure thread safety during parallel operations. + */ struct MissingIndicator { - Span missing; - RefResourceView storage; + using BitFieldT = LBitField32; + using T = typename BitFieldT::value_type; + + BitFieldT missing; + RefResourceView storage; + RefResourceView feature_offsets_expand; + static_assert(std::is_same_v); template - [[nodiscard]] std::enable_if_t, U> static InitValue(uint8_t init) { + [[nodiscard]] std::enable_if_t, U> static InitValue(bool init) { return init ? ~U{0} : U{0}; } + /** + * @param feature_offsets Offest of the first element for each feature + * @param type Type of each column (Dense or Sparce). + */ + void InitOffsetsExpand(const RefResourceView& feature_offsets, + const RefResourceView& type) { + if (feature_offsets_expand.size() != feature_offsets.size()) { + feature_offsets_expand = common::MakeFixedVecWithMalloc(feature_offsets.size(), std::size_t{0}); + } + + /* + * For missing indicator feature offsets are alligned to be a factor of + * BitFieldT::kValueSize (4 bytes). + * This is critical requariment for thread-safe access to bitfield. + * Each word processed by one thread. + */ + for (size_t fid = 1; fid < feature_offsets.size(); ++fid) { + if (type[fid - 1] == ColumnType::kDenseColumn) { + size_t n_rows = feature_offsets[fid] - feature_offsets[fid - 1]; + size_t n_rows_expand = DivRoundUp(n_rows, BitFieldT::kValueSize) * BitFieldT::kValueSize; + feature_offsets_expand[fid] = feature_offsets_expand[fid - 1] + n_rows_expand; + } else { + feature_offsets_expand[fid] = feature_offsets_expand[fid - 1]; + } + } + } + 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) { - storage = common::MakeFixedVecWithMalloc(n_elements, static_cast(init)); + MissingIndicator(const RefResourceView& feature_offsets, + const RefResourceView& type, bool init) { + InitOffsetsExpand(feature_offsets, type); + size_t n_elements = feature_offsets_expand.back(); + auto m_size = missing.ComputeStorageSize(n_elements); + storage = common::MakeFixedVecWithMalloc(m_size, InitValue(init)); this->InitView(); } - /** @brief Set the i^th element to be a valid element (instead of missing). */ - void SetValid(size_t i) { missing[i] = 0; } + /** @brief Set the i^th element corresponding to feature fid + * to be a valid element (instead of missing). */ + void SetValid(typename LBitField32::index_type i, std::size_t fid) { + missing.Clear(feature_offsets_expand[fid] + i); + } /** @brief assign the storage to the view. */ void InitView() { - missing = Span{storage.data(), static_cast(storage.size())}; + missing = LBitField32{Span{storage.data(), static_cast(storage.size())}}; } - void GrowTo(std::size_t n_elements, bool init) { + void GrowTo(const RefResourceView& feature_offsets, + const RefResourceView& type, bool init) { + InitOffsetsExpand(feature_offsets, type); + size_t n_elements = feature_offsets_expand.back(); + CHECK(storage.Resource()->Type() == ResourceHandler::kMalloc) << "[Internal Error]: Cannot grow the vector when external memory is used."; - CHECK_GE(n_elements, storage.size()); - if (n_elements == storage.size()) { + auto m_size = missing.ComputeStorageSize(n_elements); + CHECK_GE(m_size, storage.size()); + if (m_size == storage.size()) { return; } // grow the storage auto resource = std::dynamic_pointer_cast(storage.Resource()); CHECK(resource); - resource->Resize(n_elements * sizeof(uint8_t), InitValue(init)); - storage = RefResourceView{resource->DataAs(), n_elements, resource}; + resource->Resize(m_size * sizeof(T), InitValue(init)); + storage = RefResourceView{resource->DataAs(), m_size, resource}; this->InitView(); } @@ -197,7 +242,7 @@ class ColumnMatrix { ColumnBinT* begin = &local_index[feature_offsets_[fid]]; if (type_[fid] == kDenseColumn) { begin[rid] = bin_id - index_base_[fid]; - missing_.SetValid(feature_offsets_[fid] + rid); + missing_.SetValid(rid, fid); } else { begin[num_nonzeros_[fid]] = bin_id - index_base_[fid]; row_ind_[feature_offsets_[fid] + num_nonzeros_[fid]] = rid; @@ -210,7 +255,7 @@ class ColumnMatrix { ColumnBinT* begin = &local_index[feature_offsets_[fid]]; if (type_[fid] == kDenseColumn) { begin[rid] = bin_id - index_base_[fid]; - missing_.SetValid(feature_offsets_[fid] + rid); + missing_.SetValid(rid, fid); } else { begin[nnz] = bin_id - index_base_[fid]; row_ind_[feature_offsets_[fid] + nnz] = rid; @@ -319,8 +364,13 @@ class ColumnMatrix { common::Span bin_index = { reinterpret_cast(&index_[feature_offset * bins_type_size_]), column_size}; + /* + * Pass the pre-calculated starting offset missing_.feature_offsets_expand[fidx] + * in the bitfield for this specific feature (fidx). + */ return DenseColumnIter{ - bin_index, static_cast(index_base_[fidx]), missing_.missing, feature_offset}; + bin_index, static_cast(index_base_[fidx]), missing_.missing, + missing_.feature_offsets_expand[fidx]}; } // all columns are dense column and has no missing value @@ -328,7 +378,7 @@ class ColumnMatrix { template void SetIndexNoMissing(bst_idx_t base_rowid, RowBinIdxT const* row_index, const size_t n_samples, const size_t n_features, int32_t n_threads) { - missing_.GrowTo(feature_offsets_[n_features], false); + missing_.GrowTo(feature_offsets_, type_, false); DispatchBinType(bins_type_size_, [&](auto t) { using ColumnBinT = decltype(t); @@ -355,7 +405,7 @@ class ColumnMatrix { float missing, int n_threads) { auto n_features = gmat.Features(); - missing_.GrowTo(feature_offsets_[n_features], true); + missing_.GrowTo(feature_offsets_, type_, true); auto const* row_index = gmat.index.data() + gmat.row_ptr[base_rowid]; if (num_nonzeros_.empty()) { num_nonzeros_ = common::MakeFixedVecWithMalloc(n_features, std::size_t{0}); @@ -376,13 +426,30 @@ class ColumnMatrix { std::vector k_offsets(n_threads + 1, 0); size_t block_size = DivRoundUp(batch_size, n_threads); + /* + * We use bitfield as a missing indicator. To insure thread safe access to the bitfield + * each underlying word of the bitfiled should be processed by the only thread. + * So we need to allign the row-blocks. + */ + block_size = DivRoundUp(block_size, MissingIndicator::BitFieldT::kValueSize) * + MissingIndicator::BitFieldT::kValueSize; + /* + * If base_rowid > 0 we need to shift the blocks boundaries. + * Otherwise the two threads may operate with the single word of bitfield. + */ + size_t shift = MissingIndicator::BitFieldT::kValueSize - + (base_rowid % MissingIndicator::BitFieldT::kValueSize); + if (shift == MissingIndicator::BitFieldT::kValueSize) shift = 0; + // Parallel row processing for thread-local counting. #pragma omp parallel num_threads(n_threads) { exc.Run([&, is_valid]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; - size_t end = std::min(begin + block_size, batch_size); + size_t end = std::min(begin + shift + block_size, batch_size); + // Apply shift for threads > 0 to maintain word alignment across blocks. + if (tid > 0) begin += shift; for (size_t rid = begin; rid < end; ++rid) { const auto& line = batch.GetLine(rid); for (size_t i = 0; i < line.Size(); ++i) { @@ -418,7 +485,10 @@ class ColumnMatrix { exc.Run([&, is_valid, base_rowid, row_index]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; - size_t end = std::min(begin + block_size, batch_size); + size_t end = std::min(begin + shift + block_size, batch_size); + // Apply shift for threads > 0 to maintain word alignment across blocks. + if (tid > 0) begin += shift; + size_t k = 0; for (size_t rid = begin; rid < end; ++rid) { const auto& line = batch.GetLine(rid); @@ -447,7 +517,7 @@ class ColumnMatrix { void SetIndexMixedColumns(const GHistIndexMatrix& gmat) { auto n_features = gmat.Features(); - missing_ = MissingIndicator{feature_offsets_[n_features], true}; + missing_ = MissingIndicator{feature_offsets_, type_, true}; num_nonzeros_ = common::MakeFixedVecWithMalloc(n_features, std::size_t{0}); DispatchBinType(bins_type_size_, [&](auto t) { diff --git a/tests/cpp/common/test_column_matrix.cc b/tests/cpp/common/test_column_matrix.cc index 645b8c7f058f..38f64f29798e 100644 --- a/tests/cpp/common/test_column_matrix.cc +++ b/tests/cpp/common/test_column_matrix.cc @@ -136,7 +136,7 @@ TEST(ColumnMatrix, GrowMissing) { auto const& column_matrix = page.Transpose(); auto const& missing = column_matrix.Missing(); auto n = NumpyArrayIterForTest::Rows() * NumpyArrayIterForTest::Cols(); - auto expected = n; + auto expected = std::remove_reference_t::BitFieldT::ComputeStorageSize(n); auto got = missing.storage.size(); ASSERT_EQ(expected, got); DispatchBinType(column_matrix.GetTypeSize(), [&](auto dtype) { From 61b387834e24ef603c07626ed040d30d2aa16c0f Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Mon, 19 May 2025 01:03:31 -0700 Subject: [PATCH 20/35] linting --- src/common/column_matrix.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index e0e6654f2b50..1c4fd2914227 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -170,7 +170,8 @@ class ColumnMatrix { void InitOffsetsExpand(const RefResourceView& feature_offsets, const RefResourceView& type) { if (feature_offsets_expand.size() != feature_offsets.size()) { - feature_offsets_expand = common::MakeFixedVecWithMalloc(feature_offsets.size(), std::size_t{0}); + feature_offsets_expand = common::MakeFixedVecWithMalloc(feature_offsets.size(), + std::size_t{0}); } /* @@ -213,7 +214,7 @@ class ColumnMatrix { missing = LBitField32{Span{storage.data(), static_cast(storage.size())}}; } - void GrowTo(const RefResourceView& feature_offsets, + void GrowTo(const RefResourceView& feature_offsets, const RefResourceView& type, bool init) { InitOffsetsExpand(feature_offsets, type); size_t n_elements = feature_offsets_expand.back(); From 98ef5410e81c367e34848a840a2e0982cda6f9bb Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Mon, 2 Jun 2025 08:37:48 +0200 Subject: [PATCH 21/35] Update src/common/column_matrix.h Co-authored-by: Jiaming Yuan --- src/common/column_matrix.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 1c4fd2914227..e15e13d8cd81 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -428,7 +428,7 @@ class ColumnMatrix { size_t block_size = DivRoundUp(batch_size, n_threads); /* - * We use bitfield as a missing indicator. To insure thread safe access to the bitfield + * We use bitfield as a missing indicator. To ensure thread safe access to the bitfield * each underlying word of the bitfiled should be processed by the only thread. * So we need to allign the row-blocks. */ From 2b090e67eb56b33b5a0f3706c2826aeab639113b Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Mon, 2 Jun 2025 08:38:00 +0200 Subject: [PATCH 22/35] Update src/common/column_matrix.h Co-authored-by: Jiaming Yuan --- src/common/column_matrix.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index e15e13d8cd81..7820b4538ec4 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -165,7 +165,7 @@ class ColumnMatrix { /** * @param feature_offsets Offest of the first element for each feature - * @param type Type of each column (Dense or Sparce). + * @param type Type of each column (Dense or Sparse). */ void InitOffsetsExpand(const RefResourceView& feature_offsets, const RefResourceView& type) { From 8cdd7db87e42dac9dd567785787df656d0dd96ee Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Tue, 1 Jul 2025 18:04:09 +0800 Subject: [PATCH 23/35] Cleanup, typos. --- src/common/column_matrix.cc | 11 +++++------ src/common/ref_resource_view.h | 16 ++++++++-------- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/src/common/column_matrix.cc b/src/common/column_matrix.cc index 7a7cfe107d79..e686356983f5 100644 --- a/src/common/column_matrix.cc +++ b/src/common/column_matrix.cc @@ -1,20 +1,19 @@ /** - * Copyright 2017-2023, XGBoost Contributors + * Copyright 2017-2025, XGBoost Contributors * \brief Utility for fast column-wise access */ #include "column_matrix.h" -#include // for transform #include // for size_t #include // for uint64_t, uint8_t #include // for numeric_limits #include // for remove_reference_t #include // 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 +#include "../common/ref_resource_view.h" // for MakeFixedVecWithMalloc +#include "../data/gradient_index.h" // for GHistIndexMatrix +#include "io.h" // for AlignedResourceReadStream, AlignedFileWriteStream +#include "xgboost/base.h" // for bst_feaature_t namespace xgboost::common { void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold, diff --git a/src/common/ref_resource_view.h b/src/common/ref_resource_view.h index 305a51a850b8..b7f04843e3b0 100644 --- a/src/common/ref_resource_view.h +++ b/src/common/ref_resource_view.h @@ -11,7 +11,8 @@ #include // for is_reference_v, remove_reference_t, is_same_v #include // for swap, move -#include "io.h" // for ResourceHandler, AlignedResourceReadStream, MallocResource +#include "io.h" // for ResourceHandler, AlignedResourceReadStream, MallocResource +#include "threading_utils.h" // for ParallelFor #include "xgboost/logging.h" #include "xgboost/span.h" // for Span @@ -168,23 +169,22 @@ template /** * @brief Make a fixed size `RefResourceView` with malloc resource. - * Use n_threads to initilise the storage + * + * Use n_threads to initialise the storage */ template [[nodiscard]] RefResourceView MakeFixedVecWithMalloc(std::size_t n_elements, T const& init, - int n_threads) { + std::int32_t n_threads) { auto resource = std::make_shared(n_elements * sizeof(T)); auto ref = RefResourceView{resource->DataAs(), n_elements, resource}; - size_t block_size = n_elements / n_threads + (n_elements % n_threads > 0); - #pragma omp parallel num_threads(n_threads) - { - int tid = omp_get_thread_num(); + std::size_t block_size = n_elements / n_threads + (n_elements % n_threads > 0); + ParallelFor(n_threads, n_threads, [&](auto tid) { auto begin = tid * block_size; auto end = std::min((tid + 1) * block_size, n_elements); auto size = end > begin ? end - begin : 0; std::fill_n(ref.data() + begin, size, init); - } + }); return ref; } From 15aa65be57f05830a380752089bab07b313c3b84 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Tue, 1 Jul 2025 18:38:56 +0800 Subject: [PATCH 24/35] rename. --- src/common/column_matrix.cc | 4 ++-- src/common/column_matrix.h | 41 ++++++++++++++++++++----------------- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/src/common/column_matrix.cc b/src/common/column_matrix.cc index e686356983f5..90a6f6da0225 100644 --- a/src/common/column_matrix.cc +++ b/src/common/column_matrix.cc @@ -94,7 +94,7 @@ bool ColumnMatrix::Read(AlignedResourceReadStream* fi, uint32_t const* index_bas if (!common::ReadVec(fi, &missing_.storage)) { return false; } - if (!common::ReadVec(fi, &missing_.feature_offsets_expand)) { + if (!common::ReadVec(fi, &missing_.feature_offsets_padded)) { return false; } missing_.InitView(); @@ -117,7 +117,7 @@ std::size_t ColumnMatrix::Write(AlignedFileWriteStream* fo) const { bytes += common::WriteVec(fo, row_ind_); bytes += common::WriteVec(fo, feature_offsets_); bytes += common::WriteVec(fo, missing_.storage); - bytes += common::WriteVec(fo, missing_.feature_offsets_expand); + bytes += common::WriteVec(fo, missing_.feature_offsets_padded); bytes += fo->Write(bins_type_size_); bytes += fo->Write(any_missing_); diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 7820b4538ec4..0834978f4b4d 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -155,7 +155,8 @@ class ColumnMatrix { BitFieldT missing; RefResourceView storage; - RefResourceView feature_offsets_expand; + // Feature offset padded to allow concurrent access. + RefResourceView feature_offsets_padded; static_assert(std::is_same_v); template @@ -167,26 +168,28 @@ class ColumnMatrix { * @param feature_offsets Offest of the first element for each feature * @param type Type of each column (Dense or Sparse). */ - void InitOffsetsExpand(const RefResourceView& feature_offsets, + void InitOffsetsPadded(const RefResourceView& feature_offsets, const RefResourceView& type) { - if (feature_offsets_expand.size() != feature_offsets.size()) { - feature_offsets_expand = common::MakeFixedVecWithMalloc(feature_offsets.size(), + if (feature_offsets_padded.size() != feature_offsets.size()) { + CHECK(feature_offsets_padded.empty()); + feature_offsets_padded = common::MakeFixedVecWithMalloc(feature_offsets.size(), std::size_t{0}); } /* - * For missing indicator feature offsets are alligned to be a factor of + * For missing indicator feature offsets are alligned to be a factor of * BitFieldT::kValueSize (4 bytes). * This is critical requariment for thread-safe access to bitfield. * Each word processed by one thread. - */ - for (size_t fid = 1; fid < feature_offsets.size(); ++fid) { + */ + for (std::size_t fid = 1; fid < feature_offsets.size(); ++fid) { if (type[fid - 1] == ColumnType::kDenseColumn) { - size_t n_rows = feature_offsets[fid] - feature_offsets[fid - 1]; - size_t n_rows_expand = DivRoundUp(n_rows, BitFieldT::kValueSize) * BitFieldT::kValueSize; - feature_offsets_expand[fid] = feature_offsets_expand[fid - 1] + n_rows_expand; + std::size_t n_rows = feature_offsets[fid] - feature_offsets[fid - 1]; + std::size_t n_rows_padded = + DivRoundUp(n_rows, BitFieldT::kValueSize) * BitFieldT::kValueSize; + feature_offsets_padded[fid] = feature_offsets_padded[fid - 1] + n_rows_padded; } else { - feature_offsets_expand[fid] = feature_offsets_expand[fid - 1]; + feature_offsets_padded[fid] = feature_offsets_padded[fid - 1]; } } } @@ -198,16 +201,16 @@ class ColumnMatrix { */ MissingIndicator(const RefResourceView& feature_offsets, const RefResourceView& type, bool init) { - InitOffsetsExpand(feature_offsets, type); - size_t n_elements = feature_offsets_expand.back(); + this->InitOffsetsPadded(feature_offsets, type); + size_t n_elements = feature_offsets_padded.back(); auto m_size = missing.ComputeStorageSize(n_elements); storage = common::MakeFixedVecWithMalloc(m_size, InitValue(init)); this->InitView(); } - /** @brief Set the i^th element corresponding to feature fid + /** @brief Set the i^th element corresponding to feature fid * to be a valid element (instead of missing). */ void SetValid(typename LBitField32::index_type i, std::size_t fid) { - missing.Clear(feature_offsets_expand[fid] + i); + missing.Clear(feature_offsets_padded[fid] + i); } /** @brief assign the storage to the view. */ void InitView() { @@ -216,8 +219,8 @@ class ColumnMatrix { void GrowTo(const RefResourceView& feature_offsets, const RefResourceView& type, bool init) { - InitOffsetsExpand(feature_offsets, type); - size_t n_elements = feature_offsets_expand.back(); + this->InitOffsetsPadded(feature_offsets, type); + size_t n_elements = feature_offsets_padded.back(); CHECK(storage.Resource()->Type() == ResourceHandler::kMalloc) << "[Internal Error]: Cannot grow the vector when external memory is used."; @@ -366,12 +369,12 @@ class ColumnMatrix { reinterpret_cast(&index_[feature_offset * bins_type_size_]), column_size}; /* - * Pass the pre-calculated starting offset missing_.feature_offsets_expand[fidx] + * Pass the pre-calculated starting offset missing_.feature_offsets_expand[fidx] * in the bitfield for this specific feature (fidx). */ return DenseColumnIter{ bin_index, static_cast(index_base_[fidx]), missing_.missing, - missing_.feature_offsets_expand[fidx]}; + missing_.feature_offsets_padded[fidx]}; } // all columns are dense column and has no missing value From 58920a05fedc87ca440d5751c70d96d13677136b Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Tue, 1 Jul 2025 18:46:34 +0800 Subject: [PATCH 25/35] typos. --- src/common/column_matrix.h | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 0834978f4b4d..2a06eba3d176 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -405,7 +405,7 @@ class ColumnMatrix { * \brief Set column index for both dense and sparse columns */ template - void SetIndexMixedColumns(size_t base_rowid, Batch const& batch, const GHistIndexMatrix& gmat, + void SetIndexMixedColumns(bst_idx_t base_rowid, Batch const& batch, const GHistIndexMatrix& gmat, float missing, int n_threads) { auto n_features = gmat.Features(); @@ -432,18 +432,20 @@ class ColumnMatrix { /* * We use bitfield as a missing indicator. To ensure thread safe access to the bitfield - * each underlying word of the bitfiled should be processed by the only thread. - * So we need to allign the row-blocks. + * each underlying word of the bitfiled should be processed by a single thread. + * So we need to align the row-blocks. */ block_size = DivRoundUp(block_size, MissingIndicator::BitFieldT::kValueSize) * - MissingIndicator::BitFieldT::kValueSize; + MissingIndicator::BitFieldT::kValueSize; /* * If base_rowid > 0 we need to shift the blocks boundaries. * Otherwise the two threads may operate with the single word of bitfield. */ size_t shift = MissingIndicator::BitFieldT::kValueSize - (base_rowid % MissingIndicator::BitFieldT::kValueSize); - if (shift == MissingIndicator::BitFieldT::kValueSize) shift = 0; + if (shift == MissingIndicator::BitFieldT::kValueSize) { + shift = 0; + } // Parallel row processing for thread-local counting. #pragma omp parallel num_threads(n_threads) @@ -453,7 +455,9 @@ class ColumnMatrix { size_t begin = block_size * tid; size_t end = std::min(begin + shift + block_size, batch_size); // Apply shift for threads > 0 to maintain word alignment across blocks. - if (tid > 0) begin += shift; + if (tid > 0) { + begin += shift; + } for (size_t rid = begin; rid < end; ++rid) { const auto& line = batch.GetLine(rid); for (size_t i = 0; i < line.Size(); ++i) { @@ -491,7 +495,9 @@ class ColumnMatrix { size_t begin = block_size * tid; size_t end = std::min(begin + shift + block_size, batch_size); // Apply shift for threads > 0 to maintain word alignment across blocks. - if (tid > 0) begin += shift; + if (tid > 0) { + begin += shift; + } size_t k = 0; for (size_t rid = begin; rid < end; ++rid) { From 0b7037a99b0a8da71c0d48d76bf3520d50a7ba46 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin <> Date: Mon, 21 Jul 2025 00:31:25 -0700 Subject: [PATCH 26/35] update comment --- src/common/column_matrix.h | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 2a06eba3d176..970a1e819fb8 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -438,11 +438,18 @@ class ColumnMatrix { block_size = DivRoundUp(block_size, MissingIndicator::BitFieldT::kValueSize) * MissingIndicator::BitFieldT::kValueSize; /* - * If base_rowid > 0 we need to shift the blocks boundaries. - * Otherwise the two threads may operate with the single word of bitfield. + * To prevent race conditions on the bitfield, we ensure each thread operates on + * distinct 32-bit words. If a data batch (starting at `base_rowid`) doesn't align + * with a word boundary, this `shift` is calculated. It represents the number of rows + * the first thread must process to reach the next aligned word. This guarantees all + * subsequent thread workloads start on a clean boundary, making parallel updates safe. */ size_t shift = MissingIndicator::BitFieldT::kValueSize - (base_rowid % MissingIndicator::BitFieldT::kValueSize); + /* + * If `base_rowid` is already on a word boundary, the calculation results in + * `kValueSize`. In this case, no shift is needed. + */ if (shift == MissingIndicator::BitFieldT::kValueSize) { shift = 0; } From 820b79aa72f2359834013b17b6c14864acf63700 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Tue, 4 Nov 2025 18:01:18 +0100 Subject: [PATCH 27/35] Update src/common/column_matrix.h Co-authored-by: Victoriya Fedotova --- src/common/column_matrix.h | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 970a1e819fb8..70615c7c3580 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -425,6 +425,32 @@ class ColumnMatrix { size_t const batch_size = batch.Size(); // Parallel sparse batch processing + // + // This section processes the input batch in parallel across multiple threads. + // + // rows [base_rowid, batch_size + base_rowid) are divided into `n_threads` blocks. + // Threads process the assigned blocks of rows in parallel. + // + // As the indicator of the missing elements is stored in a bitfield, to ensure thread-safe + // access to the bitfield, each underlying word of the bitfield + // of size MissingIndicator::BitFieldT::kValueSize should be processed by + // a single thread. Therefore, we align the row-blocks accordingly. + // + // block_size - size of the row block assigned to each thread, divisible by the kValueSize. + // shift - adjustment applied to the starting row (base_rowid) of each thread (except for the 0-th thread) + // to ensure that each thread starts processing from a word boundary in the bitfield. + // + // 0-th thread processes rows [base_rowid, base_rowid + shift + block_size) + // 1-st thread processes rows [base_rowid + shift + block_size, base_rowid + shift + 2 * block_size) + // 2-nd thread processes rows [base_rowid + shift + 2 * block_size, base_rowid + shift + 3 * block_size) + // ... + // (n_threads-1)-th thread processes rows [base_rowid + shift + (n_threads-1) * block_size, base_rowid + batch_size) + // + // Computations are done in two passes: + // 1) Counting non-zero elements per feature per thread to determine their offsets for the next step. + // a) Counting non-zero elements per feature per thread. + // b) Aggregating counts to determine offsets. + // 2) Placing elements into the sparse structure using calculated offsets of non-zero elements. dmlc::OMPException exc; std::vector n_elements((n_threads + 1) * n_features, 0); std::vector k_offsets(n_threads + 1, 0); From 6d553fb8b3734bd54e48f68b79fc1d155dd40a7e Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Tue, 4 Nov 2025 18:03:03 +0100 Subject: [PATCH 28/35] Update src/common/column_matrix.h Co-authored-by: Victoriya Fedotova --- src/common/column_matrix.h | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 70615c7c3580..20029a69378a 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -452,8 +452,14 @@ class ColumnMatrix { // b) Aggregating counts to determine offsets. // 2) Placing elements into the sparse structure using calculated offsets of non-zero elements. dmlc::OMPException exc; - std::vector n_elements((n_threads + 1) * n_features, 0); - std::vector k_offsets(n_threads + 1, 0); + std::vector n_elements((n_threads + 1) * n_features, 0); // number of non-zero elements per feature per thread + // n_elements[tid * n_features + fid] = + // number of non-zero elements of feature fid processed by thread tid + // n_elements[n_threads * n_features + fid] = + // total number of non-zero elements of feature fid + std::vector k_offsets(n_threads + 1, 0); // offsets of non-zero elements for each thread + // k_offsets[0] = 0; + // k_offsets[tid] - starting offset of non-zero elements processed by thread tid size_t block_size = DivRoundUp(batch_size, n_threads); /* From ea5c4febcc9908b9c66ff859aecde6dc041064c1 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Tue, 4 Nov 2025 18:04:13 +0100 Subject: [PATCH 29/35] Update src/common/column_matrix.h Co-authored-by: Victoriya Fedotova --- src/common/column_matrix.h | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 20029a69378a..9c1ca7a061ee 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -460,15 +460,11 @@ class ColumnMatrix { std::vector k_offsets(n_threads + 1, 0); // offsets of non-zero elements for each thread // k_offsets[0] = 0; // k_offsets[tid] - starting offset of non-zero elements processed by thread tid - size_t block_size = DivRoundUp(batch_size, n_threads); + const auto word32size = MissingIndicator::BitFieldT::kValueSize; - /* - * We use bitfield as a missing indicator. To ensure thread safe access to the bitfield - * each underlying word of the bitfiled should be processed by a single thread. - * So we need to align the row-blocks. - */ - block_size = DivRoundUp(block_size, MissingIndicator::BitFieldT::kValueSize) * - MissingIndicator::BitFieldT::kValueSize; + size_t block_size = DivRoundUp(batch_size, n_threads); // preliminary block size + // align block_size to be multiple of kValueSize, so that each thread processes full words in the bitfield + block_size = DivRoundUp(block_size, word32size) * word32size; /* * To prevent race conditions on the bitfield, we ensure each thread operates on * distinct 32-bit words. If a data batch (starting at `base_rowid`) doesn't align From 7dc15f29ff7cbae28be9dd31e3612f6c333a65e3 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Tue, 4 Nov 2025 18:05:45 +0100 Subject: [PATCH 30/35] Update src/common/column_matrix.h Co-authored-by: Victoriya Fedotova --- src/common/column_matrix.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 9c1ca7a061ee..9a2c6cfa6eeb 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -511,6 +511,11 @@ class ColumnMatrix { exc.Rethrow(); // Parallel feature processing to aggregate counts & calculate offsets. + // + // Compute the number of non-zero elements per feature per thread (n_elements) + // and the offsets to non-zero elements per thread (k_offsets). + // + // The final values of n_elements and k_offsets will be calculated after counting the number of non-zero elements per thread. ParallelFor(n_features, n_threads, [&](auto fid) { n_elements[fid] += num_nonzeros_[fid]; for (int tid = 0; tid < n_threads; ++tid) { From 390efd1fc438d86a876f5bc0a6f1c4748e0100d7 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Tue, 4 Nov 2025 18:06:24 +0100 Subject: [PATCH 31/35] Update src/common/column_matrix.h Co-authored-by: Victoriya Fedotova --- src/common/column_matrix.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 9a2c6cfa6eeb..514be6856260 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -529,7 +529,9 @@ class ColumnMatrix { // Parallel row processing to place data using offsets into sparse structure. #pragma omp parallel num_threads(n_threads) { - std::vector nnz_offsets(n_features, 0); + std::vector nnz_offsets(n_features, 0); // offsets of non-zero elements per feature for the current thread + // nnz_offsets[fid] = + // number of non-zero elements of feature fid already processed by this thread exc.Run([&, is_valid, base_rowid, row_index]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; From 3250cc0fc595b2119d7d5a8ef277b2c8133f3634 Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Tue, 4 Nov 2025 18:06:45 +0100 Subject: [PATCH 32/35] Update src/common/column_matrix.h Co-authored-by: Victoriya Fedotova --- src/common/column_matrix.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 514be6856260..7000a84b5d8a 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -548,8 +548,8 @@ class ColumnMatrix { auto coo = line.GetElement(i); if (is_valid(coo)) { auto fid = coo.column_idx; - const uint32_t bin_id = row_index[k_offsets[tid] + k]; - size_t nnz = n_elements[tid * n_features + fid] + nnz_offsets[fid]; + const uint32_t bin_id = row_index[k_offsets[tid] + k]; // get the correct offset for this thread + size_t nnz = n_elements[tid * n_features + fid] + nnz_offsets[fid]; // calculate the correct nnz for this feature and this thread SetBinSparse(bin_id, rid + base_rowid, fid, local_index, nnz); ++k; nnz_offsets[fid] += (type_[fid] != kDenseColumn); From a8df33b83fe49c0672af5c2387d1551fc222434d Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Wed, 5 Nov 2025 09:45:22 +0100 Subject: [PATCH 33/35] linting --- src/common/column_matrix.h | 108 +++++++++++++++++++++---------------- 1 file changed, 61 insertions(+), 47 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 7000a84b5d8a..8f7f9039f0e0 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -424,46 +424,59 @@ class ColumnMatrix { ColumnBinT* local_index = reinterpret_cast(index_.data()); size_t const batch_size = batch.Size(); - // Parallel sparse batch processing - // - // This section processes the input batch in parallel across multiple threads. - // - // rows [base_rowid, batch_size + base_rowid) are divided into `n_threads` blocks. - // Threads process the assigned blocks of rows in parallel. - // - // As the indicator of the missing elements is stored in a bitfield, to ensure thread-safe - // access to the bitfield, each underlying word of the bitfield - // of size MissingIndicator::BitFieldT::kValueSize should be processed by - // a single thread. Therefore, we align the row-blocks accordingly. - // - // block_size - size of the row block assigned to each thread, divisible by the kValueSize. - // shift - adjustment applied to the starting row (base_rowid) of each thread (except for the 0-th thread) - // to ensure that each thread starts processing from a word boundary in the bitfield. - // - // 0-th thread processes rows [base_rowid, base_rowid + shift + block_size) - // 1-st thread processes rows [base_rowid + shift + block_size, base_rowid + shift + 2 * block_size) - // 2-nd thread processes rows [base_rowid + shift + 2 * block_size, base_rowid + shift + 3 * block_size) - // ... - // (n_threads-1)-th thread processes rows [base_rowid + shift + (n_threads-1) * block_size, base_rowid + batch_size) - // - // Computations are done in two passes: - // 1) Counting non-zero elements per feature per thread to determine their offsets for the next step. - // a) Counting non-zero elements per feature per thread. - // b) Aggregating counts to determine offsets. - // 2) Placing elements into the sparse structure using calculated offsets of non-zero elements. + /* Parallel sparse batch processing + * + * This section processes the input batch in parallel across multiple threads. + * + * rows [base_rowid, batch_size + base_rowid) are divided into `n_threads` blocks. + * Threads process the assigned blocks of rows in parallel. + * + * As the indicator of the missing elements is stored in a bitfield, to ensure thread-safe + * access to the bitfield, each underlying word of the bitfield + * of size MissingIndicator::BitFieldT::kValueSize should be processed by + * a single thread. Therefore, we align the row-blocks accordingly. + * + * block_size - size of the row block assigned to each thread, divisible by the kValueSize. + * shift - adjustment applied to the starting row (base_rowid) + * of each thread (except for the 0-th thread) + * to ensure that each thread starts processing from a word boundary in the bitfield. + * + * 0-th thread processes rows + * [base_rowid, base_rowid + shift + block_size) + * 1-st thread processes rows + * [base_rowid + shift + block_size, base_rowid + shift + 2 * block_size) + * 2-nd thread processes rows + * [base_rowid + shift + 2 * block_size, base_rowid + shift + 3 * block_size) + * ... + * (n_threads-1)-th thread processes rows + * [base_rowid + shift + (n_threads-1) * block_size, base_rowid + batch_size) + * + * Computations are done in two passes: + * 1) Counting non-zero elements per feature per thread + * to determine their offsets for the next step. + * a) Counting non-zero elements per feature per thread. + * b) Aggregating counts to determine offsets. + * 2) Placing elements into the sparse structure using calculated offsets of non-zero elements. + */ + + // number of non-zero elements per feature per thread dmlc::OMPException exc; - std::vector n_elements((n_threads + 1) * n_features, 0); // number of non-zero elements per feature per thread - // n_elements[tid * n_features + fid] = - // number of non-zero elements of feature fid processed by thread tid - // n_elements[n_threads * n_features + fid] = - // total number of non-zero elements of feature fid - std::vector k_offsets(n_threads + 1, 0); // offsets of non-zero elements for each thread - // k_offsets[0] = 0; - // k_offsets[tid] - starting offset of non-zero elements processed by thread tid + std::vector n_elements((n_threads + 1) * n_features, 0); + /* n_elements[tid * n_features + fid] = + * number of non-zero elements of feature fid processed by thread tid + * n_elements[n_threads * n_features + fid] = + * total number of non-zero elements of feature fid + */ + + // offsets of non-zero elements for each thread + std::vector k_offsets(n_threads + 1, 0); + // k_offsets[0] = 0; + // k_offsets[tid] - starting offset of non-zero elements processed by thread tid const auto word32size = MissingIndicator::BitFieldT::kValueSize; size_t block_size = DivRoundUp(batch_size, n_threads); // preliminary block size - // align block_size to be multiple of kValueSize, so that each thread processes full words in the bitfield + // align block_size to be multiple of kValueSize, + // so that each thread processes full words in the bitfield block_size = DivRoundUp(block_size, word32size) * word32size; /* * To prevent race conditions on the bitfield, we ensure each thread operates on @@ -472,15 +485,12 @@ class ColumnMatrix { * the first thread must process to reach the next aligned word. This guarantees all * subsequent thread workloads start on a clean boundary, making parallel updates safe. */ - size_t shift = MissingIndicator::BitFieldT::kValueSize - - (base_rowid % MissingIndicator::BitFieldT::kValueSize); + size_t shift = word32size - (base_rowid % word32size); /* * If `base_rowid` is already on a word boundary, the calculation results in * `kValueSize`. In this case, no shift is needed. */ - if (shift == MissingIndicator::BitFieldT::kValueSize) { - shift = 0; - } + if (shift == word32size) shift = 0; // Parallel row processing for thread-local counting. #pragma omp parallel num_threads(n_threads) @@ -515,7 +525,8 @@ class ColumnMatrix { // Compute the number of non-zero elements per feature per thread (n_elements) // and the offsets to non-zero elements per thread (k_offsets). // - // The final values of n_elements and k_offsets will be calculated after counting the number of non-zero elements per thread. + // The final values of n_elements and k_offsets will be calculated + // after counting the number of non-zero elements per thread. ParallelFor(n_features, n_threads, [&](auto fid) { n_elements[fid] += num_nonzeros_[fid]; for (int tid = 0; tid < n_threads; ++tid) { @@ -529,9 +540,10 @@ class ColumnMatrix { // Parallel row processing to place data using offsets into sparse structure. #pragma omp parallel num_threads(n_threads) { - std::vector nnz_offsets(n_features, 0); // offsets of non-zero elements per feature for the current thread - // nnz_offsets[fid] = - // number of non-zero elements of feature fid already processed by this thread + // offsets of non-zero elements per feature for the current thread + std::vector nnz_offsets(n_features, 0); + // nnz_offsets[fid] = + // number of non-zero elements of feature fid already processed by this thread exc.Run([&, is_valid, base_rowid, row_index]() { int tid = omp_get_thread_num(); size_t begin = block_size * tid; @@ -548,8 +560,10 @@ class ColumnMatrix { auto coo = line.GetElement(i); if (is_valid(coo)) { auto fid = coo.column_idx; - const uint32_t bin_id = row_index[k_offsets[tid] + k]; // get the correct offset for this thread - size_t nnz = n_elements[tid * n_features + fid] + nnz_offsets[fid]; // calculate the correct nnz for this feature and this thread + // get the correct offset for this thread + const uint32_t bin_id = row_index[k_offsets[tid] + k]; + // calculate the correct nnz for this feature and this thread + size_t nnz = n_elements[tid * n_features + fid] + nnz_offsets[fid]; SetBinSparse(bin_id, rid + base_rowid, fid, local_index, nnz); ++k; nnz_offsets[fid] += (type_[fid] != kDenseColumn); From 0fcd8a71a551aa0085f48cdb893bc10cba89879e Mon Sep 17 00:00:00 2001 From: Dmitry Razdoburdin Date: Wed, 5 Nov 2025 09:49:11 +0100 Subject: [PATCH 34/35] remove whitespace --- src/common/column_matrix.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index 8f7f9039f0e0..0c687073b8a4 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -541,7 +541,7 @@ class ColumnMatrix { #pragma omp parallel num_threads(n_threads) { // offsets of non-zero elements per feature for the current thread - std::vector nnz_offsets(n_features, 0); + std::vector nnz_offsets(n_features, 0); // nnz_offsets[fid] = // number of non-zero elements of feature fid already processed by this thread exc.Run([&, is_valid, base_rowid, row_index]() { From 9f1d131db994642d15354b4ccae121817944a8bf Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Thu, 29 Jan 2026 17:48:54 +0800 Subject: [PATCH 35/35] typos. --- src/common/column_matrix.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/common/column_matrix.h b/src/common/column_matrix.h index d9f906f5a386..3521f74adabb 100644 --- a/src/common/column_matrix.h +++ b/src/common/column_matrix.h @@ -166,7 +166,7 @@ class ColumnMatrix { } /** - * @param feature_offsets Offest of the first element for each feature + * @param feature_offsets Offset of the first element for each feature * @param type Type of each column (Dense or Sparse). */ void InitOffsetsPadded(const RefResourceView& feature_offsets, @@ -178,9 +178,9 @@ class ColumnMatrix { } /* - * For missing indicator feature offsets are alligned to be a factor of + * For missing indicator feature offsets are aligned to be a factor of * BitFieldT::kValueSize (4 bytes). - * This is critical requariment for thread-safe access to bitfield. + * This is critical requirement for thread-safe access to bitfield. * Each word processed by one thread. */ for (std::size_t fid = 1; fid < feature_offsets.size(); ++fid) {