Remove omp_get_max_threads in data. (#7588)

src/data/data.cc

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2015-2021 by Contributors
+ * Copyright 2015-2022 by XGBoost Contributors
  * \file data.cc
  */
 #include <dmlc/registry.h>
@@ -1001,15 +1001,14 @@ DMatrix::Create(data::IteratorAdapter<DataIterHandle, XGBCallbackDataIterNext,
                                       XGBoostBatchCSR> *adapter,
                 float missing, int nthread, const std::string &cache_prefix);
 
-SparsePage SparsePage::GetTranspose(int num_columns) const {
+SparsePage SparsePage::GetTranspose(int num_columns, int32_t n_threads) const {
   SparsePage transpose;
   common::ParallelGroupBuilder<Entry, bst_row_t> builder(&transpose.offset.HostVector(),
                                                          &transpose.data.HostVector());
-  const int nthread = omp_get_max_threads();
-  builder.InitBudget(num_columns, nthread);
+  builder.InitBudget(num_columns, n_threads);
   long batch_size = static_cast<long>(this->Size());  // NOLINT(*)
   auto page = this->GetView();
-  common::ParallelFor(batch_size, [&](long i) {  // NOLINT(*)
+  common::ParallelFor(batch_size, n_threads, [&](long i) {  // NOLINT(*)
     int tid = omp_get_thread_num();
     auto inst = page[i];
     for (const auto& entry : inst) {
@@ -1017,7 +1016,7 @@ SparsePage SparsePage::GetTranspose(int num_columns) const {
     }
   });
   builder.InitStorage();
-  common::ParallelFor(batch_size, [&](long i) {  // NOLINT(*)
+  common::ParallelFor(batch_size, n_threads, [&](long i) {  // NOLINT(*)
     int tid = omp_get_thread_num();
     auto inst = page[i];
     for (const auto& entry : inst) {
@@ -1059,8 +1058,6 @@ uint64_t SparsePage::Push(const AdapterBatchT& batch, float missing, int nthread
   constexpr bool kIsRowMajor = AdapterBatchT::kIsRowMajor;
   // Allow threading only for row-major case as column-major requires O(nthread*batch_size) memory
   nthread = kIsRowMajor ? nthread : 1;
-  // Set number of threads but keep old value so we can reset it after
-  int nthread_original = common::OmpSetNumThreadsWithoutHT(&nthread);
   if (!kIsRowMajor) {
     CHECK_EQ(nthread, 1);
   }
@@ -1085,7 +1082,6 @@ uint64_t SparsePage::Push(const AdapterBatchT& batch, float missing, int nthread
   expected_rows = kIsRowMajor ? batch_size : expected_rows;
   uint64_t max_columns = 0;
   if (batch_size == 0) {
-    omp_set_num_threads(nthread_original);
     return max_columns;
   }
   const size_t thread_size = batch_size / nthread;
@@ -1154,7 +1150,6 @@ uint64_t SparsePage::Push(const AdapterBatchT& batch, float missing, int nthread
     });
   }
   exec.Rethrow();
-  omp_set_num_threads(nthread_original);
 
   return max_columns;
 }

src/data/gradient_index.cc

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2021 by Contributors
+ * Copyright 2017-2022 by XGBoost Contributors
  * \brief Data type for fast histogram aggregation.
  */
 #include <algorithm>
@@ -126,17 +126,16 @@ void GHistIndexMatrix::PushBatch(SparsePage const &batch,
   });
 }
 
-void GHistIndexMatrix::Init(DMatrix *p_fmat, int max_bins, bool sorted_sketch,
+void GHistIndexMatrix::Init(DMatrix *p_fmat, int max_bins, bool sorted_sketch, int32_t n_threads,
                             common::Span<float> hess) {
   // We use sorted sketching for approx tree method since it's more efficient in
   // computation time (but higher memory usage).
-  cut = common::SketchOnDMatrix(p_fmat, max_bins, sorted_sketch, hess);
+  cut = common::SketchOnDMatrix(p_fmat, max_bins, n_threads, sorted_sketch, hess);
 
   max_num_bins = max_bins;
-  const int32_t nthread = omp_get_max_threads();
   const uint32_t nbins = cut.Ptrs().back();
   hit_count.resize(nbins, 0);
-  hit_count_tloc_.resize(nthread * nbins, 0);
+  hit_count_tloc_.resize(n_threads * nbins, 0);
 
   this->p_fmat = p_fmat;
   size_t new_size = 1;
@@ -154,7 +153,7 @@ void GHistIndexMatrix::Init(DMatrix *p_fmat, int max_bins, bool sorted_sketch,
   auto ft = p_fmat->Info().feature_types.ConstHostSpan();
 
   for (const auto &batch : p_fmat->GetBatches<SparsePage>()) {
-    this->PushBatch(batch, ft, rbegin, prev_sum, nbins, nthread);
+    this->PushBatch(batch, ft, rbegin, prev_sum, nbins, n_threads);
     prev_sum = row_ptr[rbegin + batch.Size()];
     rbegin += batch.Size();
   }

src/data/gradient_index.h

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2021 by Contributors
+ * Copyright 2017-2022 by XGBoost Contributors
  * \brief Data type for fast histogram aggregation.
  */
 #ifndef XGBOOST_DATA_GRADIENT_INDEX_H_
@@ -19,9 +19,8 @@ namespace xgboost {
  *  index for CPU histogram.  On GPU ellpack page is used.
  */
 class GHistIndexMatrix {
-  void PushBatch(SparsePage const &batch, common::Span<FeatureType const> ft,
-                 size_t rbegin, size_t prev_sum, uint32_t nbins,
-                 int32_t n_threads);
+  void PushBatch(SparsePage const& batch, common::Span<FeatureType const> ft, size_t rbegin,
+                 size_t prev_sum, uint32_t nbins, int32_t n_threads);
 
  public:
   /*! \brief row pointer to rows by element position */
@@ -37,11 +36,13 @@ class GHistIndexMatrix {
   size_t base_rowid{0};
 
   GHistIndexMatrix() = default;
-  GHistIndexMatrix(DMatrix* x, int32_t max_bin, bool sorted_sketch, common::Span<float> hess = {}) {
-    this->Init(x, max_bin, sorted_sketch, hess);
+  GHistIndexMatrix(DMatrix* x, int32_t max_bin, bool sorted_sketch, int32_t n_threads,
+                   common::Span<float> hess = {}) {
+    this->Init(x, max_bin, sorted_sketch, n_threads, hess);
   }
   // Create a global histogram matrix, given cut
-  void Init(DMatrix* p_fmat, int max_num_bins, bool sorted_sketch, common::Span<float> hess);
+  void Init(DMatrix* p_fmat, int max_num_bins, bool sorted_sketch, int32_t n_threads,
+            common::Span<float> hess);
   void Init(SparsePage const& page, common::Span<FeatureType const> ft,
             common::HistogramCuts const& cuts, int32_t max_bins_per_feat, bool is_dense,
             int32_t n_threads);

src/data/simple_dmatrix.cc

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2014~2021 by Contributors
+ * Copyright 2014~2022 by XGBoost Contributors
  * \file simple_dmatrix.cc
  * \brief the input data structure for gradient boosting
  * \author Tianqi Chen
@@ -55,7 +55,7 @@ BatchSet<SparsePage> SimpleDMatrix::GetRowBatches() {
 BatchSet<CSCPage> SimpleDMatrix::GetColumnBatches() {
   // column page doesn't exist, generate it
   if (!column_page_) {
-    column_page_.reset(new CSCPage(sparse_page_->GetTranspose(info_.num_col_)));
+    column_page_.reset(new CSCPage(sparse_page_->GetTranspose(info_.num_col_, ctx_.Threads())));
   }
   auto begin_iter =
       BatchIterator<CSCPage>(new SimpleBatchIteratorImpl<CSCPage>(column_page_));
@@ -66,7 +66,7 @@ BatchSet<SortedCSCPage> SimpleDMatrix::GetSortedColumnBatches() {
   // Sorted column page doesn't exist, generate it
   if (!sorted_column_page_) {
     sorted_column_page_.reset(
-        new SortedCSCPage(sparse_page_->GetTranspose(info_.num_col_)));
+        new SortedCSCPage(sparse_page_->GetTranspose(info_.num_col_, ctx_.Threads())));
     sorted_column_page_->SortRows();
   }
   auto begin_iter = BatchIterator<SortedCSCPage>(
@@ -99,7 +99,8 @@ BatchSet<GHistIndexMatrix> SimpleDMatrix::GetGradientIndex(const BatchParam& par
     CHECK_EQ(param.gpu_id, -1);
     // Used only by approx.
     auto sorted_sketch = param.regen;
-    gradient_index_.reset(new GHistIndexMatrix(this, param.max_bin, sorted_sketch, param.hess));
+    gradient_index_.reset(
+        new GHistIndexMatrix(this, param.max_bin, sorted_sketch, this->ctx_.Threads(), param.hess));
     batch_param_ = param;
     CHECK_EQ(batch_param_.hess.data(), param.hess.data());
   }
@@ -110,6 +111,8 @@ BatchSet<GHistIndexMatrix> SimpleDMatrix::GetGradientIndex(const BatchParam& par
 
 template <typename AdapterT>
 SimpleDMatrix::SimpleDMatrix(AdapterT* adapter, float missing, int nthread) {
+  this->ctx_.nthread = nthread;
+
   std::vector<uint64_t> qids;
   uint64_t default_max = std::numeric_limits<uint64_t>::max();
   uint64_t last_group_id = default_max;
@@ -124,7 +127,7 @@ SimpleDMatrix::SimpleDMatrix(AdapterT* adapter, float missing, int nthread) {
   // Iterate over batches of input data
   while (adapter->Next()) {
     auto& batch = adapter->Value();
-    auto batch_max_columns = sparse_page_->Push(batch, missing, nthread);
+    auto batch_max_columns = sparse_page_->Push(batch, missing, ctx_.Threads());
     inferred_num_columns = std::max(batch_max_columns, inferred_num_columns);
     total_batch_size += batch.Size();
     // Append meta information if available

src/data/simple_dmatrix.h

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2015-2021 by Contributors
+ * Copyright 2015-2022 by XGBoost Contributors
  * \file simple_dmatrix.h
  * \brief In-memory version of DMatrix.
  * \author Tianqi Chen
@@ -61,6 +61,9 @@ class SimpleDMatrix : public DMatrix {
   bool SparsePageExists() const override {
     return true;
   }
+
+ private:
+  GenericParameter ctx_;
 };
 }  // namespace data
 }  // namespace xgboost

src/data/sparse_page_dmatrix.cc

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2014-2021 by Contributors
+ * Copyright 2014-2022 by Contributors
  * \file sparse_page_dmatrix.cc
  * \brief The external memory version of Page Iterator.
  * \author Tianqi Chen
@@ -164,7 +164,8 @@ BatchSet<GHistIndexMatrix> SparsePageDMatrix::GetGradientIndex(const BatchParam&
     // all index here.
     if (!ghist_index_page_ || (param != batch_param_ && param != BatchParam{})) {
       this->InitializeSparsePage();
-      ghist_index_page_.reset(new GHistIndexMatrix{this, param.max_bin, param.regen});
+      ghist_index_page_.reset(
+          new GHistIndexMatrix{this, param.max_bin, param.regen, ctx_.Threads()});
       this->InitializeSparsePage();
       batch_param_ = param;
     }
@@ -181,7 +182,8 @@ BatchSet<GHistIndexMatrix> SparsePageDMatrix::GetGradientIndex(const BatchParam&
     MakeCache(this, ".gradient_index.page", cache_prefix_, &cache_info_);
     // Use sorted sketch for approx.
     auto sorted_sketch = param.regen;
-    auto cuts = common::SketchOnDMatrix(this, param.max_bin, sorted_sketch, param.hess);
+    auto cuts =
+        common::SketchOnDMatrix(this, param.max_bin, ctx_.Threads(), sorted_sketch, param.hess);
     this->InitializeSparsePage();  // reset after use.
 
     batch_param_ = param;

src/data/sparse_page_source.h

@@ -1,5 +1,5 @@
 /*!
- *  Copyright (c) 2014-2021 by Contributors
+ *  Copyright 2014-2022 by XGBoost Contributors
  * \file sparse_page_source.h
  */
 #ifndef XGBOOST_DATA_SPARSE_PAGE_SOURCE_H_
@@ -311,7 +311,7 @@ class CSCPageSource : public PageSourceIncMixIn<CSCPage> {
       auto const &csr = source_->Page();
       this->page_.reset(new CSCPage{});
       // we might be able to optimize this by merging transpose and pushcsc
-      this->page_->PushCSC(csr->GetTranspose(n_features_));
+      this->page_->PushCSC(csr->GetTranspose(n_features_, nthreads_));
       page_->SetBaseRowId(csr->base_rowid);
       this->WriteCache();
     }
@@ -336,7 +336,7 @@ class SortedCSCPageSource : public PageSourceIncMixIn<SortedCSCPage> {
       auto const &csr = this->source_->Page();
       this->page_.reset(new SortedCSCPage{});
       // we might be able to optimize this by merging transpose and pushcsc
-      this->page_->PushCSC(csr->GetTranspose(n_features_));
+      this->page_->PushCSC(csr->GetTranspose(n_features_, nthreads_));
       CHECK_EQ(this->page_->Size(), n_features_);
       CHECK_EQ(this->page_->data.Size(), csr->data.Size());
       this->page_->SortRows();