Unify max nodes. (#5497)

dmlc · Apr 10, 2020 · 866a477 · 866a477
1 parent bd653fa
commit 866a477
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 22 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,