diff --git a/src/tree/split_evaluator.h b/src/tree/split_evaluator.h index 8cdf88834..ba3533e84 100644 --- a/src/tree/split_evaluator.h +++ b/src/tree/split_evaluator.h @@ -13,6 +13,7 @@ #include #include #include +#include #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::max()); - upper_bounds_.Resize(p.MaxNodes(), std::numeric_limits::max()); + // Initialised to some small size, can grow if needed + lower_bounds_.Resize(256, -std::numeric_limits::max()); + upper_bounds_.Resize(256, std::numeric_limits::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::max()); + } + if (upper_bounds_.Size() <= max_nidx) { + upper_bounds_.Resize(max_nidx * 2 + 1, std::numeric_limits::max()); + } + common::Transform<>::Init( [=] XGBOOST_DEVICE(size_t, common::Span lower, common::Span upper, diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu index 88978142e..ae209cdaf 100644 --- 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) {