Consistent use of context to specify number of threads. (#8733)

- Use context in all tests.
- Use context in R.
- Use context in C API DMatrix initialization. (0 threads is used as dft).

tests/cpp/tree/hist/test_evaluate_splits.cc

@@ -1,5 +1,5 @@
-/*!
- * Copyright 2021-2022 by XGBoost Contributors
+/**
+ * Copyright 2021-2023 by XGBoost Contributors
  */
 #include <gtest/gtest.h>
 #include <xgboost/base.h>
@@ -14,9 +14,7 @@ namespace xgboost {
 namespace tree {
 void TestEvaluateSplits(bool force_read_by_column) {
   int static constexpr kRows = 8, kCols = 16;
-  auto orig = omp_get_max_threads();
   int32_t n_threads = std::min(omp_get_max_threads(), 4);
-  omp_set_num_threads(n_threads);
   auto sampler = std::make_shared<common::ColumnSampler>();
 
   TrainParam param;
@@ -32,7 +30,7 @@ void TestEvaluateSplits(bool force_read_by_column) {
 
   size_t constexpr kMaxBins = 4;
   // dense, no missing values
-  GHistIndexMatrix gmat(dmat.get(), kMaxBins, 0.5, false, common::OmpGetNumThreads(0));
+  GHistIndexMatrix gmat(dmat.get(), kMaxBins, 0.5, false, AllThreadsForTest());
   common::RowSetCollection row_set_collection;
   std::vector<size_t> &row_indices = *row_set_collection.Data();
   row_indices.resize(kRows);
@@ -80,8 +78,6 @@ void TestEvaluateSplits(bool force_read_by_column) {
       right.SetSubstract(GradStats{total_gpair}, left);
     }
   }
-
-  omp_set_num_threads(orig);
 }
 
 TEST(HistEvaluator, Evaluate) {
@@ -122,7 +118,7 @@ TEST_F(TestPartitionBasedSplit, CPUHist) {
   // check the evaluator is returning the optimal split
   std::vector<FeatureType> ft{FeatureType::kCategorical};
   auto sampler = std::make_shared<common::ColumnSampler>();
-  HistEvaluator<CPUExpandEntry> evaluator{param_, info_, common::OmpGetNumThreads(0), sampler};
+  HistEvaluator<CPUExpandEntry> evaluator{param_, info_, AllThreadsForTest(), sampler};
   evaluator.InitRoot(GradStats{total_gpair_});
   RegTree tree;
   std::vector<CPUExpandEntry> entries(1);
@@ -152,7 +148,7 @@ auto CompareOneHotAndPartition(bool onehot) {
 
   auto sampler = std::make_shared<common::ColumnSampler>();
   auto evaluator =
-      HistEvaluator<CPUExpandEntry>{param, dmat->Info(), common::OmpGetNumThreads(0), sampler};
+      HistEvaluator<CPUExpandEntry>{param, dmat->Info(), AllThreadsForTest(), sampler};
   std::vector<CPUExpandEntry> entries(1);
 
   for (auto const &gmat : dmat->GetBatches<GHistIndexMatrix>({32, param.sparse_threshold})) {
@@ -203,7 +199,7 @@ TEST_F(TestCategoricalSplitWithMissing, HistEvaluator) {
   info.num_col_ = 1;
   info.feature_types = {FeatureType::kCategorical};
   auto evaluator =
-      HistEvaluator<CPUExpandEntry>{param_, info, common::OmpGetNumThreads(0), sampler};
+      HistEvaluator<CPUExpandEntry>{param_, info, AllThreadsForTest(), sampler};
   evaluator.InitRoot(GradStats{parent_sum_});
 
   std::vector<CPUExpandEntry> entries(1);

tests/cpp/tree/hist/test_histogram.cc

@@ -1,7 +1,8 @@
-/*!
- * Copyright 2018-2022 by Contributors
+/**
+ * Copyright 2018-2023 by Contributors
  */
 #include <gtest/gtest.h>
+#include <xgboost/context.h>  // Context
 
 #include <limits>
 
@@ -375,6 +376,7 @@ TEST(CPUHistogram, Categorical) {
 }
 namespace {
 void TestHistogramExternalMemory(BatchParam batch_param, bool is_approx, bool force_read_by_column) {
+  Context ctx;
   size_t constexpr kEntries = 1 << 16;
   auto m = CreateSparsePageDMatrix(kEntries, "cache");
 
@@ -417,7 +419,7 @@ void TestHistogramExternalMemory(BatchParam batch_param, bool is_approx, bool fo
         1, [&](size_t nidx_in_set) { return partition_size.at(nidx_in_set); },
         256};
 
-    multi_build.Reset(total_bins, batch_param, common::OmpGetNumThreads(0), rows_set.size(), false);
+    multi_build.Reset(total_bins, batch_param, ctx.Threads(), rows_set.size(), false);
 
     size_t page_idx{0};
     for (auto const &page : m->GetBatches<GHistIndexMatrix>(batch_param)) {
@@ -438,17 +440,16 @@ void TestHistogramExternalMemory(BatchParam batch_param, bool is_approx, bool fo
     common::RowSetCollection row_set_collection;
     InitRowPartitionForTest(&row_set_collection, n_samples);
 
-    single_build.Reset(total_bins, batch_param, common::OmpGetNumThreads(0), 1, false);
+    single_build.Reset(total_bins, batch_param, ctx.Threads(), 1, false);
     SparsePage concat;
     std::vector<float> hess(m->Info().num_row_, 1.0f);
     for (auto const& page : m->GetBatches<SparsePage>()) {
       concat.Push(page);
     }
 
-    auto cut = common::SketchOnDMatrix(m.get(), batch_param.max_bin, common::OmpGetNumThreads(0),
-                                       false, hess);
+    auto cut = common::SketchOnDMatrix(m.get(), batch_param.max_bin, ctx.Threads(), false, hess);
     GHistIndexMatrix gmat(concat, {}, cut, batch_param.max_bin, false,
-                          std::numeric_limits<double>::quiet_NaN(), common::OmpGetNumThreads(0));
+                          std::numeric_limits<double>::quiet_NaN(), ctx.Threads());
     single_build.BuildHist(0, gmat, &tree, row_set_collection, nodes, {}, h_gpair, force_read_by_column);
     single_page = single_build.Histogram()[0];
   }

tests/cpp/tree/test_prediction_cache.cc

@@ -1,5 +1,5 @@
-/*!
- * Copyright 2021-2022 by XGBoost contributors
+/**
+ * Copyright 2021-2023 by XGBoost contributors
  */
 #include <gtest/gtest.h>
 #include <xgboost/host_device_vector.h>
@@ -62,7 +62,6 @@ class TestPredictionCache : public ::testing::Test {
 
   void RunTest(std::string updater_name) {
     {
-      omp_set_num_threads(1);
       Context ctx;
       ctx.InitAllowUnknown(Args{{"nthread", "8"}});
       if (updater_name == "grow_gpu_hist") {

tests/cpp/tree/test_regen.cc

@@ -1,11 +1,12 @@
-/*!
- * Copyright 2022 XGBoost contributors
+/**
+ * Copyright 2022-2023 XGBoost contributors
  */
 #include <gtest/gtest.h>
 
 #include "../../../src/data/adapter.h"
 #include "../../../src/data/simple_dmatrix.h"
 #include "../helpers.h"
+#include "xgboost/context.h"
 
 namespace xgboost {
 namespace {
@@ -50,7 +51,7 @@ class RegenTest : public ::testing::Test {
     auto dense = RandomDataGenerator{kRows, kCols, 0.5}.GenerateArrayInterface(&storage);
     auto adapter = data::ArrayAdapter(StringView{dense});
     p_fmat_ = std::shared_ptr<DMatrix>(new DMatrixForTest{
-        &adapter, std::numeric_limits<float>::quiet_NaN(), common::OmpGetNumThreads(0)});
+        &adapter, std::numeric_limits<float>::quiet_NaN(), AllThreadsForTest()});
 
     p_fmat_->Info().labels.Reshape(256, 1);
     auto labels = p_fmat_->Info().labels.Data();