From 866a477319a47b529ace503924fa8136382f8ac9 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Fri, 10 Apr 2020 19:26:35 +0800
Subject: [PATCH] Unify max nodes. (#5497)

---
 src/tree/constraints.cu         | 11 +----------
 src/tree/param.h                | 14 ++++++++++++++
 src/tree/updater_gpu_common.cuh |  5 -----
 src/tree/updater_gpu_hist.cu    |  8 +-------
 4 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/src/tree/constraints.cu b/src/tree/constraints.cu
index 857dcfd55..371ee790c 100644
--- a/src/tree/constraints.cu
+++ b/src/tree/constraints.cu
@@ -53,16 +53,7 @@ void FeatureInteractionConstraint::Configure(
   }
 
   // --- Initialize allowed features attached to nodes.
-  if (param.max_depth == 0 && param.max_leaves == 0) {
-    LOG(FATAL) << "Max leaves and max depth cannot both be unconstrained for gpu_hist.";
-  }
-  int32_t n_nodes {0};
-  if (param.max_depth != 0) {
-    n_nodes = std::pow(2, param.max_depth + 1);
-  } else {
-    n_nodes = param.max_leaves * 2 - 1;
-  }
-  CHECK_NE(n_nodes, 0);
+  int32_t n_nodes { param.MaxNodes() };
   node_constraints_.resize(n_nodes);
   node_constraints_storage_.resize(n_nodes);
   for (auto& n : node_constraints_storage_) {
diff --git a/src/tree/param.h b/src/tree/param.h
index 35eeec111..f8c961cf7 100644
--- a/src/tree/param.h
+++ b/src/tree/param.h
@@ -230,6 +230,20 @@ struct TrainParam : public XGBoostParameter<TrainParam> {
     CHECK_GT(ret, 0U);
     return ret;
   }
+
+  bst_node_t MaxNodes() const {
+    if (this->max_depth == 0 && this->max_leaves == 0) {
+      LOG(FATAL) << "Max leaves and max depth cannot both be unconstrained.";
+    }
+    bst_node_t n_nodes{0};
+    if (this->max_leaves > 0) {
+      n_nodes = this->max_leaves * 2 - 1;
+    } else {
+      n_nodes = (1 << (this->max_depth + 1)) - 1;
+    }
+    CHECK_NE(n_nodes, 0);
+    return n_nodes;
+  }
 };
 
 /*! \brief Loss functions */
diff --git a/src/tree/updater_gpu_common.cuh b/src/tree/updater_gpu_common.cuh
index 87008f8da..63da94ada 100644
--- a/src/tree/updater_gpu_common.cuh
+++ b/src/tree/updater_gpu_common.cuh
@@ -178,10 +178,5 @@ struct SumCallbackOp {
     return old_prefix;
   }
 };
-
-// Total number of nodes in tree, given depth
-XGBOOST_DEVICE inline int MaxNodesDepth(int depth) {
-  return (1 << (depth + 1)) - 1;
-}
 }  // namespace tree
 }  // namespace xgboost
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 2f78730a4..043fe4d71 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -927,13 +927,7 @@ struct GPUHistMakerDevice {
 
 template <typename GradientSumT>
 inline void GPUHistMakerDevice<GradientSumT>::InitHistogram() {
-  CHECK(!(param.max_leaves == 0 && param.max_depth == 0))
-      << "Max leaves and max depth cannot both be unconstrained for "
-      "gpu_hist.";
-
-  int max_nodes =
-      param.max_leaves > 0 ? param.max_leaves * 2 : MaxNodesDepth(param.max_depth);
-
+  bst_node_t max_nodes { param.MaxNodes() };
   ba.Allocate(device_id,
               &prediction_cache, n_rows,
               &node_sum_gradients_d, max_nodes,