Support Padding accross all dimensions (#3305)

This patch will provide support to pad accross all dimensions of input tensor

Signed-off-by: shubham <shub98.gupta@samsung.com>

diff --git a/runtimes/pure_arm_compute/src/internal/layers/SimplePadLayer.cc b/runtimes/pure_arm_compute/src/internal/layers/SimplePadLayer.cc
index 65bb512..6f0bdc3 100644 (file)
--- a/runtimes/pure_arm_compute/src/internal/layers/SimplePadLayer.cc
+++ b/runtimes/pure_arm_compute/src/internal/layers/SimplePadLayer.cc
@@ -17,6 +17,41 @@
 #include "internal/layers/SimplePadLayer.h"
 #include <arm_compute/runtime/CL/CLScheduler.h>
 
+namespace
+{
+bool validate_arg(const ::arm_compute::ITensor *input, const ::arm_compute::ITensor *output,
+                  const ::arm_compute::ITensor *padding_size,
+                  const ::arm_compute::Coordinates &axises)
+{
+  const int input_batch = input->info()->tensor_shape()[axises[0]];
+  const int input_height = input->info()->tensor_shape()[axises[1]];
+  const int input_width = input->info()->tensor_shape()[axises[2]];
+  const int input_depth = input->info()->tensor_shape()[axises[3]];
+
+  const int output_batch = output->info()->tensor_shape()[axises[0]];
+  const int output_height = output->info()->tensor_shape()[axises[1]];
+  const int output_width = output->info()->tensor_shape()[axises[2]];
+  const int output_depth = output->info()->tensor_shape()[axises[3]];
+
+  auto pad_batch_up = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 0}));
+  auto pad_batch_down = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({1, 0}));
+  auto pad_height_top = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 1}));
+  auto pad_height_bottom = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({1, 1}));
+  auto pad_width_left = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 2}));
+  auto pad_width_right = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({1, 2}));
+  auto pad_depth_front = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 3}));
+  auto pad_depth_back = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({1, 3}));
+
+  const int padded_batch = input_batch + pad_batch_up + pad_batch_down;
+  const int padded_height = input_height + pad_height_top + pad_height_bottom;
+  const int padded_width = input_width + pad_width_left + pad_width_right;
+  const int padded_depth = input_depth + pad_depth_front + pad_depth_back;
+
+  return (padded_batch == output_batch) && (padded_height == output_height) &&
+         (padded_width == output_width) && (padded_depth == output_depth);
+}
+} // namespace
+
 void SimplePadLayer::configure(::arm_compute::ITensor *input, ::arm_compute::ITensor *output,
                                ::arm_compute::ITensor *padding_size,
                                const ::arm_compute::Coordinates &axises)
@@ -46,56 +81,50 @@ inline void ApplyPadding(const ::arm_compute::ITensor *input_data,
                          const ::arm_compute::TensorShape &output_shape,
                          const ::arm_compute::Coordinates &axises)
 {
+
+  assert(validate_arg(input_data, output_data, padding_size, axises) &&
+         "Padded Input shape does not match to output shape");
+
+  const int input_batch = input_shape[axises[0]];
   const int input_height = input_shape[axises[1]];
   const int input_width = input_shape[axises[2]];
+  const int input_depth = input_shape[axises[3]];
 
-  const int batch = output_shape[axises[0]];
+  const int output_batch = output_shape[axises[0]];
   const int output_height = output_shape[axises[1]];
   const int output_width = output_shape[axises[2]];
-  const int depth = output_shape[axises[3]];
-
-  // Supports only Spatial padding
-  // Padding size for top, bottom, left and right are required.
-  auto pad_top = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 1}));
-  auto pad_bottom = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({1, 1}));
-  auto pad_left = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 2}));
-  auto pad_right = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({1, 2}));
-
-  const int padded_height = input_height + pad_top + pad_bottom;
-  const int padded_width = input_width + pad_left + pad_right;
-
-  { // new block for assertions
-    assert(input_shape[axises[0]] == output_shape[axises[0]]);
-    assert(padded_height == output_height);
-    assert(padded_width == output_width);
-    assert(input_shape[axises[3]] == output_shape[axises[3]]);
-  }
+  const int output_depth = output_shape[axises[3]];
 
-  for (int in_b = 0; in_b < batch; ++in_b)
+  // Padding size for Up, Top, Left and Front are required.
+  auto pad_batch_up = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 0}));
+  auto pad_height_top = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 1}));
+  auto pad_width_left = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 2}));
+  auto pad_depth_front = *reinterpret_cast<const int32_t *>(padding_size->ptr_to_element({0, 3}));
+
+  for (int out_b = 0; out_b < output_batch; ++out_b)
   {
-    for (int in_h = 0; in_h < padded_height; ++in_h)
+    for (int out_h = 0; out_h < output_height; ++out_h)
     {
-      for (int in_w = 0; in_w < padded_width; ++in_w)
+      for (int out_w = 0; out_w < output_width; ++out_w)
       {
-        for (int in_d = 0; in_d < depth; ++in_d)
+        for (int out_d = 0; out_d < output_depth; ++out_d)
         {
-          const int out_d = in_d;
-          const int out_h = in_h;
-          const int out_w = in_w;
-          const int out_b = in_b;
-
           auto output_id = asARMComputeCoordinates(
               ::arm_compute::Coordinates{out_b, out_h, out_w, out_d}, axises);
 
-          if (in_h < pad_top || in_h >= (input_height + pad_top) || in_w < pad_left ||
-              in_w >= (pad_left + input_width))
+          if (out_b < pad_batch_up || out_b >= (input_batch + pad_batch_up) ||
+              out_h < pad_height_top || out_h >= (input_height + pad_height_top) ||
+              out_w < pad_width_left || out_w >= (input_width + pad_width_left) ||
+              out_d < pad_depth_front || out_d >= (input_depth + pad_depth_front))
           {
             *reinterpret_cast<T *>(output_data->ptr_to_element(output_id)) = 0;
           }
           else
           {
             auto input_id = asARMComputeCoordinates(
-                ::arm_compute::Coordinates{in_b, in_h - pad_top, in_w - pad_left, in_d}, axises);
+                ::arm_compute::Coordinates{out_b - pad_batch_up, out_h - pad_height_top,
+                                           out_w - pad_width_left, out_d - pad_depth_front},
+                axises);
             *reinterpret_cast<T *>(output_data->ptr_to_element(output_id)) =
                 *reinterpret_cast<T *>(input_data->ptr_to_element(input_id));
           }