Do not initialise data structures to maximum possible tree size. (#7919)

dmlc · May 19, 2022 · f6babc8 · f6babc8
1 parent 6f424d8
commit f6babc8
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/src/tree/split_evaluator.h b/src/tree/split_evaluator.h
@@ -13,6 +13,7 @@
 #include <utility>
 #include <vector>
 #include <limits>
+#include <algorithm>
 
 #include "xgboost/tree_model.h"
 #include "xgboost/host_device_vector.h"
@@ -49,8 +50,9 @@ class TreeEvaluator {
     } else {
       monotone_.HostVector() = p.monotone_constraints;
       monotone_.HostVector().resize(n_features, 0);
-      lower_bounds_.Resize(p.MaxNodes(), -std::numeric_limits<float>::max());
-      upper_bounds_.Resize(p.MaxNodes(), std::numeric_limits<float>::max());
+      // Initialised to some small size, can grow if needed
+      lower_bounds_.Resize(256, -std::numeric_limits<float>::max());
+      upper_bounds_.Resize(256, std::numeric_limits<float>::max());
       has_constraint_ = true;
     }
 
@@ -157,6 +159,15 @@ class TreeEvaluator {
     if (!has_constraint_) {
       return;
     }
+
+    auto max_nidx = std::max(leftid, rightid);
+    if (lower_bounds_.Size() <= max_nidx) {
+      lower_bounds_.Resize(max_nidx * 2 + 1, -std::numeric_limits<float>::max());
+    }
+    if (upper_bounds_.Size() <= max_nidx) {
+      upper_bounds_.Resize(max_nidx * 2 + 1, std::numeric_limits<float>::max());
+    }
+
     common::Transform<>::Init(
         [=] XGBOOST_DEVICE(size_t, common::Span<float> lower,
                            common::Span<float> upper,

diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
@@ -223,7 +223,7 @@ struct GPUHistMakerDevice {
       // Copy assigning an empty vector causes an exception in MSVC debug builds
       monotone_constraints = param.monotone_constraints;
     }
-    node_sum_gradients.resize(param.MaxNodes());
+    node_sum_gradients.resize(256);
 
     // Init histogram
     hist.Init(ctx_->gpu_id, page->Cuts().TotalBins());
@@ -614,12 +614,17 @@ struct GPUHistMakerDevice {
     }
     evaluator_.ApplyTreeSplit(candidate, p_tree);
 
-    node_sum_gradients[tree[candidate.nid].LeftChild()] = candidate.split.left_sum;
-    node_sum_gradients[tree[candidate.nid].RightChild()] = candidate.split.right_sum;
+    const auto& parent = tree[candidate.nid];
+    std::size_t max_nidx = std::max(parent.LeftChild(), parent.RightChild());
+    // Grow as needed
+    if (node_sum_gradients.size() <= max_nidx) {
+      node_sum_gradients.resize(max_nidx * 2 + 1);
+    }
+    node_sum_gradients[parent.LeftChild()] = candidate.split.left_sum;
+    node_sum_gradients[parent.RightChild()] = candidate.split.right_sum;
 
-    interaction_constraints.Split(candidate.nid, tree[candidate.nid].SplitIndex(),
-                                  tree[candidate.nid].LeftChild(),
-                                  tree[candidate.nid].RightChild());
+    interaction_constraints.Split(candidate.nid, parent.SplitIndex(), parent.LeftChild(),
+                                  parent.RightChild());
   }
 
   GPUExpandEntry InitRoot(RegTree* p_tree, dh::AllReducer* reducer) {