Support column split in approx tree method (#8847)

tests/cpp/tree/test_histmaker.cc

@@ -1,5 +1,4 @@
 #include <gtest/gtest.h>
-
 #include <xgboost/tree_model.h>
 #include <xgboost/tree_updater.h>
 
@@ -8,26 +7,34 @@
 namespace xgboost {
 namespace tree {
 
-TEST(GrowHistMaker, InteractionConstraint) {
-  size_t constexpr kRows = 32;
-  size_t constexpr kCols = 16;
+std::shared_ptr<DMatrix> GenerateDMatrix(std::size_t rows, std::size_t cols){
+  return RandomDataGenerator{rows, cols, 0.6f}.Seed(3).GenerateDMatrix();
+}
 
-  Context ctx;
-
-  auto p_dmat = RandomDataGenerator{kRows, kCols, 0.6f}.Seed(3).GenerateDMatrix();
-
-  HostDeviceVector<GradientPair> gradients (kRows);
-  std::vector<GradientPair>& h_gradients = gradients.HostVector();
+std::unique_ptr<HostDeviceVector<GradientPair>> GenerateGradients(std::size_t rows) {
+  auto p_gradients = std::make_unique<HostDeviceVector<GradientPair>>(rows);
+  auto& h_gradients = p_gradients->HostVector();
 
   xgboost::SimpleLCG gen;
   xgboost::SimpleRealUniformDistribution<bst_float> dist(0.0f, 1.0f);
 
-  for (size_t i = 0; i < kRows; ++i) {
-    bst_float grad = dist(&gen);
-    bst_float hess = dist(&gen);
-    h_gradients[i] = GradientPair(grad, hess);
+  for (auto i = 0; i < rows; ++i) {
+    auto grad = dist(&gen);
+    auto hess = dist(&gen);
+    h_gradients[i] = GradientPair{grad, hess};
   }
 
+  return p_gradients;
+}
+
+TEST(GrowHistMaker, InteractionConstraint)
+{
+  auto constexpr kRows = 32;
+  auto constexpr kCols = 16;
+  auto p_dmat = GenerateDMatrix(kRows, kCols);
+  auto p_gradients = GenerateGradients(kRows);
+
+  Context ctx;
   {
     // With constraints
     RegTree tree;
@@ -39,7 +46,7 @@ TEST(GrowHistMaker, InteractionConstraint) {
         {"interaction_constraints", "[[0, 1]]"},
         {"num_feature", std::to_string(kCols)}});
     std::vector<HostDeviceVector<bst_node_t>> position(1);
-    updater->Update(&gradients, p_dmat.get(), position, {&tree});
+    updater->Update(p_gradients.get(), p_dmat.get(), position, {&tree});
 
     ASSERT_EQ(tree.NumExtraNodes(), 4);
     ASSERT_EQ(tree[0].SplitIndex(), 1);
@@ -56,7 +63,7 @@ TEST(GrowHistMaker, InteractionConstraint) {
         TreeUpdater::Create("grow_histmaker", &ctx, ObjInfo{ObjInfo::kRegression})};
     updater->Configure(Args{{"num_feature", std::to_string(kCols)}});
     std::vector<HostDeviceVector<bst_node_t>> position(1);
-    updater->Update(&gradients, p_dmat.get(), position, {&tree});
+    updater->Update(p_gradients.get(), p_dmat.get(), position, {&tree});
 
     ASSERT_EQ(tree.NumExtraNodes(), 10);
     ASSERT_EQ(tree[0].SplitIndex(), 1);
@@ -66,5 +73,53 @@ TEST(GrowHistMaker, InteractionConstraint) {
   }
 }
 
+namespace {
+void TestColumnSplit(int32_t rows, int32_t cols, RegTree const& expected_tree) {
+  auto p_dmat = GenerateDMatrix(rows, cols);
+  auto p_gradients = GenerateGradients(rows);
+  Context ctx;
+  std::unique_ptr<TreeUpdater> updater{
+      TreeUpdater::Create("grow_histmaker", &ctx, ObjInfo{ObjInfo::kRegression})};
+  updater->Configure(Args{{"num_feature", std::to_string(cols)}});
+  std::vector<HostDeviceVector<bst_node_t>> position(1);
+
+  std::unique_ptr<DMatrix> sliced{
+      p_dmat->SliceCol(collective::GetWorldSize(), collective::GetRank())};
+
+  RegTree tree;
+  tree.param.num_feature = cols;
+  updater->Update(p_gradients.get(), sliced.get(), position, {&tree});
+
+  EXPECT_EQ(tree.NumExtraNodes(), 10);
+  EXPECT_EQ(tree[0].SplitIndex(), 1);
+
+  EXPECT_NE(tree[tree[0].LeftChild()].SplitIndex(), 0);
+  EXPECT_NE(tree[tree[0].RightChild()].SplitIndex(), 0);
+
+  EXPECT_EQ(tree, expected_tree);
+}
+}  // anonymous namespace
+
+TEST(GrowHistMaker, ColumnSplit) {
+  auto constexpr kRows = 32;
+  auto constexpr kCols = 16;
+
+  RegTree expected_tree;
+  expected_tree.param.num_feature = kCols;
+  {
+    auto p_dmat = GenerateDMatrix(kRows, kCols);
+    auto p_gradients = GenerateGradients(kRows);
+    Context ctx;
+    std::unique_ptr<TreeUpdater> updater{
+        TreeUpdater::Create("grow_histmaker", &ctx, ObjInfo{ObjInfo::kRegression})};
+    updater->Configure(Args{{"num_feature", std::to_string(kCols)}});
+    std::vector<HostDeviceVector<bst_node_t>> position(1);
+    updater->Update(p_gradients.get(), p_dmat.get(), position, {&expected_tree});
+  }
+
+  auto constexpr kWorldSize = 2;
+  RunWithInMemoryCommunicator(kWorldSize, TestColumnSplit, kRows, kCols, std::cref(expected_tree));
+}
+
 }  // namespace tree
 }  // namespace xgboost