Optimize cpu sketch allreduce for sparse data. (#6009)

* Bypass RABIT serialization reducer and use custom allgather based merging.

tests/cpp/common/test_quantile.cc

@@ -5,14 +5,122 @@
 
 namespace xgboost {
 namespace common {
+
+TEST(Quantile, LoadBalance) {
+  size_t constexpr kRows = 1000, kCols = 100;
+  auto m = RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix();
+  std::vector<bst_feature_t> cols_ptr;
+  for (auto const &page : m->GetBatches<SparsePage>()) {
+    cols_ptr = HostSketchContainer::LoadBalance(page, kCols, 13);
+  }
+  size_t n_cols = 0;
+  for (size_t i = 1; i < cols_ptr.size(); ++i) {
+    n_cols += cols_ptr[i] - cols_ptr[i - 1];
+  }
+  CHECK_EQ(n_cols, kCols);
+}
+
+void TestDistributedQuantile(size_t rows, size_t cols) {
+  std::string msg {"Skipping AllReduce test"};
+  int32_t constexpr kWorkers = 4;
+  InitRabitContext(msg, kWorkers);
+  auto world = rabit::GetWorldSize();
+  if (world != 1) {
+    ASSERT_EQ(world, kWorkers);
+  } else {
+    return;
+  }
+
+  std::vector<MetaInfo> infos(2);
+  auto& h_weights = infos.front().weights_.HostVector();
+  h_weights.resize(rows);
+  SimpleLCG lcg;
+  SimpleRealUniformDistribution<float> dist(3, 1000);
+  std::generate(h_weights.begin(), h_weights.end(), [&]() { return dist(&lcg); });
+  std::vector<bst_row_t> column_size(cols, rows);
+  size_t n_bins = 64;
+
+  // Generate cuts for distributed environment.
+  auto sparsity = 0.5f;
+  auto rank = rabit::GetRank();
+  HostSketchContainer sketch_distributed(column_size, n_bins, false);
+  auto m = RandomDataGenerator{rows, cols, sparsity}
+               .Seed(rank)
+               .Lower(.0f)
+               .Upper(1.0f)
+               .GenerateDMatrix();
+  for (auto const &page : m->GetBatches<SparsePage>()) {
+    sketch_distributed.PushRowPage(page, m->Info());
+  }
+  HistogramCuts distributed_cuts;
+  sketch_distributed.MakeCuts(&distributed_cuts);
+
+  // Generate cuts for single node environment
+  rabit::Finalize();
+  CHECK_EQ(rabit::GetWorldSize(), 1);
+  std::for_each(column_size.begin(), column_size.end(), [=](auto& size) { size *= world; });
+  HostSketchContainer sketch_on_single_node(column_size, n_bins, false);
+  for (auto rank = 0; rank < world; ++rank) {
+    auto m = RandomDataGenerator{rows, cols, sparsity}
+                 .Seed(rank)
+                 .Lower(.0f)
+                 .Upper(1.0f)
+                 .GenerateDMatrix();
+    for (auto const &page : m->GetBatches<SparsePage>()) {
+      sketch_on_single_node.PushRowPage(page, m->Info());
+    }
+  }
+
+  HistogramCuts single_node_cuts;
+  sketch_on_single_node.MakeCuts(&single_node_cuts);
+
+  auto const& sptrs = single_node_cuts.Ptrs();
+  auto const& dptrs = distributed_cuts.Ptrs();
+  auto const& svals = single_node_cuts.Values();
+  auto const& dvals = distributed_cuts.Values();
+  auto const& smins = single_node_cuts.MinValues();
+  auto const& dmins = distributed_cuts.MinValues();
+
+  ASSERT_EQ(sptrs.size(), dptrs.size());
+  for (size_t i = 0; i < sptrs.size(); ++i) {
+    ASSERT_EQ(sptrs[i], dptrs[i]);
+  }
+
+  ASSERT_EQ(svals.size(), dvals.size());
+  for (size_t i = 0; i < svals.size(); ++i) {
+    ASSERT_NEAR(svals[i], dvals[i], 2e-2f);
+  }
+
+  ASSERT_EQ(smins.size(), dmins.size());
+  for (size_t i = 0; i < smins.size(); ++i) {
+    ASSERT_FLOAT_EQ(smins[i], dmins[i]);
+  }
+}
+
+TEST(Quantile, DistributedBasic) {
+#if defined(__unix__)
+  constexpr size_t kRows = 10, kCols = 10;
+  TestDistributedQuantile(kRows, kCols);
+#endif
+}
+
+TEST(Quantile, Distributed) {
+#if defined(__unix__)
+  constexpr size_t kRows = 1000, kCols = 200;
+  TestDistributedQuantile(kRows, kCols);
+#endif
+}
+
 TEST(Quantile, SameOnAllWorkers) {
+#if defined(__unix__)
   std::string msg{"Skipping Quantile AllreduceBasic test"};
-  size_t constexpr kWorkers = 4;
+  int32_t constexpr kWorkers = 4;
   InitRabitContext(msg, kWorkers);
   auto world = rabit::GetWorldSize();
   if (world != 1) {
     CHECK_EQ(world, kWorkers);
   } else {
+    LOG(WARNING) << msg;
     return;
   }
 
@@ -72,6 +180,8 @@ TEST(Quantile, SameOnAllWorkers) {
           }
         }
       });
+  rabit::Finalize();
+#endif  // defined(__unix__)
 }
 }  // namespace common
 }  // namespace xgboost