ReduceLayer && Pooling Layer && parallelfor (#268)

Author: AndreySorokin7

include/layers/PoolingLayer.hpp

@@ -2,15 +2,15 @@
 #include <algorithm>
 #include <cmath>
 #include <cstdlib>
+#include <limits>
 #include <numeric>
 #include <stdexcept>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "layers/Layer.hpp"
-#include "tbb/blocked_range2d.h"
-#include "tbb/parallel_for.h"
+#include "parallel/parallel.hpp"
 
 namespace it_lab_ai {
 
@@ -64,7 +64,6 @@ class PoolingLayer : public Layer {
   Shape dilations_;
   bool ceil_mode_;
   std::string poolingType_;
-  ImplType implType_;
 };
 
 inline size_t coord_size(size_t coord, const Shape& shape) {
@@ -98,12 +97,13 @@ class PoolingLayerImpl : public LayerImpl<ValueType> {
   PoolingLayerImpl(const Shape& input_shape, const Shape& pooling_shape,
                    const std::string& pooling_type = "average")
       : PoolingLayerImpl(input_shape, pooling_shape, {2, 2}, {0, 0, 0, 0},
-                         {1, 1}, false, pooling_type) {}
+                         {1, 1}, false, pooling_type, ParBackend::kSeq) {}
   PoolingLayerImpl(const Shape& input_shape, const Shape& pooling_shape,
                    const Shape& strides = {2, 2},
                    const Shape& pads = {0, 0, 0, 0},
                    const Shape& dilations = {1, 1}, bool ceil_mode = false,
-                   const std::string& pooling_type = "average");
+                   const std::string& pooling_type = "average",
+                   ParBackend parallel_backend = ParBackend::kSeq);
   PoolingLayerImpl(const PoolingLayerImpl& c) = default;
   PoolingLayerImpl& operator=(const PoolingLayerImpl& c) = default;
   std::vector<ValueType> run(
@@ -116,20 +116,22 @@ class PoolingLayerImpl : public LayerImpl<ValueType> {
   Shape dilations_;
   bool ceil_mode_;
   PoolingType poolingType_;
+  ParBackend parallel_backend_;
 };
 
 template <typename ValueType>
 PoolingLayerImpl<ValueType>::PoolingLayerImpl(
     const Shape& input_shape, const Shape& pooling_shape, const Shape& strides,
     const Shape& pads, const Shape& dilations, bool ceil_mode,
-    const std::string& pooling_type)
+    const std::string& pooling_type, ParBackend parallel_backend)
     : LayerImpl<ValueType>(input_shape, input_shape),
       poolingShape_(pooling_shape),
       strides_(strides),
       pads_(pads),
       dilations_(dilations),
       ceil_mode_(ceil_mode),
-      poolingType_(kAverage) {
+      poolingType_(kAverage),
+      parallel_backend_(parallel_backend) {
   if (pooling_shape[0] == 0 && pooling_shape[1] == 0) {
     poolingShape_ = Shape({input_shape[input_shape.dims() - 2],
                            input_shape[input_shape.dims() - 1]});
@@ -208,221 +210,113 @@ std::vector<ValueType> PoolingLayerImpl<ValueType>::run(
   int batch_dim = this->inputShape_.dims() > spatial_dims ? 0 : -1;
   int channel_dim = this->inputShape_.dims() > spatial_dims + 1 ? 1 : -1;
 
-  for (size_t n = 0; n < (batch_dim >= 0 ? this->outputShape_[batch_dim] : 1);
-       n++) {
-    for (size_t c = 0;
-         c < (channel_dim >= 0 ? this->outputShape_[channel_dim] : 1); c++) {
-      for (size_t h = 0;
-           h < this->outputShape_[this->outputShape_.dims() - spatial_dims];
-           h++) {
-        for (size_t w = 0;
-             w < (spatial_dims > 1
-                      ? this->outputShape_[this->outputShape_.dims() -
-                                           spatial_dims + 1]
-                      : 1);
-             w++) {
-          std::vector<ValueType> pooling_buf;
-
-          int start_h =
-              static_cast<int>(h * strides_[0]) - static_cast<int>(pads_[0]);
-          int start_w = spatial_dims > 1 ? static_cast<int>(w * strides_[1]) -
-                                               static_cast<int>(pads_[2])
-                                         : 0;
-
-          for (size_t kh = 0; kh < poolingShape_[0]; kh++) {
-            for (size_t kw = 0; kw < (spatial_dims > 1 ? poolingShape_[1] : 1);
-                 kw++) {
-              int pos_h = start_h + static_cast<int>(kh * dilations_[0]);
-              int pos_w = spatial_dims > 1
-                              ? start_w + static_cast<int>(kw * dilations_[1])
-                              : 0;
-
-              if (pos_h >= 0 &&
-                  pos_h < static_cast<int>(
-                              this->inputShape_[this->inputShape_.dims() -
-                                                spatial_dims]) &&
-                  (spatial_dims <= 1 ||
-                   (pos_w >= 0 &&
-                    pos_w < static_cast<int>(
-                                this->inputShape_[this->inputShape_.dims() -
-                                                  spatial_dims + 1])))) {
-                std::vector<size_t> input_coords(this->inputShape_.dims(), 0);
-                if (batch_dim >= 0) input_coords[batch_dim] = n;
-                if (channel_dim >= 0) input_coords[channel_dim] = c;
-                input_coords[this->inputShape_.dims() - spatial_dims] = pos_h;
-                if (spatial_dims > 1) {
-                  input_coords[this->inputShape_.dims() - spatial_dims + 1] =
-                      pos_w;
-                }
+  size_t out_h = this->outputShape_[this->outputShape_.dims() - spatial_dims];
+  size_t out_w =
+      spatial_dims > 1
+          ? this->outputShape_[this->outputShape_.dims() - spatial_dims + 1]
+          : 1;
+  size_t out_n = batch_dim >= 0 ? this->outputShape_[batch_dim] : 1;
+  size_t out_c = channel_dim >= 0 ? this->outputShape_[channel_dim] : 1;
+
+  size_t total_work = out_n * out_c * out_h * out_w;
+
+  parallel::Options options;
+  options.backend = parallel_backend_;
+
+  parallel::parallel_for(
+      total_work,
+      [&](size_t idx) {
+        size_t tmp = idx;
+        size_t w = 0;
+        if (spatial_dims > 1) {
+          w = tmp % out_w;
+          tmp /= out_w;
+        }
+        size_t h = tmp % out_h;
+        tmp /= out_h;
 
-                size_t input_index = this->inputShape_.get_index(input_coords);
-                pooling_buf.push_back(input[input_index]);
+        size_t c = 0;
+        if (channel_dim >= 0) {
+          c = tmp % out_c;
+          tmp /= out_c;
+        }
+        size_t n = batch_dim >= 0 ? tmp : 0;
+
+        int start_h =
+            static_cast<int>(h * strides_[0]) - static_cast<int>(pads_[0]);
+        int start_w = spatial_dims > 1 ? static_cast<int>(w * strides_[1]) -
+                                             static_cast<int>(pads_[2])
+                                       : 0;
+
+        auto sum = ValueType(0);
+        ValueType max_val = std::numeric_limits<ValueType>::lowest();
+        size_t count = 0;
+
+        for (size_t kh = 0; kh < poolingShape_[0]; kh++) {
+          for (size_t kw = 0; kw < (spatial_dims > 1 ? poolingShape_[1] : 1);
+               kw++) {
+            int pos_h = start_h + static_cast<int>(kh * dilations_[0]);
+            int pos_w = spatial_dims > 1
+                            ? start_w + static_cast<int>(kw * dilations_[1])
+                            : 0;
+
+            if (pos_h >= 0 &&
+                pos_h < static_cast<int>(
+                            this->inputShape_[this->inputShape_.dims() -
+                                              spatial_dims]) &&
+                (spatial_dims <= 1 ||
+                 (pos_w >= 0 &&
+                  pos_w < static_cast<int>(
+                              this->inputShape_[this->inputShape_.dims() -
+                                                spatial_dims + 1])))) {
+              std::vector<size_t> input_coords(this->inputShape_.dims(), 0);
+              if (batch_dim >= 0) input_coords[batch_dim] = n;
+              if (channel_dim >= 0) input_coords[channel_dim] = c;
+              input_coords[this->inputShape_.dims() - spatial_dims] = pos_h;
+              if (spatial_dims > 1) {
+                input_coords[this->inputShape_.dims() - spatial_dims + 1] =
+                    pos_w;
               }
-            }
-          }
 
-          std::vector<size_t> output_coords(this->outputShape_.dims(), 0);
-          if (batch_dim >= 0) output_coords[batch_dim] = n;
-          if (channel_dim >= 0) output_coords[channel_dim] = c;
-          output_coords[this->outputShape_.dims() - spatial_dims] = h;
-          if (spatial_dims > 1) {
-            output_coords[this->outputShape_.dims() - spatial_dims + 1] = w;
-          }
+              size_t input_index = this->inputShape_.get_index(input_coords);
+              ValueType val = input[input_index];
 
-          size_t output_index = this->outputShape_.get_index(output_coords);
-
-          if (!pooling_buf.empty()) {
-            switch (this->poolingType_) {
-              case kAverage:
-                res[output_index] = avg_pooling(pooling_buf);
-                break;
-              case kMax:
-                res[output_index] = max_pooling(pooling_buf);
-                break;
-              default:
-                throw std::runtime_error("Unknown pooling type");
+              if (this->poolingType_ == kMax) {
+                if (count == 0 || val > max_val) {
+                  max_val = val;
+                }
+              } else {
+                sum += val;
+              }
+              ++count;
             }
           }
         }
-      }
-    }
-  }
-
-  return res;
-}
-
-template <typename ValueType>
-class PoolingLayerImplTBB : public PoolingLayerImpl<ValueType> {
- public:
-  PoolingLayerImplTBB(const Shape& input_shape, const Shape& pooling_shape,
-                      const Shape& strides = {2, 2},
-                      const Shape& pads = {0, 0, 0, 0},
-                      const Shape& dilations = {1, 1}, bool ceil_mode = false,
-                      const std::string& pooling_type = "average")
-      : PoolingLayerImpl<ValueType>(input_shape, pooling_shape, strides, pads,
-                                    dilations, ceil_mode, pooling_type) {}
-  std::vector<ValueType> run(
-      const std::vector<ValueType>& input) const override;
-};
-
-template <typename ValueType>
-std::vector<ValueType> PoolingLayerImplTBB<ValueType>::run(
-    const std::vector<ValueType>& input) const {
-  if (input.size() != this->inputShape_.count()) {
-    throw std::invalid_argument("Input size doesn't fit pooling layer");
-  }
 
-  std::vector<ValueType> res(this->outputShape_.count(), ValueType(0));
+        if (count == 0) return;
 
-  size_t spatial_dims = this->poolingShape_.dims();
-  int batch_dim = this->inputShape_.dims() > spatial_dims ? 0 : -1;
-  int channel_dim = this->inputShape_.dims() > spatial_dims + 1 ? 1 : -1;
+        std::vector<size_t> output_coords(this->outputShape_.dims(), 0);
+        if (batch_dim >= 0) output_coords[batch_dim] = n;
+        if (channel_dim >= 0) output_coords[channel_dim] = c;
+        output_coords[this->outputShape_.dims() - spatial_dims] = h;
+        if (spatial_dims > 1) {
+          output_coords[this->outputShape_.dims() - spatial_dims + 1] = w;
+        }
 
-  oneapi::tbb::parallel_for(
-      oneapi::tbb::blocked_range<size_t>(
-          0, batch_dim >= 0 ? this->outputShape_[batch_dim] : 1),
-      [&](const oneapi::tbb::blocked_range<size_t>& r1) {
-        for (size_t n = r1.begin(); n < r1.end(); n++) {
-          oneapi::tbb::parallel_for(
-              oneapi::tbb::blocked_range<size_t>(
-                  0, channel_dim >= 0 ? this->outputShape_[channel_dim] : 1),
-              [&](const oneapi::tbb::blocked_range<size_t>& r2) {
-                for (size_t c = r2.begin(); c < r2.end(); c++) {
-                  for (size_t h = 0;
-                       h < this->outputShape_[this->outputShape_.dims() -
-                                              spatial_dims];
-                       h++) {
-                    for (size_t w = 0;
-                         w <
-                         (spatial_dims > 1
-                              ? this->outputShape_[this->outputShape_.dims() -
-                                                   spatial_dims + 1]
-                              : 1);
-                         w++) {
-                      std::vector<ValueType> pooling_buf;
-
-                      int start_h = static_cast<int>(h * this->strides_[0]) -
-                                    static_cast<int>(this->pads_[0]);
-                      int start_w =
-                          spatial_dims > 1
-                              ? static_cast<int>(w * this->strides_[1]) -
-                                    static_cast<int>(this->pads_[2])
-                              : 0;
-
-                      for (size_t kh = 0; kh < this->poolingShape_[0]; kh++) {
-                        for (size_t kw = 0;
-                             kw <
-                             (spatial_dims > 1 ? this->poolingShape_[1] : 1);
-                             kw++) {
-                          int pos_h = start_h + static_cast<int>(
-                                                    kh * this->dilations_[0]);
-                          int pos_w =
-                              spatial_dims > 1
-                                  ? start_w + static_cast<int>(
-                                                  kw * this->dilations_[1])
-                                  : 0;
-
-                          if (pos_h >= 0 &&
-                              pos_h < static_cast<int>(
-                                          this->inputShape_[this->inputShape_
-                                                                .dims() -
-                                                            spatial_dims]) &&
-                              (spatial_dims <= 1 ||
-                               (pos_w >= 0 &&
-                                pos_w < static_cast<int>(
-                                            this->inputShape_
-                                                [this->inputShape_.dims() -
-                                                 spatial_dims + 1])))) {
-                            std::vector<size_t> input_coords(
-                                this->inputShape_.dims(), 0);
-                            if (batch_dim >= 0) input_coords[batch_dim] = n;
-                            if (channel_dim >= 0) input_coords[channel_dim] = c;
-                            input_coords[this->inputShape_.dims() -
-                                         spatial_dims] = pos_h;
-                            if (spatial_dims > 1) {
-                              input_coords[this->inputShape_.dims() -
-                                           spatial_dims + 1] = pos_w;
-                            }
-
-                            size_t input_index =
-                                this->inputShape_.get_index(input_coords);
-                            pooling_buf.push_back(input[input_index]);
-                          }
-                        }
-                      }
-
-                      std::vector<size_t> output_coords(
-                          this->outputShape_.dims(), 0);
-                      if (batch_dim >= 0) output_coords[batch_dim] = n;
-                      if (channel_dim >= 0) output_coords[channel_dim] = c;
-                      output_coords[this->outputShape_.dims() - spatial_dims] =
-                          h;
-                      if (spatial_dims > 1) {
-                        output_coords[this->outputShape_.dims() - spatial_dims +
-                                      1] = w;
-                      }
-
-                      size_t output_index =
-                          this->outputShape_.get_index(output_coords);
-
-                      if (!pooling_buf.empty()) {
-                        switch (this->poolingType_) {
-                          case kAverage:
-                            res[output_index] = avg_pooling(pooling_buf);
-                            break;
-                          case kMax:
-                            res[output_index] = max_pooling(pooling_buf);
-                            break;
-                          default:
-                            throw std::runtime_error("Unknown pooling type");
-                        }
-                      }
-                    }
-                  }
-                }
-              });
+        size_t output_index = this->outputShape_.get_index(output_coords);
+
+        switch (this->poolingType_) {
+          case kAverage:
+            res[output_index] = sum / static_cast<ValueType>(count);
+            break;
+          case kMax:
+            res[output_index] = max_val;
+            break;
+          default:
+            throw std::runtime_error("Unknown pooling type");
         }
-      });
+      },
+      options);
 
   return res;
 }