speed up multi-threading sum (#3485)

* speed up multi-threading sum

* Apply suggestions from code review

include/LightGBM/utils/threading.h

@@ -71,7 +71,7 @@ class Threading {
       const std::function<void(int, INDEX_T, INDEX_T)>& inner_fun) {
     int n_block = 1;
     INDEX_T num_inner = end - start;
-    BlockInfo<INDEX_T>(end - start, min_block_size, &n_block, &num_inner);
+    BlockInfo<INDEX_T>(num_inner, min_block_size, &n_block, &num_inner);
     OMP_INIT_EX();
 #pragma omp parallel for schedule(static, 1)
     for (int i = 0; i < n_block; ++i) {
@@ -93,20 +93,16 @@ class Threading {
       VAL1_T* res1, VAL2_T* res2) {
     int n_block = 1;
     INDEX_T num_inner = end - start;
-    BlockInfoForceSize<INDEX_T>(end - start, min_block_size, &n_block,
+    BlockInfoForceSize<INDEX_T>(num_inner, min_block_size, &n_block,
                                 &num_inner);
     std::vector<VAL1_T> val_1s(n_block, static_cast<VAL1_T>(0));
     std::vector<VAL2_T> val_2s(n_block, static_cast<VAL2_T>(0));
-    OMP_INIT_EX();
-#pragma omp parallel for schedule(static, 1)
+#pragma omp parallel for schedule(static)
     for (int i = 0; i < n_block; ++i) {
-      OMP_LOOP_EX_BEGIN();
       INDEX_T inner_start = start + num_inner * i;
       INDEX_T inner_end = std::min(end, inner_start + num_inner);
       inner_fun(i, inner_start, inner_end, &val_1s[i], &val_2s[i]);
-      OMP_LOOP_EX_END();
     }
-    OMP_THROW_EX();
     *res1 = 0;
     *res2 = 0;
     for (int i = 0; i < n_block; ++i) {

src/treelearner/leaf_splits.hpp

@@ -60,47 +60,65 @@ class LeafSplits {
   }
 
   /*!
-  * \brief Init splits on current leaf, it will traverse all data to sum up the results
-  * \param gradients
-  * \param hessians
-  */
+   * \brief Init splits on the current leaf, it will traverse all data to sum up the results
+   * \param gradients
+   * \param hessians
+   */
   void Init(const score_t* gradients, const score_t* hessians) {
     num_data_in_leaf_ = num_data_;
     leaf_index_ = 0;
     data_indices_ = nullptr;
-    Threading::SumReduction<data_size_t, double, double>(
-        0, num_data_in_leaf_, 2048,
-        [=](int, data_size_t start, data_size_t end, double* s1, double* s2) {
-          *s1 = *s2 = 0;
-          for (data_size_t i = start; i < end; ++i) {
-            *s1 += gradients[i];
-            *s2 += hessians[i];
-          }
-        },
-        &sum_gradients_, &sum_hessians_);
+    if (num_data_in_leaf_ < 4096) {
+      sum_gradients_ = 0.0f;
+      sum_hessians_ = 0.0f;
+      for (data_size_t i = 0; i < num_data_in_leaf_; ++i) {
+        sum_gradients_ += gradients[i];
+        sum_hessians_ += hessians[i];
+      }
+    } else {
+      Threading::SumReduction<data_size_t, double, double>(
+          0, num_data_in_leaf_, 2048,
+          [=](int, data_size_t start, data_size_t end, double* s1, double* s2) {
+            for (data_size_t i = start; i < end; ++i) {
+              *s1 += gradients[i];
+              *s2 += hessians[i];
+            }
+          },
+          &sum_gradients_, &sum_hessians_);
+    }
   }
 
   /*!
-  * \brief Init splits on current leaf of partial data.
-  * \param leaf Index of current leaf
-  * \param data_partition current data partition
-  * \param gradients
-  * \param hessians
-  */
-  void Init(int leaf, const DataPartition* data_partition, const score_t* gradients, const score_t* hessians) {
+   * \brief Init splits on current leaf of partial data.
+   * \param leaf Index of current leaf
+   * \param data_partition current data partition
+   * \param gradients
+   * \param hessians
+   */
+  void Init(int leaf, const DataPartition* data_partition,
+            const score_t* gradients, const score_t* hessians) {
     leaf_index_ = leaf;
     data_indices_ = data_partition->GetIndexOnLeaf(leaf, &num_data_in_leaf_);
-    Threading::SumReduction<data_size_t, double, double>(
-        0, num_data_in_leaf_, 2048,
-        [=](int, data_size_t start, data_size_t end, double* s1, double* s2) {
-          *s1 = *s2 = 0;
-          for (data_size_t i = start; i < end; ++i) {
-            data_size_t idx = data_indices_[i];
-            *s1 += gradients[idx];
-            *s2 += hessians[idx];
-          }
-        },
-        &sum_gradients_, &sum_hessians_);
+    if (num_data_in_leaf_ < 4096) {
+      sum_gradients_ = 0.0f;
+      sum_hessians_ = 0.0f;
+      for (data_size_t i = 0; i < num_data_in_leaf_; ++i) {
+        const data_size_t idx = data_indices_[i];
+        sum_gradients_ += gradients[idx];
+        sum_hessians_ += hessians[idx];
+      }
+    } else {
+      Threading::SumReduction<data_size_t, double, double>(
+          0, num_data_in_leaf_, 2048,
+          [=](int, data_size_t start, data_size_t end, double* s1, double* s2) {
+            for (data_size_t i = start; i < end; ++i) {
+              const data_size_t idx = data_indices_[i];
+              *s1 += gradients[idx];
+              *s2 += hessians[idx];
+            }
+          },
+          &sum_gradients_, &sum_hessians_);
+    }
   }