reduce: add num threads recommend (#46)

bytedance · pleasantrabbit · Aug 13, 2020 · Feb 19, 2020 · Feb 19, 2020 · Feb 19, 2020
commit b4ecc6a01a39644c92e341c3191a2d28f45a8761
diff --git a/byteps/common/compressor/compressor.h b/byteps/common/compressor/compressor.h
@@ -18,7 +18,8 @@
 
 #include <memory>
 
-#include "../cpu_reducer.h"
+#include "../common.h"
+#include "../logging.h"
 #include "common.h"
 
 namespace byteps {
@@ -52,10 +53,7 @@ namespace compressor {
 class Compressor {
  public:
   Compressor(size_t size, DataType dtype)
-      : _size(size),
-        _dtype(dtype),
-        _buf(new byte_t[size]),
-        _cpu_reducer(new CpuReducer(nullptr)){};
+      : _size(size), _dtype(dtype), _buf(new byte_t[size]){};
   virtual ~Compressor() = default;
 
   /*!
@@ -126,9 +124,6 @@ class Compressor {
 
   /*! \brief buffer to store compressed grad */
   std::unique_ptr<byte_t[]> _buf;
-
-  /*! \brief CPU reducer */
-  std::unique_ptr<CpuReducer> _cpu_reducer;
 };
 
 }  // namespace compressor

diff --git a/byteps/common/compressor/error_feedback.h b/byteps/common/compressor/error_feedback.h
@@ -16,6 +16,7 @@
 #ifndef BYTEPS_COMPRESSOR_ERROR_FEEDBACK_H
 #define BYTEPS_COMPRESSOR_ERROR_FEEDBACK_H
 
+#include "../cpu_reducer.h"
 #include "compressor.h"
 
 namespace byteps {
@@ -51,6 +52,7 @@ class ErrorFeedback : public Compressor {
   ErrorFeedback(size_t size, DataType dtype, std::unique_ptr<Compressor> cptr)
       : Compressor(size, dtype),
         _error(new byte_t[size]()),
+        _cpu_reducer(new CpuReducer(nullptr)),
         _cptr(std::move(cptr)) {}
   virtual ~ErrorFeedback() = default;
 
@@ -85,6 +87,8 @@ class ErrorFeedback : public Compressor {
   /*! \brief buffer of error */
   std::unique_ptr<byte_t[]> _error;
 
+  std::unique_ptr<CpuReducer> _cpu_reducer;
+
  private:
   /*! \brief compressor pointer */
   std::unique_ptr<Compressor> _cptr;

diff --git a/byteps/common/compressor/impl/dithering.cc b/byteps/common/compressor/impl/dithering.cc
@@ -14,6 +14,7 @@
 // =============================================================================
 
 #include <cmath>
+#include <cstring>
 
 #include "../compressor_registry.h"
 #include "dithering.h"
@@ -54,12 +55,10 @@ tensor_t DitheringCompressor::CompressImpl(index_t* dst, const scalar_t* src,
   // normalize
   double scale = 0.0;
   if (_ntype == NomalizeType::MAX) {
-#pragma omp parallel for simd reduction(max : scale)
     for (size_t i = 0; i < len; i++) {
       scale = scale > std::abs(src[i]) ? scale : std::abs(src[i]);
     }
   } else if (_ntype == NomalizeType::L2) {
-#pragma omp parallel for simd num_threads(4) reduction(+ : scale)
     for (size_t i = 0; i < len; ++i) {
       scale += src[i] * src[i];
     }

diff --git a/byteps/common/compressor/impl/nesterov_momentum.cc b/byteps/common/compressor/impl/nesterov_momentum.cc
@@ -30,7 +30,7 @@ CompressorRegistry::Register reg(
       auto cptr = CompressorRegistry::Create(kwargs_clone, size, dtype);
       BPS_CHECK_NE(cptr, nullptr);
       // find \mu
-      auto mu = HyperParamFinder<float>(kwargs, "compressor_mu");
+      auto mu = HyperParamFinder<float>(kwargs, "momentum_mu");
       return std::unique_ptr<NesterovMomentumCompressor>(
           new NesterovMomentumCompressor(size, dtype, std::move(cptr), mu));
     });

diff --git a/byteps/common/compressor/impl/onebit.cc b/byteps/common/compressor/impl/onebit.cc
@@ -13,6 +13,8 @@
 // limitations under the License.
 // =============================================================================
 
+#include <cstring>
+
 #include "onebit.h"
 #include "../compressor_registry.h"
 
@@ -41,14 +43,13 @@ tensor_t OnebitCompressor::CompressImpl(index_t* dst, const scalar_t* src,
   float scale = 1.0f;
   if (_use_scale) {
     double sum = 0.0f;
-#pragma omp parallel for simd num_threads(4) reduction(+ : sum)
     for (size_t i = 0; i < len; ++i) {
       sum += std::abs(src[i]);
     }
     scale = sum / len;
   }
 
-#pragma omp parallel for simd num_threads(4)
+#pragma omp parallel for simd
   for (size_t i = 0; i < chunk_len; ++i) {
     index_t x = src[i * PACKING_SIZE] < 0;
     for (size_t j = 1; j < PACKING_SIZE; ++j) {
@@ -86,7 +87,7 @@ tensor_t OnebitCompressor::DecompressImpl(scalar_t* dst, const index_t* src,
     std::memcpy(ptr, src, compressed_size);
   }
 
-#pragma omp parallel for simd num_threads(4)
+#pragma omp parallel for simd
   for (int i = chunk_len - 1; i >= 0; --i) {
     index_t x = ptr[i];
     for (int j = PACKING_SIZE - 1; j >= 0; --j) {
@@ -119,7 +120,7 @@ void OnebitCompressor::FastUpdateErrorImpl(scalar_t* error, scalar_t* corrected,
   auto* pf = reinterpret_cast<const float*>(compressed + chunk_len);
   float scale = *pf;
 
-#pragma omp parallel for simd num_threads(4)
+#pragma omp parallel for simd
   for (int i = chunk_len - 1; i >= 0; --i) {
     index_t x = compressed[i];
     for (int j = PACKING_SIZE - 1; j >= 0; --j) {

diff --git a/byteps/common/compressor/impl/randomk.cc b/byteps/common/compressor/impl/randomk.cc
@@ -13,6 +13,8 @@
 // limitations under the License.
 // =============================================================================
 
+#include <cstring>
+
 #include "randomk.h"
 #include "../compressor_registry.h"
 

diff --git a/byteps/common/compressor/impl/topk.cc b/byteps/common/compressor/impl/topk.cc
@@ -13,6 +13,7 @@
 // limitations under the License.
 // =============================================================================
 
+#include <cstring>
 #include <queue>
 
 #include "../compressor_registry.h"

diff --git a/byteps/common/compressor/momentum.h b/byteps/common/compressor/momentum.h
@@ -16,6 +16,7 @@
 #ifndef BYTEPS_COMPRESSOR_MOMENTUM_H
 #define BYTEPS_COMPRESSOR_MOMENTUM_H
 
+#include "../cpu_reducer.h"
 #include "compressor.h"
 
 namespace byteps {
@@ -47,6 +48,7 @@ class Momentum : public Compressor {
       : Compressor(size, dtype),
         _mom(new byte_t[size]()),
         _mu(mu),
+        _cpu_reducer(new CpuReducer(nullptr)),
         _cptr(std::move(cptr)){};
   virtual ~Momentum() = default;
 
@@ -80,6 +82,8 @@ class Momentum : public Compressor {
   /*! \brief momentum factor */
   float _mu;
 
+  std::unique_ptr<CpuReducer> _cpu_reducer;
+
  private:
   /*! \brief compressor pointer */
   std::unique_ptr<Compressor> _cptr;

diff --git a/byteps/common/cpu_reducer.cc b/byteps/common/cpu_reducer.cc
@@ -38,12 +38,18 @@ CpuReducer::CpuReducer(std::shared_ptr<BytePSComm> comm) {
     _comm = nullptr;
   }
 #endif
-
   if (getenv("BYTEPS_OMP_THREAD_PER_GPU")) {
     _num_threads = atoi(getenv("BYTEPS_OMP_THREAD_PER_GPU"));
   } else {
     _num_threads = 4;
   }
+
+  if (getenv("BYTEPS_OMP_THREAD_THRESHOLD")) {
+    _single_thread_threshold = atoi(getenv("BYTEPS_OMP_THREAD_THRESHOLD"));
+  } else {
+    _single_thread_threshold = (1 << 16);
+  }
+
   return;
 }
 
@@ -86,7 +92,7 @@ int CpuReducer::sum(void* dst, const void* src, size_t len, DataType dtype) {
 
 template <typename T>
 int CpuReducer::_sum(T* dst, const T* src, size_t len) {
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < len / (size_t)sizeof(T); ++i) {
     dst[i] = dst[i] + src[i];
   }
@@ -101,7 +107,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src, size_t len) {
 
 #if __AVX__ && __F16C__
   if (is_avx_and_f16c()) {
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
     for (size_t i = 0; i < (size_t)(len / 8) * 8; i += 8) {
       // convert in & inout to m256
       __m256 in_m256 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i*)(in + i)));
@@ -126,7 +132,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src, size_t len) {
     Float2HalfBits(&inout_float, inout + i);
   }
 #else
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < (size_t)len; ++i) {
     float in_float;
     float inout_float;
@@ -177,7 +183,7 @@ int CpuReducer::sum(void* dst, const void* src1, const void* src2, size_t len,
 
 template <typename T>
 int CpuReducer::_sum(T* dst, const T* src1, const T* src2, size_t len) {
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < len / (size_t)sizeof(T); ++i) {
     dst[i] = src1[i] + src2[i];
   }
@@ -194,7 +200,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src1, const void* src2,
 
 #if __AVX__ && __F16C__
   if (is_avx_and_f16c()) {
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
     for (size_t i = 0; i < (size_t)(len / 8) * 8; i += 8) {
       // convert in1 & in2 to m256
       __m256 in_m256 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i*)(in1 + i)));
@@ -219,7 +225,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src1, const void* src2,
     Float2HalfBits(&out_float, out + i);
   }
 #else
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < (size_t)len; ++i) {
     float in1_float;
     float in2_float;
@@ -264,7 +270,7 @@ int CpuReducer::sum(void* dst, const void* src, size_t len, DataType dtype,
 
 template <typename T>
 int CpuReducer::_sum(T* dst, const T* src, size_t len, float alpha) {
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < len / (size_t)sizeof(T); ++i) {
     dst[i] = dst[i] + alpha * src[i];
   }
@@ -284,7 +290,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src, size_t len,
 
   if (is_avx_and_f16c()) {
     __m256 __mm256_alpha = _mm256_loadu_ps(mm256_alpha);
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
     for (size_t i = 0; i < (size_t)(len / 8) * 8; i += 8) {
       // convert in & inout to m256
       __m256 in_m256 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i*)(in + i)));
@@ -310,7 +316,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src, size_t len,
     Float2HalfBits(&inout_float, inout + i);
   }
 #else
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < (size_t)len; ++i) {
     float in_float;
     float inout_float;
@@ -362,7 +368,7 @@ int CpuReducer::sum(void* dst, const void* src1, const void* src2, size_t len,
 template <typename T>
 int CpuReducer::_sum(T* dst, const T* src1, const T* src2, size_t len,
                      float alpha) {
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < len / (size_t)sizeof(T); ++i) {
     dst[i] = src1[i] + alpha * src2[i];
   }
@@ -383,7 +389,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src1, const void* src2,
 
   if (is_avx_and_f16c()) {
     __m256 __mm256_alpha = _mm256_loadu_ps(mm256_alpha);
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
     for (size_t i = 0; i < (size_t)(len / 8) * 8; i += 8) {
       // convert in1 & in2 to m256
       __m256 in1_m256 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i*)(in1 + i)));
@@ -409,7 +415,7 @@ int CpuReducer::_sum_float16(void* dst, const void* src1, const void* src2,
     Float2HalfBits(&out_float, out + i);
   }
 #else
-#pragma omp parallel for simd num_threads(_num_threads)
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
   for (size_t i = 0; i < (size_t)len; ++i) {
     float in1_float;
     float in2_float;
@@ -424,8 +430,16 @@ int CpuReducer::_sum_float16(void* dst, const void* src1, const void* src2,
 }
 
 int CpuReducer::copy(void* dst, const void* src, size_t len) {
-  std::memcpy(dst, src, len);
+  auto in = (float*)src;
+  auto out = (float*)dst;
+#pragma omp parallel for simd num_threads(GetRecommendNumThreads(len))
+  for (size_t i = 0; i < len / 4; ++i) {
+    out[i] = in[i];
+  }
+  if (len % 4) {
+    std::memcpy(out + len / 4, in + len / 4, len % 4);
+  }
+  return 0;
 }
-
 }  // namespace common
 }  // namespace byteps
diff --git a/byteps/common/cpu_reducer.h b/byteps/common/cpu_reducer.h
@@ -63,6 +63,14 @@ class CpuReducer {
   DataType GetDataType(int dtype) { return static_cast<DataType>(dtype); }
 
  private:
+  size_t GetRecommendNumThreads(size_t len) const {
+    if (len < _single_thread_threshold) {
+      return 1;
+    } else {
+      return _num_threads;
+    }
+  }
+
 #if __AVX__ && __F16C__
   // Query CPUID to determine AVX and F16C runtime support.
   bool is_avx_and_f16c() {
@@ -200,6 +208,7 @@ class CpuReducer {
 
   std::shared_ptr<BytePSComm> _comm;
   int _num_threads;
+  size_t _single_thread_threshold; 
 };
 
 }  // namespace common