Serialize expand entry for allgather. (#9702)

src/tree/gpu_hist/expand_entry.cuh

@@ -1,31 +1,36 @@
-/*!
- * Copyright 2020 by XGBoost Contributors
+/**
+ * Copyright 2020-2023, XGBoost Contributors
  */
 #ifndef EXPAND_ENTRY_CUH_
 #define EXPAND_ENTRY_CUH_
-#include <xgboost/span.h>
+
+#include <limits>   // for numeric_limits
+#include <utility>  // for move
 
 #include "../param.h"
 #include "../updater_gpu_common.cuh"
+#include "xgboost/base.h"  // for bst_node_t
 
-namespace xgboost {
-namespace tree {
-
+namespace xgboost::tree {
 struct GPUExpandEntry {
-  int nid;
-  int depth;
+  bst_node_t nid;
+  bst_node_t depth;
   DeviceSplitCandidate split;
 
-  float base_weight { std::numeric_limits<float>::quiet_NaN() };
-  float left_weight { std::numeric_limits<float>::quiet_NaN() };
-  float right_weight { std::numeric_limits<float>::quiet_NaN() };
+  float base_weight{std::numeric_limits<float>::quiet_NaN()};
+  float left_weight{std::numeric_limits<float>::quiet_NaN()};
+  float right_weight{std::numeric_limits<float>::quiet_NaN()};
 
   GPUExpandEntry() = default;
-  XGBOOST_DEVICE GPUExpandEntry(int nid, int depth, DeviceSplitCandidate split,
-                             float base, float left, float right)
-      : nid(nid), depth(depth), split(std::move(split)), base_weight{base},
-        left_weight{left}, right_weight{right} {}
-  bool IsValid(const TrainParam& param, int num_leaves) const {
+  XGBOOST_DEVICE GPUExpandEntry(bst_node_t nid, bst_node_t depth, DeviceSplitCandidate split,
+                                float base, float left, float right)
+      : nid(nid),
+        depth(depth),
+        split(std::move(split)),
+        base_weight{base},
+        left_weight{left},
+        right_weight{right} {}
+  [[nodiscard]] bool IsValid(TrainParam const& param, bst_node_t num_leaves) const {
     if (split.loss_chg <= kRtEps) return false;
     if (split.left_sum.GetQuantisedHess() == 0 || split.right_sum.GetQuantisedHess() == 0) {
       return false;
@@ -42,17 +47,11 @@ struct GPUExpandEntry {
     return true;
   }
 
-  bst_float GetLossChange() const {
-    return split.loss_chg;
-  }
+  [[nodiscard]] float GetLossChange() const { return split.loss_chg; }
 
-  int GetNodeId() const {
-    return nid;
-  }
+  [[nodiscard]] bst_node_t GetNodeId() const { return nid; }
 
-  int GetDepth() const {
-    return depth;
-  }
+  [[nodiscard]] bst_node_t GetDepth() const { return depth; }
 
   friend std::ostream& operator<<(std::ostream& os, const GPUExpandEntry& e) {
     os << "GPUExpandEntry: \n";
@@ -63,9 +62,69 @@ struct GPUExpandEntry {
     os << "right_sum: " << e.split.right_sum << "\n";
     return os;
   }
-};
 
-}  // namespace tree
-}  // namespace xgboost
+  void Save(Json* p_out) const {
+    auto& out = *p_out;
+
+    out["nid"] = Integer{this->nid};
+    out["depth"] = Integer{this->depth};
+    // GPU specific
+    out["base_weight"] = this->base_weight;
+    out["left_weight"] = this->left_weight;
+    out["right_weight"] = this->right_weight;
+
+    /**
+     * Handle split
+     */
+    out["split"] = Object{};
+    auto& split = out["split"];
+    split["loss_chg"] = this->split.loss_chg;
+    split["sindex"] = Integer{this->split.findex};
+    split["split_value"] = this->split.fvalue;
+
+    // cat
+    split["thresh"] = Integer{this->split.thresh};
+    split["is_cat"] = Boolean{this->split.is_cat};
+    /**
+     * Gradients
+     */
+    auto save = [&](std::string const& name, GradientPairInt64 const& sum) {
+      out[name] = I64Array{2};
+      auto& array = get<I64Array>(out[name]);
+      array[0] = sum.GetQuantisedGrad();
+      array[1] = sum.GetQuantisedHess();
+    };
+    save("left_sum", this->split.left_sum);
+    save("right_sum", this->split.right_sum);
+  }
+
+  void Load(Json const& in) {
+    this->nid = get<Integer const>(in["nid"]);
+    this->depth = get<Integer const>(in["depth"]);
+    // GPU specific
+    this->base_weight = get<Number const>(in["base_weight"]);
+    this->left_weight = get<Number const>(in["left_weight"]);
+    this->right_weight = get<Number const>(in["right_weight"]);
+
+    /**
+     * Handle split
+     */
+    auto const& split = in["split"];
+    this->split.loss_chg = get<Number const>(split["loss_chg"]);
+    this->split.findex = get<Integer const>(split["sindex"]);
+    this->split.fvalue = get<Number const>(split["split_value"]);
+    // cat
+    this->split.thresh = get<Integer const>(split["thresh"]);
+    this->split.is_cat = get<Boolean const>(split["is_cat"]);
+    /**
+     * Gradients
+     */
+    auto const& left_sum = get<I64Array const>(in["left_sum"]);
+    this->split.left_sum = GradientPairInt64{left_sum[0], left_sum[1]};
+    auto const& right_sum = get<I64Array const>(in["right_sum"]);
+    this->split.right_sum = GradientPairInt64{right_sum[0], right_sum[1]};
+  }
+};
+}  // namespace xgboost::tree
 
 #endif  // EXPAND_ENTRY_CUH_