Share the model view for CPU and GPU. (#11759)

- Extract the device model view and share it with the CPU.
- Ensure the model is locked when pulling tree data.
- Use the host device vector to store the tree group.

Author: trivialfis

plugin/sycl/predictor/predictor.cc

@@ -90,10 +90,11 @@ class DeviceModel {
 
     int num_group = model.learner_model_param->num_output_group;
     if (num_group > 1) {
-      tree_group.Resize(model.tree_info.size());
+      tree_group.Resize(model.tree_info.Size());
       auto& tree_group_host = tree_group.HostVector();
-      for (size_t tree_idx = 0; tree_idx < model.tree_info.size(); tree_idx++)
-        tree_group_host[tree_idx] = model.tree_info[tree_idx];
+      auto const& tree_group_in = model.tree_info.ConstHostVector();
+      for (size_t tree_idx = 0; tree_idx < tree_group_in.size(); tree_idx++)
+        tree_group_host[tree_idx] = tree_group_in[tree_idx];
     }
   }
 };

src/common/device_vector.cuh

@@ -467,7 +467,9 @@ class DeviceUVectorImpl {
     this->size_ = n;
     this->capacity_ = n;
 
-    std::swap(this->data_, new_ptr);
+    this->data_ = std::move(new_ptr);
+    // swap failed with CTK12.8
+    // std::swap(this->data_, new_ptr);
   }
   // Resize with init
   void resize(std::size_t n, T const &v) {  // NOLINT

src/data/cat_container.h

@@ -18,9 +18,10 @@
 #include "xgboost/base.h"                // for bst_cat_t
 #include "xgboost/data.h"                // for Entry
 #include "xgboost/host_device_vector.h"  // for HostDeviceVector
-#include "xgboost/json.h"                // for Json
 
 namespace xgboost {
+class Json;
+
 /**
  * @brief Error policy class used to interface with the encoder implementaion.
  */

src/gbm/gbtree.cc

@@ -145,6 +145,21 @@ void GBTree::Configure(Args const& cfg) {
   }
 }
 
+void GBTreeModel::InitTreesToUpdate() {
+  if (trees_to_update.empty()) {
+    for (auto& tree : trees) {
+      trees_to_update.push_back(std::move(tree));
+    }
+
+    trees.clear();
+    param.num_trees = 0;
+    tree_info.HostVector().clear();
+
+    iteration_indptr.clear();
+    iteration_indptr.push_back(0);
+  }
+}
+
 void GPUCopyGradient(Context const*, linalg::Matrix<GradientPair> const*, bst_group_t,
                      linalg::Matrix<GradientPair>*)
 #if defined(XGBOOST_USE_CUDA)
@@ -444,7 +459,9 @@ void GBTree::Slice(bst_layer_t begin, bst_layer_t end, bst_layer_t step, Gradien
 
   auto& out_indptr = out_model.iteration_indptr;
   TreesOneGroup& out_trees = out_model.trees;
-  std::vector<int32_t>& out_trees_info = out_model.tree_info;
+  auto& out_tree_info = out_model.tree_info.HostVector();
+
+  auto const& in_tree_info = this->model_.tree_info.ConstHostVector();
 
   bst_layer_t n_layers = (end - begin) / step;
   out_indptr.resize(n_layers + 1, 0);
@@ -462,8 +479,8 @@ void GBTree::Slice(bst_layer_t begin, bst_layer_t end, bst_layer_t step, Gradien
         std::unique_ptr<RegTree> new_tree{this->model_.trees.at(in_tree_idx)->Copy()};
         out_trees.emplace_back(std::move(new_tree));
 
-        bst_group_t group = this->model_.tree_info[in_tree_idx];
-        out_trees_info.push_back(group);
+        bst_group_t group = in_tree_info[in_tree_idx];
+        out_tree_info.push_back(group);
 
         out_model.iteration_indptr[out_l + 1]++;
       });
@@ -735,7 +752,7 @@ class Dart : public GBTree {
     auto layer_trees = [&]() {
       return model_.param.num_parallel_tree * model_.learner_model_param->OutputLength();
     };
-
+    auto const& h_tree_info = this->model_.tree_info.ConstHostVector();
     for (bst_tree_t i = tree_begin; i < tree_end; i += 1) {
       if (training && std::binary_search(idx_drop_.cbegin(), idx_drop_.cend(), i)) {
         continue;
@@ -749,7 +766,7 @@ class Dart : public GBTree {
 
       // Multiple the weight to output prediction.
       auto w = this->weight_drop_.at(i);
-      auto group = model_.tree_info.at(i);
+      auto group = h_tree_info.at(i);
       CHECK_EQ(p_out_preds->predictions.Size(), predts.predictions.Size());
 
       size_t n_rows = p_fmat->Info().num_row_;
@@ -815,6 +832,7 @@ class Dart : public GBTree {
       CHECK(success) << msg;
     };
 
+    auto const& h_tree_info = this->model_.tree_info.ConstHostVector();
     // Inplace predict is not used for training, so no need to drop tree.
     for (bst_tree_t i = tree_begin; i < tree_end; ++i) {
       predict_impl(i);
@@ -837,7 +855,7 @@ class Dart : public GBTree {
       }
       // Multiple the tree weight
       auto w = this->weight_drop_.at(i);
-      auto group = model_.tree_info.at(i);
+      auto group = h_tree_info.at(i);
       CHECK_EQ(predts.predictions.Size(), p_out_preds->predictions.Size());
 
       size_t n_rows = p_fmat->Info().num_row_;

src/gbm/gbtree_model.cc

@@ -19,7 +19,7 @@ namespace xgboost::gbm {
 namespace {
 // For creating the tree indptr from old models.
 void MakeIndptr(GBTreeModel* out_model) {
-  auto const& tree_info = out_model->tree_info;
+  auto const& tree_info = out_model->tree_info.ConstHostVector();
   if (tree_info.empty()) {
     return;
   }
@@ -41,7 +41,7 @@ void MakeIndptr(GBTreeModel* out_model) {
 // Validate the consistency of the model.
 void Validate(GBTreeModel const& model) {
   CHECK_EQ(model.trees.size(), model.param.num_trees);
-  CHECK_EQ(model.tree_info.size(), model.param.num_trees);
+  CHECK_EQ(model.tree_info.Size(), model.param.num_trees);
   // True even if the model is empty since we should always have 0 as the first element.
   CHECK_EQ(model.iteration_indptr.back(), model.param.num_trees);
 }
@@ -61,9 +61,10 @@ void GBTreeModel::SaveModel(Json* p_out) const {
     trees_json[t] = std::move(jtree);
   });
 
-  std::vector<Json> tree_info_json(tree_info.size());
-  for (size_t i = 0; i < tree_info.size(); ++i) {
-    tree_info_json[i] = Integer(tree_info[i]);
+  auto const& h_tree_info = tree_info.ConstHostVector();
+  std::vector<Json> tree_info_json(tree_info.Size());
+  for (size_t i = 0; i < h_tree_info.size(); ++i) {
+    tree_info_json[i] = Integer(h_tree_info[i]);
   }
 
   out["trees"] = Array(std::move(trees_json));
@@ -91,7 +92,8 @@ void GBTreeModel::LoadModel(Json const& in) {
 
   auto const& tree_info_json = get<Array const>(jmodel.at("tree_info"));
   CHECK_EQ(tree_info_json.size(), param.num_trees);
-  tree_info.resize(param.num_trees);
+  auto& h_tree_info = this->tree_info.HostVector();
+  h_tree_info.resize(param.num_trees);
 
   common::ParallelFor(param.num_trees, ctx_->Threads(), [&](auto t) {
     auto tree_id = get<Integer const>(trees_json[t]["id"]);
@@ -100,7 +102,7 @@ void GBTreeModel::LoadModel(Json const& in) {
   });
 
   for (bst_tree_t i = 0; i < param.num_trees; ++i) {
-    tree_info[i] = get<Integer const>(tree_info_json[i]);
+    h_tree_info[i] = get<Integer const>(tree_info_json[i]);
   }
 
   auto indptr_it = jmodel.find("iteration_indptr");
@@ -144,10 +146,16 @@ bst_tree_t GBTreeModel::CommitModel(TreesOneIter&& new_trees) {
 }
 
 void GBTreeModel::CommitModelGroup(TreesOneGroup&& new_trees, bst_target_t group_idx) {
+  auto& h_tree_info = this->tree_info.HostVector();
   for (auto& new_tree : new_trees) {
     trees.push_back(std::move(new_tree));
-    tree_info.push_back(group_idx);
+    h_tree_info.push_back(group_idx);
   }
   param.num_trees += static_cast<int>(new_trees.size());
 }
+
+common::Span<bst_target_t const> GBTreeModel::TreeGroups(DeviceOrd device) const {
+  return device.IsCPU() ? this->tree_info.ConstHostSpan()
+                        : (this->tree_info.SetDevice(device), this->tree_info.ConstDeviceSpan());
+}
 }  // namespace xgboost::gbm

src/gbm/gbtree_model.h

@@ -6,45 +6,42 @@
 #ifndef XGBOOST_GBM_GBTREE_MODEL_H_
 #define XGBOOST_GBM_GBTREE_MODEL_H_
 
-#include <dmlc/io.h>
 #include <dmlc/parameter.h>
-#include <xgboost/context.h>
-#include <xgboost/learner.h>
-#include <xgboost/model.h>
-#include <xgboost/parameter.h>
-#include <xgboost/tree_model.h>
 
 #include <memory>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "../common/threading_utils.h"
 #include "../data/cat_container.h"  // for CatContainer
+#include "xgboost/context.h"
+#include "xgboost/learner.h"
+#include "xgboost/model.h"
+#include "xgboost/tree_model.h"
 
 namespace xgboost {
 
 class Json;
 
 namespace gbm {
 /**
- * \brief Container for all trees built (not update) for one group.
+ * @brief Container for all trees built (not update) for one group.
  */
 using TreesOneGroup = std::vector<std::unique_ptr<RegTree>>;
 /**
- * \brief Container for all trees built (not update) for one iteration.
+ * @brief Container for all trees built (not update) for one iteration.
  */
 using TreesOneIter = std::vector<TreesOneGroup>;
 
-/*! \brief model parameters */
+/** @brief GBTree model parameters. */
 struct GBTreeModelParam : public dmlc::Parameter<GBTreeModelParam> {
  public:
   /**
-   * \brief number of trees
+   * @brief The number of trees.
    */
   std::int32_t num_trees{0};
   /**
-   * \brief Number of trees for a forest.
+   * @brief Number of trees for a single forest.
    */
   std::int32_t num_parallel_tree{1};
 
@@ -73,20 +70,8 @@ struct GBTreeModel : public Model {
       param.UpdateAllowUnknown(cfg);
     }
   }
-
-  void InitTreesToUpdate() {
-    if (trees_to_update.empty()) {
-      for (auto& tree : trees) {
-        trees_to_update.push_back(std::move(tree));
-      }
-      trees.clear();
-      param.num_trees = 0;
-      tree_info.clear();
-
-      iteration_indptr.clear();
-      iteration_indptr.push_back(0);
-    }
-  }
+  /** @brief Move existing trees into the update queue. */
+  void InitTreesToUpdate();
 
   void SaveModel(Json* p_out) const override;
   void LoadModel(Json const& p_out) override;
@@ -114,9 +99,9 @@ struct GBTreeModel : public Model {
     return static_cast<std::int32_t>(iteration_indptr.size() - 1);
   }
 
-  // base margin
+  /** @brief Global model properties. */
   LearnerModelParam const* learner_model_param;
-  // model parameter
+  /** @brief GBTree model parameters. */
   GBTreeModelParam param;
   /*! \brief vector of trees stored in the model */
   std::vector<std::unique_ptr<RegTree>> trees;
@@ -125,7 +110,7 @@ struct GBTreeModel : public Model {
   /**
    * @brief Group index for trees.
    */
-  std::vector<int> tree_info;
+  HostDeviceVector<bst_target_t> tree_info;
   /**
    * @brief Number of trees accumulated for each iteration.
    */
@@ -137,6 +122,10 @@ struct GBTreeModel : public Model {
   void Cats(std::shared_ptr<CatContainer> cats) { this->cats_ = cats; }
 
   auto const* Ctx() const { return this->ctx_; }
+  /**
+   * @brief Getter for the tree group index.
+   */
+  common::Span<bst_target_t const> TreeGroups(DeviceOrd device) const;
 
  private:
   /**