[sycl] add partitioning and related tests (#10080)

Co-authored-by: Dmitry Razdoburdin <>

tests/cpp/plugin/test_sycl_gradient_index.cc

@@ -49,7 +49,8 @@ TEST(SyclGradientIndex, Init) {
 
   auto p_fmat = RandomDataGenerator{n_rows, n_columns, 0.3}.GenerateDMatrix();
 
-  sycl::DeviceMatrix dmat(qu, p_fmat.get());
+  sycl::DeviceMatrix dmat;
+  dmat.Init(qu, p_fmat.get());
 
   int max_bins = 256;
   common::GHistIndexMatrix gmat_sycl;

tests/cpp/plugin/test_sycl_partition_builder.cc

@@ -13,6 +13,108 @@
 
 namespace xgboost::sycl::common {
 
+void TestPartitioning(float sparsity, int max_bins) {
+  const size_t num_rows = 16;
+  const size_t num_columns = 1;
+
+  Context ctx;
+  ctx.UpdateAllowUnknown(Args{{"device", "sycl"}});
+
+  DeviceManager device_manager;
+  auto qu = device_manager.GetQueue(ctx.Device());
+
+  auto p_fmat = RandomDataGenerator{num_rows, num_columns, sparsity}.GenerateDMatrix();
+  sycl::DeviceMatrix dmat;
+  dmat.Init(qu, p_fmat.get());
+
+  common::GHistIndexMatrix gmat;
+  gmat.Init(qu, &ctx, dmat, max_bins);
+
+  RowSetCollection row_set_collection;
+  auto& row_indices = row_set_collection.Data();
+  row_indices.Resize(&qu, num_rows);
+  size_t* p_row_indices = row_indices.Data();
+
+  qu.submit([&](::sycl::handler& cgh) {
+    cgh.parallel_for<>(::sycl::range<1>(num_rows),
+                       [p_row_indices](::sycl::item<1> pid) {
+      const size_t idx = pid.get_id(0);
+      p_row_indices[idx] = idx;
+    });
+  }).wait_and_throw();
+  row_set_collection.Init();
+
+  RegTree tree;
+  tree.ExpandNode(0, 0, 0, false, 0, 0, 0, 0, 0, 0, 0);
+
+  const size_t n_nodes = row_set_collection.Size();
+  PartitionBuilder partition_builder;
+  partition_builder.Init(&qu, n_nodes, [&](size_t nid) {
+    return row_set_collection[nid].Size();
+  });
+
+  std::vector<tree::ExpandEntry> nodes;
+  nodes.emplace_back(tree::ExpandEntry(0, tree.GetDepth(0)));
+
+  ::sycl::event event;
+  std::vector<int32_t> split_conditions = {2};
+  partition_builder.Partition(gmat, nodes, row_set_collection,
+                    split_conditions, &tree, &event);
+  qu.wait_and_throw();
+
+  size_t* data_result = const_cast<size_t*>(row_set_collection[0].begin);
+  partition_builder.MergeToArray(0, data_result, &event);
+  qu.wait_and_throw();
+
+  bst_float split_pt = gmat.cut.Values()[split_conditions[0]];
+
+  std::vector<uint8_t> ridx_left(num_rows, 0);
+  std::vector<uint8_t> ridx_right(num_rows, 0);
+  for (auto &batch : gmat.p_fmat->GetBatches<SparsePage>()) {
+    const auto& data_vec = batch.data.HostVector();
+    const auto& offset_vec = batch.offset.HostVector();
+
+    size_t begin = offset_vec[0];
+    for (size_t idx = 0; idx < offset_vec.size() - 1; ++idx) {
+      size_t end = offset_vec[idx + 1];
+      if (begin < end) {
+        const auto& entry = data_vec[begin];
+        if (entry.fvalue < split_pt) {
+          ridx_left[idx] = 1;
+        } else {
+          ridx_right[idx] = 1;
+        }
+      } else {
+        // missing value
+        if (tree[0].DefaultLeft()) {
+          ridx_left[idx] = 1;
+        } else {
+          ridx_right[idx] = 1;
+        }
+      }
+      begin = end;
+    }
+  }
+  auto n_left  = std::accumulate(ridx_left.begin(),  ridx_left.end(),  0);
+  auto n_right = std::accumulate(ridx_right.begin(), ridx_right.end(), 0);
+
+  std::vector<size_t> row_indices_host(num_rows);
+  qu.memcpy(row_indices_host.data(), row_indices.Data(), num_rows * sizeof(size_t));
+  qu.wait_and_throw();
+
+  ASSERT_EQ(n_left,  partition_builder.GetNLeftElems(0));
+  for (size_t i = 0; i < n_left; ++i) {
+    auto idx = row_indices_host[i];
+    ASSERT_EQ(ridx_left[idx], 1);
+  }
+
+  ASSERT_EQ(n_right, partition_builder.GetNRightElems(0));
+  for (size_t i = 0; i < n_right; ++i) {
+    auto idx = row_indices_host[num_rows - 1 - i];
+    ASSERT_EQ(ridx_right[idx], 1);
+  }
+}
+
 TEST(SyclPartitionBuilder, BasicTest) {
   constexpr size_t kNodes = 5;
   // Number of rows for each node
@@ -67,7 +169,7 @@ TEST(SyclPartitionBuilder, BasicTest) {
   std::vector<size_t> v(*std::max_element(rows.begin(), rows.end()));
   size_t row_id = 0;
   for(size_t nid = 0; nid < kNodes; ++nid) {
-    builder.MergeToArray(nid, v.data(), event);
+    builder.MergeToArray(nid, v.data(), &event);
     qu.wait();
 
     // Check that row_id for left side are correct
@@ -88,4 +190,20 @@ TEST(SyclPartitionBuilder, BasicTest) {
   }
 }
 
+TEST(SyclPartitionBuilder, PartitioningSparce) {
+  TestPartitioning(0.3, 256);
+}
+
+TEST(SyclPartitionBuilder, PartitioningDence8Bits) {
+  TestPartitioning(0.0, 256);
+}
+
+TEST(SyclPartitionBuilder, PartitioningDence16Bits) {
+  TestPartitioning(0.0, 256 + 1);
+}
+
+TEST(SyclPartitionBuilder, PartitioningDence32Bits) {
+  TestPartitioning(0.0, (1u << 16) + 1);
+}
+
 }  // namespace xgboost::common