From b7b028d912654915bf3beb03a618f85d1913e305 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Shubham=20Gupta/System=20SW=20/SRI-Bangalore/Engineer/?=
 =?utf8?q?=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?= <shub98.gupta@samsung.com>
Date: Wed, 17 Oct 2018 10:49:14 +0530
Subject: [PATCH] Fix offset bug for BatchToSpace (#3204)

For issue #3105

This patch will fix offset bug from BatchToSpace

Signed-off-by: shubham <shub98.gupta@samsung.com>
---
 .../src/internal/layers/SimpleBatchToSpaceNd.cc    | 45 +++++++---------------
 .../src/internal/layers/SimpleBatchToSpaceNd.h     |  4 +-
 2 files changed, 17 insertions(+), 32 deletions(-)

diff --git a/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.cc b/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.cc
index d485e8a..4344b65 100644
--- a/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.cc
+++ b/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.cc
@@ -20,7 +20,7 @@
 
 void SimpleBatchToSpaceND::configure(::arm_compute::ITensor *input, ::arm_compute::ITensor *output,
                                      const int32_t *block_size,
-                                     const ::arm_compute::Coordinates &axises = {3, 1, 0, 2})
+                                     const ::arm_compute::Coordinates &axises)
 {
   const auto rank = axises.num_dimensions();
   assert(rank == 4);
@@ -34,26 +34,10 @@ void SimpleBatchToSpaceND::configure(::arm_compute::ITensor *input, ::arm_comput
   _axises = axises;
 }
 
-inline int32_t Offset4D(const ::arm_compute::TensorShape &shape, int32_t b, int32_t h, int32_t w,
-                        int32_t d, const ::arm_compute::Coordinates &axises)
-{
-  // b, h, w, d >= 0
-  size_t indexes[4];
-  indexes[axises[0]] = b;
-  indexes[axises[1]] = h;
-  indexes[axises[2]] = w;
-  indexes[axises[3]] = d;
-
-  int32_t offset = indexes[3] * shape[2] * shape[1] * shape[0];
-  offset += indexes[2] * shape[1] * shape[0];
-  offset += indexes[1] * shape[0];
-  offset += indexes[0];
-  return offset;
-}
-
 template <typename T>
-inline void BatchToSpaceND(const T *input_data, const ::arm_compute::TensorShape &input_shape,
-                           const int32_t *block_size_data, T *output_data,
+inline void BatchToSpaceND(const ::arm_compute::ITensor *input,
+                           const ::arm_compute::TensorShape &input_shape,
+                           const int32_t *block_size_data, ::arm_compute::ITensor *output,
                            const ::arm_compute::TensorShape &output_shape,
                            const ::arm_compute::Coordinates &axises)
 {
@@ -82,10 +66,13 @@ inline void BatchToSpaceND(const T *input_data, const ::arm_compute::TensorShape
               ((out_h % block_size_data[0]) * block_size_data[1] + out_w % block_size_data[1]) *
                   output_batch;
 
-          const int output_index = Offset4D(output_shape, out_b, out_h, out_w, out_d, axises);
-          const int input_index = Offset4D(input_shape, in_b, in_h, in_w, in_d, axises);
+          auto input_id =
+              asARMComputeCoordinates(::arm_compute::Coordinates{in_b, in_h, in_w, in_d}, axises);
+          auto output_id = asARMComputeCoordinates(
+              ::arm_compute::Coordinates{out_b, out_h, out_w, out_d}, axises);
 
-          output_data[output_index] = input_data[input_index];
+          *reinterpret_cast<T *>(output->ptr_to_element(output_id)) =
+              *reinterpret_cast<T *>(input->ptr_to_element(input_id));
         }
       }
     }
@@ -101,20 +88,16 @@ void SimpleBatchToSpaceND::run()
     CAST_CL(_output)->map(q);
   }
 
-  auto input_buf = _input->buffer();
-  auto output_buf = _output->buffer();
   switch (_input->info()->data_type())
   {
     case ::arm_compute::DataType::U8:
     case ::arm_compute::DataType::QASYMM8:
-      BatchToSpaceND(reinterpret_cast<const uint8_t *>(input_buf), _input->info()->tensor_shape(),
-                     _block_size, reinterpret_cast<uint8_t *>(output_buf),
-                     _output->info()->tensor_shape(), _axises);
+      BatchToSpaceND<uint8_t>(_input, _input->info()->tensor_shape(), _block_size, _output,
+                              _output->info()->tensor_shape(), _axises);
       break;
     case ::arm_compute::DataType::F32:
-      BatchToSpaceND(reinterpret_cast<const float *>(input_buf), _input->info()->tensor_shape(),
-                     _block_size, reinterpret_cast<float *>(output_buf),
-                     _output->info()->tensor_shape(), _axises);
+      BatchToSpaceND<float>(_input, _input->info()->tensor_shape(), _block_size, _output,
+                            _output->info()->tensor_shape(), _axises);
       break;
     default:
       ARM_COMPUTE_ERROR("DataType not supported");
diff --git a/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.h b/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.h
index 52a1d35..e4f3f42 100644
--- a/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.h
+++ b/runtimes/pure_arm_compute/src/internal/layers/SimpleBatchToSpaceNd.h
@@ -18,6 +18,7 @@
 #define __SIMPLE_BATCH_TO_SPACE_ND_H__
 
 #include "internal/arm_compute.h"
+#include "internal/arm_compute/Cast.h"
 #include <arm_compute/core/ITensor.h>
 #include <arm_compute/runtime/IFunction.h>
 
@@ -31,7 +32,8 @@ public:
    * @param[in]  block_size  Block size.
    */
   void configure(::arm_compute::ITensor *input, ::arm_compute::ITensor *output,
-                 const int32_t *block_size, const ::arm_compute::Coordinates &axises);
+                 const int32_t *block_size,
+                 const ::arm_compute::Coordinates &axises = getARMComputeAxises(4));
 
   void run() override;
 
-- 
2.7.4