Replace the device model. (#11752)

trivialfis · web-flow · commit 425d136e81a4 · 2025-10-16T17:21:12.000+08:00
- Support leaf, QDM, inplace prediction for multi-target.
- Replace device model with tree-internal storage.
- Cleanup predict kernel.
diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
@@ -262,12 +262,18 @@ class RegTree : public Model {
   Node& operator[](bst_node_t nidx) { return nodes_.HostVector()[nidx]; }
 
  public:
-  /*! \brief get const reference to nodes */
-  [[nodiscard]] const std::vector<Node>& GetNodes() const { return nodes_.ConstHostVector(); }
+  /** @brief Get const reference to nodes */
+  [[nodiscard]] common::Span<Node const> GetNodes(DeviceOrd device) const {
+    CHECK(!this->IsMultiTarget());
+    return device.IsCPU() ? nodes_.ConstHostSpan()
+                          : (nodes_.SetDevice(device), nodes_.ConstDeviceSpan());
+  }
 
-  /*! \brief get const reference to stats */
-  [[nodiscard]] const std::vector<RTreeNodeStat>& GetStats() const {
-    return stats_.ConstHostVector();
+  /** @brief Get const reference to stats */
+  [[nodiscard]] common::Span<RTreeNodeStat const> GetStats(DeviceOrd device) const {
+    CHECK(!this->IsMultiTarget());
+    return device.IsCPU() ? stats_.ConstHostSpan()
+                          : (stats_.SetDevice(device), stats_.ConstDeviceSpan());
   }
 
   /*! \brief get node statistics given nid */
diff --git a/plugin/sycl/predictor/predictor.cc b/plugin/sycl/predictor/predictor.cc
@@ -75,13 +75,13 @@ class DeviceModel {
       if (model.trees[tree_idx]->HasCategoricalSplit()) {
         LOG(FATAL) << "Categorical features are not yet supported by sycl";
       }
-      n_nodes += model.trees[tree_idx]->GetNodes().size();
+      n_nodes += model.trees[tree_idx]->Size();
       first_node_position_host[tree_idx - tree_begin + 1] = n_nodes;
     }
 
     nodes.Resize(n_nodes);
     for (int tree_idx = tree_begin; tree_idx < tree_end; tree_idx++) {
-      auto& src_nodes = model.trees[tree_idx]->GetNodes();
+      auto const& src_nodes = model.trees[tree_idx]->GetNodes(DeviceOrd::CPU());
       size_t n_nodes_shift = first_node_position_host[tree_idx - tree_begin];
       for (size_t node_idx = 0; node_idx < src_nodes.size(); node_idx++) {
         nodes.HostVector()[node_idx + n_nodes_shift] = static_cast<Node>(src_nodes[node_idx]);
diff --git a/python-package/xgboost/testing/ordinal.py b/python-package/xgboost/testing/ordinal.py
@@ -434,9 +434,11 @@ def run_thread_safety(DMatrixT: Type) -> bool:
         return True
 
     futures = []
+    n_cpus = os.cpu_count()
+    assert n_cpus is not None
     for dm in (DMatrix, QuantileDMatrix):
-        with ThreadPoolExecutor(max_workers=10) as e:
-            for _ in range(10):
+        with ThreadPoolExecutor(max_workers=max(n_cpus, 10)) as e:
+            for _ in range(32):
                 fut = e.submit(run_thread_safety, dm)
                 futures.append(fut)
 
diff --git a/src/predictor/cpu_predictor.cc b/src/predictor/cpu_predictor.cc
@@ -640,7 +640,7 @@ class ColumnSplitHelper {
     tree_offsets_.resize(n_trees);
     for (decltype(tree_begin) i = 0; i < n_trees; i++) {
       auto const &tree = *model_.trees[tree_begin_ + i];
-      tree_sizes_[i] = tree.GetNodes().size();
+      tree_sizes_[i] = tree.Size();
     }
     // std::exclusive_scan (only available in c++17) equivalent to get tree offsets.
     tree_offsets_[0] = 0;
diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu
diff --git a/src/tree/hist/evaluate_splits.h b/src/tree/hist/evaluate_splits.h
@@ -447,7 +447,7 @@ class HistEvaluator {
                              tree[candidate.nid].SplitIndex(), left_weight, right_weight);
     evaluator = tree_evaluator_.GetEvaluator();
 
-    snode_.resize(tree.GetNodes().size());
+    snode_.resize(tree.Size());
     snode_.at(left_child).stats = candidate.split.left_sum;
     snode_.at(left_child).root_gain =
         evaluator.CalcGain(candidate.nid, *param_, GradStats{candidate.split.left_sum});
diff --git a/src/tree/tree_view.cc b/src/tree/tree_view.cc
@@ -31,8 +31,8 @@ auto DispatchWeight(DeviceOrd device, RegTree const* tree) {
 
 ScalarTreeView::ScalarTreeView(Context const* ctx, RegTree const* tree)
     : CategoriesMixIn{tree->GetCategoriesMatrix(ctx->Device())},
-      nodes{tree->GetNodes().data()},
-      stats{tree->GetStats().data()},
+      nodes{tree->GetNodes(ctx->Device()).data()},
+      stats{tree->GetStats(ctx->Device()).data()},
       n{tree->NumNodes()} {
   CHECK(!tree->IsMultiTarget());
 }
diff --git a/src/tree/tree_view.h b/src/tree/tree_view.h
@@ -74,20 +74,22 @@ struct CategoriesMixIn {
   RegTree::CategoricalSplitMatrix cats;
 
   [[nodiscard]] XGBOOST_DEVICE bool HasCategoricalSplit() const { return !cats.categories.empty(); }
-  [[nodiscard]] XGBOOST_DEVICE RegTree::CategoricalSplitMatrix GetCategoriesMatrix() const {
+  [[nodiscard]] XGBOOST_DEVICE RegTree::CategoricalSplitMatrix const& GetCategoriesMatrix() const {
     return cats;
   }
   /**
    * @brief Get the bit storage of categories used by a node.
    */
-  [[nodiscard]] common::Span<uint32_t const> NodeCats(bst_node_t nidx) const {
+  [[nodiscard]] XGBOOST_DEVICE common::Span<uint32_t const> NodeCats(bst_node_t nidx) const {
     auto node_ptr = this->GetCategoriesMatrix().node_ptr;
     auto categories = this->GetCategoriesMatrix().categories;
     auto segment = node_ptr[nidx];
     auto node_cats = categories.subspan(segment.beg, segment.size);
     return node_cats;
   }
-  [[nodiscard]] FeatureType SplitType(bst_node_t nidx) const { return cats.split_type[nidx]; }
+  [[nodiscard]] XGBOOST_DEVICE FeatureType SplitType(bst_node_t nidx) const {
+    return cats.split_type[nidx];
+  }
 };
 
 /**
@@ -142,20 +144,20 @@ struct ScalarTreeView : public WalkTreeMixIn<ScalarTreeView>, public CategoriesM
   }
 
   [[nodiscard]] RTreeNodeStat const& Stat(bst_node_t nidx) const { return stats[nidx]; }
-  [[nodiscard]] auto SumHess(bst_node_t nidx) const { return stats[nidx].sum_hess; }
-  [[nodiscard]] auto LossChg(bst_node_t nidx) const { return stats[nidx].loss_chg; }
+  [[nodiscard]] XGBOOST_DEVICE auto SumHess(bst_node_t nidx) const { return stats[nidx].sum_hess; }
+  [[nodiscard]] XGBOOST_DEVICE auto LossChg(bst_node_t nidx) const { return stats[nidx].loss_chg; }
 
   XGBOOST_DEVICE explicit ScalarTreeView(RegTree::Node const* nodes, RTreeNodeStat const* stats,
                                          RegTree::CategoricalSplitMatrix cats, bst_node_t n_nodes)
       : CategoriesMixIn{std::move(cats)}, nodes{nodes}, stats{stats}, n{n_nodes} {}
 
-  /** @brief Create a device view, not implemented yet. */
+  /** @brief Create a device view */
   explicit ScalarTreeView(Context const* ctx, RegTree const* tree);
   /** @brief Create a host view */
   explicit ScalarTreeView(RegTree const* tree)
       : CategoriesMixIn{tree->GetCategoriesMatrix(DeviceOrd::CPU())},
-        nodes{tree->GetNodes().data()},
-        stats{tree->GetStats().data()},
+        nodes{tree->GetNodes(DeviceOrd::CPU()).data()},
+        stats{tree->GetStats(DeviceOrd::CPU()).data()},
         n{tree->NumNodes()} {
     CHECK(!tree->IsMultiTarget());
   }
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
@@ -669,7 +669,8 @@ struct GPUHistMakerDevice {
     // Use the nodes from tree, the leaf value might be changed by the objective since the
     // last update tree call.
     dh::CachingDeviceUVector<RegTree::Node> nodes;
-    dh::CopyTo(p_tree->GetNodes(), &nodes, this->ctx_->CUDACtx()->Stream());
+    // We can remove the CPU copy once we refactor the GPU hist to use the device tree.
+    dh::CopyTo(p_tree->GetNodes(DeviceOrd::CPU()), &nodes, this->ctx_->CUDACtx()->Stream());
     common::Span<RegTree::Node> d_nodes = dh::ToSpan(nodes);
     CHECK_EQ(out_preds_d.Shape(1), 1);
     dh::LaunchN(d_position.size(), ctx_->CUDACtx()->Stream(),
diff --git a/tests/cpp/helpers.h b/tests/cpp/helpers.h
@@ -445,6 +445,10 @@ class CudaArrayIterForTest : public ArrayIterForTest {
  public:
   explicit CudaArrayIterForTest(float sparsity, size_t rows = Rows(), size_t cols = Cols(),
                                 size_t batches = Batches());
+  explicit CudaArrayIterForTest(Context const* ctx, HostDeviceVector<float> const& data,
+                                std::size_t n_samples, bst_feature_t n_features,
+                                std::size_t n_batches)
+      : ArrayIterForTest{ctx, data, n_samples, n_features, n_batches} {};
   int Next() override;
   ~CudaArrayIterForTest() override = default;
 };
diff --git a/tests/cpp/plugin/test_sycl_hist_updater.cc b/tests/cpp/plugin/test_sycl_hist_updater.cc
@@ -500,7 +500,7 @@ void TestHistUpdaterExpandWithLossGuide(const xgboost::tree::TrainParam& param)
 
   updater.TestExpandWithLossGuide(gmat, p_fmat.get(), &tree, gpair);
 
-  const auto& nodes = tree.GetNodes();
+  const auto& nodes = tree.GetNodes(DeviceOrd::CPU());
   std::vector<float> ans(data.size());
   for (size_t data_idx = 0; data_idx < data.size(); ++data_idx) {
       size_t node_idx = 0;
@@ -544,7 +544,7 @@ void TestHistUpdaterExpandWithDepthWise(const xgboost::tree::TrainParam& param)
 
   updater.TestExpandWithDepthWise(gmat, p_fmat.get(), &tree, gpair);
 
-  const auto& nodes = tree.GetNodes();
+  const auto& nodes = tree.GetNodes(DeviceOrd::CPU());
   std::vector<float> ans(data.size());
   for (size_t data_idx = 0; data_idx < data.size(); ++data_idx) {
       size_t node_idx = 0;
diff --git a/tests/cpp/predictor/test_cpu_predictor.cc b/tests/cpp/predictor/test_cpu_predictor.cc
@@ -31,7 +31,7 @@ void CheckArrayLayout(const RegTree& tree, ArrayLayoutT buffer, int max_depth, i
   const auto& split_cond = buffer.SplitCond();
   const auto& default_left = buffer.DefaultLeft();
   const auto& nidx_in_tree = buffer.NidxInTree();
-  const auto& nodes = tree.GetNodes();
+  const auto& nodes = tree.GetNodes(DeviceOrd::CPU());
 
   if (depth == max_depth) {
     ASSERT_EQ(nidx_in_tree[nid_array - (1u << max_depth) + 1], nid);
diff --git a/tests/cpp/predictor/test_predictor.cc b/tests/cpp/predictor/test_predictor.cc
@@ -771,13 +771,22 @@ void TestVectorLeafPrediction(Context const *ctx) {
   };
   auto test_inplace = [&](float expected, HostDeviceVector<float> const*p_data) {
       PredictionCacheEntry predt_cache;
-      auto p_fmat = GetDMatrixFromData(p_data->ConstHostVector(), kRows, kCols);
+      std::shared_ptr<DMatrix> p_fmat = GetDMatrixFromData(p_data->ConstHostVector(), kRows, kCols);
       predictor->InitOutPredictions(p_fmat->Info(), &predt_cache.predictions, model);
+      if (ctx->IsCUDA()) {
+        // pull data to device.
+        p_data->SetDevice(ctx->Device());
+        p_data->ConstDeviceSpan();
+      }
       auto arr = GetArrayInterface(p_data, kRows, kCols);
       std::string str;
       Json::Dump(arr, &str);
       auto proxy = std::shared_ptr<DMatrix>(new data::DMatrixProxy{});
-      dynamic_cast<data::DMatrixProxy *>(proxy.get())->SetArray(str.data());
+      if (ctx->IsCUDA()) {
+        dynamic_cast<data::DMatrixProxy *>(proxy.get())->SetCudaArray(str.c_str());
+      } else {
+        dynamic_cast<data::DMatrixProxy *>(proxy.get())->SetArray(str.c_str());
+      }
       predictor->InplacePredict(proxy, model, std::numeric_limits<float>::quiet_NaN(), &predt_cache,
                                 0, 1);
       auto const &h_predt = predt_cache.predictions.HostVector();
@@ -797,11 +806,18 @@ void TestVectorLeafPrediction(Context const *ctx) {
 
     predictor->InitOutPredictions(p_fmat->Info(), &predt_cache.predictions, model);
 
-    auto iter = NumpyArrayIterForTest{ctx, *p_data, kRows, static_cast<bst_feature_t>(kCols),
-                                      static_cast<std::size_t>(1)};
+    std::unique_ptr<ArrayIterForTest> iter;
+    if (ctx->IsCUDA()) {
+      iter.reset(new CudaArrayIterForTest{ctx, *p_data, kRows, static_cast<bst_feature_t>(kCols),
+                                          static_cast<std::size_t>(1)});
+    } else {
+      iter.reset(new NumpyArrayIterForTest{ctx, *p_data, kRows, static_cast<bst_feature_t>(kCols),
+                                           static_cast<std::size_t>(1)});
+    }
+
     p_fmat = std::make_shared<data::IterativeDMatrix>(
-        &iter, iter.Proxy(), nullptr, Reset, Next, std::numeric_limits<float>::quiet_NaN(), 0, 256,
-        std::numeric_limits<std::int64_t>::max());
+        iter.get(), iter->Proxy(), nullptr, Reset, Next, std::numeric_limits<float>::quiet_NaN(), 0,
+        256, std::numeric_limits<std::int64_t>::max());
 
     predictor->InitOutPredictions(p_fmat->Info(), &predt_cache.predictions, model);
     predictor->PredictBatch(p_fmat.get(), &predt_cache, model, 0, 1);
@@ -817,18 +833,14 @@ void TestVectorLeafPrediction(Context const *ctx) {
   auto mt_tree = model.trees.front()->HostMtView();
   HostDeviceVector<float> data(kRows * kCols, mt_tree.SplitCond(RegTree::kRoot) + 1.0);
   test_batch(2.5, &data);
-  if (!ctx->IsCUDA()) {
-    test_inplace(2.5, &data);
-    test_ghist(2.5, &data);
-  }
+  test_inplace(2.5, &data);
+  test_ghist(2.5, &data);
 
   // go to left
   data.HostVector().assign(data.Size(), mt_tree.SplitCond(RegTree::kRoot) - 1.0);
   test_batch(1.5, &data);
-  if (!ctx->IsCUDA()) {
-    test_inplace(1.5, &data);
-    test_ghist(1.5, &data);
-  }
+  test_inplace(1.5, &data);
+  test_ghist(1.5, &data);
 }
 
 void ShapExternalMemoryTest::Run(Context const *ctx, bool is_qdm, bool is_interaction) {
diff --git a/tests/cpp/tree/test_tree_model.cc b/tests/cpp/tree/test_tree_model.cc
@@ -54,10 +54,10 @@ TEST(Tree, AllocateNode) {
                   /*left_sum=*/0.0f, /*right_sum=*/0.0f);
   ASSERT_EQ(tree.NumExtraNodes(), 2);
 
-  auto& nodes = tree.GetNodes();
-  ASSERT_FALSE(nodes.at(1).IsDeleted());
-  ASSERT_TRUE(nodes.at(1).IsLeaf());
-  ASSERT_TRUE(nodes.at(2).IsLeaf());
+  auto nodes = tree.GetNodes(DeviceOrd::CPU());
+  ASSERT_FALSE(nodes[1].IsDeleted());
+  ASSERT_TRUE(nodes[1].IsLeaf());
+  ASSERT_TRUE(nodes[2].IsLeaf());
 }
 
 TEST(Tree, ExpandCategoricalFeature) {
@@ -66,7 +66,7 @@ TEST(Tree, ExpandCategoricalFeature) {
     RegTree tree;
     tree.ExpandCategorical(0, 0, {}, true, 1.0, 2.0, 3.0, 11.0, 2.0,
                            /*left_sum=*/3.0, /*right_sum=*/4.0);
-    ASSERT_EQ(tree.GetNodes().size(), 3ul);
+    ASSERT_EQ(tree.Size(), 3ul);
     ASSERT_EQ(tree.GetNumLeaves(), 2);
     ASSERT_EQ(tree.GetSplitTypes(ctx.Device()).size(), 3ul);
     ASSERT_EQ(tree.GetSplitTypes(ctx.Device())[0], FeatureType::kCategorical);

Original file line number	Diff line number	Diff line change
`@@ -640,7 +640,7 @@ class ColumnSplitHelper {`
`640`	`640`	`tree_offsets_.resize(n_trees);`
`641`	`641`	`for (decltype(tree_begin) i = 0; i < n_trees; i++) {`
`642`	`642`	`auto const &tree = *model_.trees[tree_begin_ + i];`
`643`		`- tree_sizes_[i] = tree.GetNodes().size();`
	`643`	`+ tree_sizes_[i] = tree.Size();`
`644`	`644`	`}`
`645`	`645`	`// std::exclusive_scan (only available in c++17) equivalent to get tree offsets.`
`646`	`646`	`tree_offsets_[0] = 0;`