updated based on comments

microsoft · May 23, 2020 · 7e34e36 · 7e34e36
1 parent 64abf8c
commit 7e34e36
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/src/treelearner/data_parallel_tree_learner.cpp b/src/treelearner/data_parallel_tree_learner.cpp
@@ -29,7 +29,14 @@ void DataParallelTreeLearner<TREELEARNER_T>::Init(const Dataset* train_data, boo
   // allocate buffer for communication
   size_t buffer_size = this->train_data_->NumTotalBin() * kHistEntrySize;
 
-  input_buffer_.resize(buffer_size);
+  auto max_cat_threshold = this->config_->max_cat_threshold;
+  // input_buffer_ needs to be able to hold smaller and larger best splits in SyncUpGlobalBestSplit
+  int splitInfoSize = SplitInfo::Size(max_cat_threshold) * 2;
+  if (buffer_size < splitInfoSize) {
+    input_buffer_.resize(splitInfoSize);
+  } else {
+    input_buffer_.resize(buffer_size);
+  }
   output_buffer_.resize(buffer_size);
 
   is_feature_aggregated_.resize(this->num_features_);
@@ -231,11 +238,6 @@ void DataParallelTreeLearner<TREELEARNER_T>::FindBestSplitsFromHistograms(const
   }
 
   // sync global best info
-  auto max_cat_threshold = this->config_->max_cat_threshold;
-  int splitInfoSize = SplitInfo::Size(max_cat_threshold);
-  if (input_buffer_.size() < splitInfoSize * 2) {
-    input_buffer_.resize(splitInfoSize * 2);
-  }
   SyncUpGlobalBestSplit(input_buffer_.data(), input_buffer_.data(), &smaller_best_split, &larger_best_split, this->config_->max_cat_threshold);
 
   // set best split

diff --git a/src/treelearner/split_info.hpp b/src/treelearner/split_info.hpp
@@ -49,7 +49,7 @@ struct SplitInfo {
   bool default_left = true;
   int8_t monotone_type = 0;
   inline static int Size(int max_cat_threshold) {
-    return 2 * sizeof(int) + sizeof(uint32_t) + sizeof(bool) + sizeof(double) * 9 + sizeof(data_size_t) * 2 + max_cat_threshold * sizeof(uint32_t) + sizeof(int8_t);
+    return 2 * sizeof(int) + sizeof(uint32_t) + sizeof(bool) + sizeof(double) * 7 + sizeof(data_size_t) * 2 + max_cat_threshold * sizeof(uint32_t) + sizeof(int8_t);
   }
 
   inline void CopyTo(char* buffer) const {