check _mm_prefetch in cmake (#2694)

* try to use _mm_prefetch anywhere

* refine

* fix bug

* remove the unneeded prefetch

CMakeLists.txt

@@ -115,6 +115,21 @@ if(USE_HDFS)
     SET(HDFS_CXX_LIBRARIES ${HDFS_LIB} ${JAVA_JVM_LIBRARY})
 endif(USE_HDFS)
 
+include(CheckCXXSourceCompiles)
+check_cxx_source_compiles("
+#include <xmmintrin.h>
+int main() {
+  int a = 0;
+  _mm_prefetch(&a, _MM_HINT_NTA);
+  return 0;
+}
+" MM_PREFETCH)
+
+if(${MM_PREFETCH})
+  message(STATUS "Use _mm_prefetch")
+  ADD_DEFINITIONS(-DMM_PREFETCH)
+endif()
+
 if(UNIX OR MINGW OR CYGWIN)
     SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -pthread -O3 -Wextra -Wall -Wno-ignored-attributes -Wno-unknown-pragmas -Wno-return-type")
     if(USE_SWIG)

include/LightGBM/meta.h

@@ -12,7 +12,7 @@
 #include <utility>
 #include <vector>
 
-#if (defined(_MSC_VER) && (defined(_M_IX86) || defined(_M_AMD64))) || defined(__INTEL_COMPILER)
+#if (defined(_MSC_VER) && (defined(_M_IX86) || defined(_M_AMD64))) || defined(__INTEL_COMPILER) || MM_PREFETCH
   #include <xmmintrin.h>
   #define PREFETCH_T0(addr) _mm_prefetch(reinterpret_cast<const char*>(addr), _MM_HINT_T0)
 #elif defined(__GNUC__)
@@ -71,6 +71,9 @@ typedef void(*AllgatherFunction)(char* input, comm_size_t input_size, const comm
 
 #define NO_SPECIFIC (-1)
 
+// Prefetch size is usually 64 bytes
+const int kCacheLineSize = 64;
+
 }  // namespace LightGBM
 
 #endif   // LightGBM_META_H_

src/io/dense_bin.hpp

@@ -71,12 +71,17 @@ class DenseBin: public Bin {
   void ConstructHistogram(const data_size_t* data_indices, data_size_t start, data_size_t end,
     const score_t* ordered_gradients, const score_t* ordered_hessians,
     HistogramBinEntry* out) const override {
-
-    const data_size_t prefetch_size = 32 / sizeof(VAL_T);
-    for (data_size_t i = start; i < end; i++) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + data_indices[i + prefetch_size]);
-      }
+    const data_size_t pf_offset = 64 / sizeof(VAL_T);
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize / sizeof(VAL_T);
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + data_indices[i + pf_offset]);
+      const VAL_T bin = data_[data_indices[i]];
+      out[bin].sum_gradients += ordered_gradients[i];
+      out[bin].sum_hessians += ordered_hessians[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const VAL_T bin = data_[data_indices[i]];
       out[bin].sum_gradients += ordered_gradients[i];
       out[bin].sum_hessians += ordered_hessians[i];
@@ -87,11 +92,17 @@ class DenseBin: public Bin {
   void ConstructHistogram(data_size_t start, data_size_t end,
     const score_t* ordered_gradients, const score_t* ordered_hessians,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32 / sizeof(VAL_T);
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + i + prefetch_size);
-      }
+    const data_size_t pf_offset = 64 / sizeof(VAL_T);
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize / sizeof(VAL_T);
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + i + pf_offset);
+      const VAL_T bin = data_[i];
+      out[bin].sum_gradients += ordered_gradients[i];
+      out[bin].sum_hessians += ordered_hessians[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const VAL_T bin = data_[i];
       out[bin].sum_gradients += ordered_gradients[i];
       out[bin].sum_hessians += ordered_hessians[i];
@@ -102,11 +113,16 @@ class DenseBin: public Bin {
   void ConstructHistogram(const data_size_t* data_indices, data_size_t start, data_size_t end,
     const score_t* ordered_gradients,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32 / sizeof(VAL_T);
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + data_indices[i + prefetch_size]);
-      }
+    const data_size_t pf_offset = 64 / sizeof(VAL_T);
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize / sizeof(VAL_T);
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + data_indices[i + pf_offset]);
+      const VAL_T bin = data_[data_indices[i]];
+      out[bin].sum_gradients += ordered_gradients[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const VAL_T bin = data_[data_indices[i]];
       out[bin].sum_gradients += ordered_gradients[i];
       ++out[bin].cnt;
@@ -116,11 +132,16 @@ class DenseBin: public Bin {
   void ConstructHistogram(data_size_t start, data_size_t end,
     const score_t* ordered_gradients,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32 / sizeof(VAL_T);
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + i + prefetch_size);
-      }
+    const data_size_t pf_offset = 64 / sizeof(VAL_T);
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize / sizeof(VAL_T);
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + i + pf_offset);
+      const VAL_T bin = data_[i];
+      out[bin].sum_gradients += ordered_gradients[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const VAL_T bin = data_[i];
       out[bin].sum_gradients += ordered_gradients[i];
       ++out[bin].cnt;

src/io/dense_nbits_bin.hpp

@@ -76,11 +76,18 @@ class Dense4bitsBin : public Bin {
   void ConstructHistogram(const data_size_t* data_indices, data_size_t start, data_size_t end,
     const score_t* ordered_gradients, const score_t* ordered_hessians,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32;
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + (data_indices[i + prefetch_size] >> 1));
-      }
+    const data_size_t pf_offset = 64;
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize;
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + (data_indices[i + pf_offset] >> 1));
+      const data_size_t idx = data_indices[i];
+      const auto bin = (data_[idx >> 1] >> ((idx & 1) << 2)) & 0xf;
+      out[bin].sum_gradients += ordered_gradients[i];
+      out[bin].sum_hessians += ordered_hessians[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const data_size_t idx = data_indices[i];
       const auto bin = (data_[idx >> 1] >> ((idx & 1) << 2)) & 0xf;
       out[bin].sum_gradients += ordered_gradients[i];
@@ -92,11 +99,17 @@ class Dense4bitsBin : public Bin {
   void ConstructHistogram(data_size_t start, data_size_t end,
     const score_t* ordered_gradients, const score_t* ordered_hessians,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32;
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + ((i + prefetch_size) >> 1));
-      }
+    const data_size_t pf_offset = 64;
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize;
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + ((i + pf_offset) >> 1));
+      const auto bin = (data_[i >> 1] >> ((i & 1) << 2)) & 0xf;
+      out[bin].sum_gradients += ordered_gradients[i];
+      out[bin].sum_hessians += ordered_hessians[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const auto bin = (data_[i >> 1] >> ((i & 1) << 2)) & 0xf;
       out[bin].sum_gradients += ordered_gradients[i];
       out[bin].sum_hessians += ordered_hessians[i];
@@ -107,11 +120,17 @@ class Dense4bitsBin : public Bin {
   void ConstructHistogram(const data_size_t* data_indices, data_size_t start, data_size_t end,
     const score_t* ordered_gradients,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32;
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + (data_indices[i + prefetch_size] >> 1));
-      }
+    const data_size_t pf_offset = 64;
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize;
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + (data_indices[i + pf_offset] >> 1));
+      const data_size_t idx = data_indices[i];
+      const auto bin = (data_[idx >> 1] >> ((idx & 1) << 2)) & 0xf;
+      out[bin].sum_gradients += ordered_gradients[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const data_size_t idx = data_indices[i];
       const auto bin = (data_[idx >> 1] >> ((idx & 1) << 2)) & 0xf;
       out[bin].sum_gradients += ordered_gradients[i];
@@ -122,11 +141,16 @@ class Dense4bitsBin : public Bin {
   void ConstructHistogram(data_size_t start, data_size_t end,
     const score_t* ordered_gradients,
     HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 32;
-    for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(data_.data() + ((i + prefetch_size) >> 1));
-      }
+    const data_size_t pf_offset = 64;
+    const data_size_t pf_end = end - pf_offset - kCacheLineSize;
+    data_size_t i = start;
+    for (; i < pf_end; i++) {
+      PREFETCH_T0(data_.data() + ((i + pf_offset) >> 1));
+      const auto bin = (data_[i >> 1] >> ((i & 1) << 2)) & 0xf;
+      out[bin].sum_gradients += ordered_gradients[i];
+      ++out[bin].cnt;
+    }
+    for (; i < end; i++) {
       const auto bin = (data_[i >> 1] >> ((i & 1) << 2)) & 0xf;
       out[bin].sum_gradients += ordered_gradients[i];
       ++out[bin].cnt;

src/io/ordered_sparse_bin.hpp

@@ -84,41 +84,30 @@ class OrderedSparseBin: public OrderedBin {
 
   void ConstructHistogram(int leaf, const score_t* gradient, const score_t* hessian,
                           HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 4;
     // get current leaf boundary
     const data_size_t start = leaf_start_[leaf];
     const data_size_t end = start + leaf_cnt_[leaf];
     for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(ordered_pair_.data() + i + prefetch_size);
-        PREFETCH_T0(gradient + ordered_pair_[i + prefetch_size].ridx);
-        PREFETCH_T0(hessian + ordered_pair_[i + prefetch_size].ridx);
-      }
-      const VAL_T bin0 = ordered_pair_[i].bin;
-      const auto g0 = gradient[ordered_pair_[i].ridx];
-      const auto h0 = hessian[ordered_pair_[i].ridx];
+      const VAL_T bin = ordered_pair_[i].bin;
+      const auto g = gradient[ordered_pair_[i].ridx];
+      const auto h = hessian[ordered_pair_[i].ridx];
 
-      out[bin0].sum_gradients += g0;
-      out[bin0].sum_hessians += h0;
-      ++out[bin0].cnt;
+      out[bin].sum_gradients += g;
+      out[bin].sum_hessians += h;
+      ++out[bin].cnt;
     }
   }
 
   void ConstructHistogram(int leaf, const score_t* gradient,
                           HistogramBinEntry* out) const override {
-    const data_size_t prefetch_size = 4;
     // get current leaf boundary
     const data_size_t start = leaf_start_[leaf];
     const data_size_t end = start + leaf_cnt_[leaf];
     for (data_size_t i = start; i < end; ++i) {
-      if (i + prefetch_size < end) {
-        PREFETCH_T0(ordered_pair_.data() + i + prefetch_size);
-        PREFETCH_T0(gradient + ordered_pair_[i + prefetch_size].ridx);
-      }
-      const VAL_T bin0 = ordered_pair_[i].bin;
-      const auto g0 = gradient[ordered_pair_[i].ridx];
-      out[bin0].sum_gradients += g0;
-      ++out[bin0].cnt;
+      const VAL_T bin = ordered_pair_[i].bin;
+      const auto g = gradient[ordered_pair_[i].ridx];
+      out[bin].sum_gradients += g;
+      ++out[bin].cnt;
     }
   }