Make data_layout an attribute of the Scale function

Resolves COMPMID-4208

Change-Id: I61ca670134a005462ad0528a5aff9507a90860e7
Signed-off-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/4942
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>

6 files changed, 46 insertions, 41 deletions

diff --git a/arm_compute/core/KernelDescriptors.h b/arm_compute/core/KernelDescriptors.h
index ea46bfa5a6..e381220695 100644
--- a/arm_compute/core/KernelDescriptors.h
+++ b/arm_compute/core/KernelDescriptors.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020 Arm Limited.
+ * Copyright (c) 2019-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -73,8 +73,8 @@ struct GEMMKernelInfo
         int32_t             ina_offset,
         int32_t             inb_offset)
         : m(im), n(in), k(ik), depth_output_gemm3d(idepth_output_gemm3d), reinterpret_input_as_3d(ireinterpret_input_as_3d), broadcast_bias(ibroadcast_bias), fp_mixed_precision(ifp_mixed_precision),
-          has_pad_y(ihas_pad_y), activation_info(iactivation_info), mult_transpose1xW_width(inmult_transpose1xW_width), mult_interleave4x4_height(imult_interleave4x4_height), lhs_info(ilhs_info), rhs_info(irhs_info),
-          a_offset(ina_offset), b_offset(inb_offset)
+          has_pad_y(ihas_pad_y), activation_info(iactivation_info), mult_transpose1xW_width(inmult_transpose1xW_width), mult_interleave4x4_height(imult_interleave4x4_height), lhs_info(ilhs_info),
+          rhs_info(irhs_info), a_offset(ina_offset), b_offset(inb_offset)
     {
     }
 
@@ -182,19 +182,22 @@ struct ScaleKernelInfo
      * @param[in] sampling_policy       (Optional) Sampling policy used by the interpolation. Defaults to @ref SamplingPolicy::CENTER
      * @param[in] use_padding           (Optional) Is padding in use or not. Defaults to true.
      * @param[in] align_corners         (Optional) Align corners of input and output, only affecting bilinear policy with TOP_LEFT sampling policy. Defaults to false.
+     * @param[in] data_layout           (Optional) Data layout used by the layer. Defaults to @ref DataLayout::UNKNOWN
      */
     ScaleKernelInfo(InterpolationPolicy interpolation_policy,
                     BorderMode          border_mode,
                     PixelValue          constant_border_value = PixelValue(),
                     SamplingPolicy      sampling_policy       = SamplingPolicy::CENTER,
                     bool                use_padding           = true,
-                    bool                align_corners         = false)
+                    bool                align_corners         = false,
+                    DataLayout          data_layout           = DataLayout::UNKNOWN)
         : interpolation_policy{ interpolation_policy },
           border_mode{ border_mode },
           constant_border_value{ constant_border_value },
           sampling_policy{ sampling_policy },
           use_padding{ use_padding },
-          align_corners{ align_corners }
+          align_corners{ align_corners },
+          data_layout{ data_layout }
     {
     }
 
@@ -204,6 +207,7 @@ struct ScaleKernelInfo
     SamplingPolicy      sampling_policy;       /**< Sampling policy used by the interpolation. */
     bool                use_padding;           /**< Indication of using padding */
     bool                align_corners;         /**< Align corners of input and output */
+    DataLayout          data_layout;           /**< Data layout to use */
 };
 
 struct ThresholdKernelInfo
diff --git a/src/core/CL/kernels/CLScaleKernel.cpp b/src/core/CL/kernels/CLScaleKernel.cpp
index f3d2fa12d5..c2c78c8f6d 100644
--- a/src/core/CL/kernels/CLScaleKernel.cpp
+++ b/src/core/CL/kernels/CLScaleKernel.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -41,15 +41,15 @@
 #include <set>
 #include <string>
 
-using namespace arm_compute;
-
+namespace arm_compute
+{
 namespace
 {
-inline std::pair<float, float> calculate_scale_factors(const ITensorInfo &input, const ITensorInfo &output, bool align_corners)
+inline std::pair<float, float> calculate_scale_factors(const ITensorInfo &input, const ITensorInfo &output, const ScaleKernelInfo &info)
 {
-    DataLayout data_layout = input.data_layout();
-    const int  idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
-    const int  idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+    const DataLayout data_layout = info.data_layout == DataLayout::UNKNOWN ? input.data_layout() : info.data_layout;
+    const int        idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
+    const int        idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
 
     // Compute the ratio between source width/height and destination width/height
     const unsigned int input_width   = input.dimension(idx_width);
@@ -57,8 +57,8 @@ inline std::pair<float, float> calculate_scale_factors(const ITensorInfo &input,
     const unsigned int output_width  = output.dimension(idx_width);
     const unsigned int output_height = output.dimension(idx_height);
 
-    float wr = arm_compute::scale_utils::calculate_resize_ratio(input_width, output_width, align_corners);
-    float hr = arm_compute::scale_utils::calculate_resize_ratio(input_height, output_height, align_corners);
+    float wr = arm_compute::scale_utils::calculate_resize_ratio(input_width, output_width, info.align_corners);
+    float hr = arm_compute::scale_utils::calculate_resize_ratio(input_height, output_height, info.align_corners);
 
     return std::make_pair(wr, hr);
 }
@@ -73,11 +73,9 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, c
     ARM_COMPUTE_RETURN_ERROR_ON(output == input);
     ARM_COMPUTE_RETURN_ERROR_ON(info.align_corners && !arm_compute::scale_utils::is_align_corners_allowed_sampling_policy(info.sampling_policy));
 
-    const bool will_use_align_corners = info.align_corners;
-
     float wr = 0.f;
     float hr = 0.f;
-    std::tie(wr, hr) = calculate_scale_factors(*input, *output, will_use_align_corners);
+    std::tie(wr, hr) = calculate_scale_factors(*input, *output, info);
 
     ARM_COMPUTE_RETURN_ERROR_ON(info.interpolation_policy == InterpolationPolicy::AREA && (wr > 1.f || hr > 1.f));
 
@@ -86,10 +84,10 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, c
 
 std::pair<Status, Window> validate_and_configure_window(ITensorInfo *input, ITensorInfo *output, const ScaleKernelInfo &info, BorderSize &border)
 {
-    Window       win{};
-    bool         window_changed{};
-    unsigned int num_elems_processed_per_iteration = 0;
-    DataLayout   data_layout                       = input->data_layout();
+    Window           win{};
+    bool             window_changed{};
+    unsigned int     num_elems_processed_per_iteration = 0;
+    const DataLayout data_layout                       = info.data_layout == DataLayout::UNKNOWN ? input->data_layout() : info.data_layout;
 
     switch(data_layout)
     {
@@ -141,7 +139,8 @@ BorderSize CLScaleKernel::border_size() const
 Status CLScaleKernel::validate(const ITensorInfo *input, const ITensorInfo *output, const ScaleKernelInfo &info)
 {
     ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments(input, output, info));
-    BorderSize border = BorderSize(static_cast<size_t>(input->data_layout() == DataLayout::NCHW));
+    const DataLayout data_layout = info.data_layout == DataLayout::UNKNOWN ? input->data_layout() : info.data_layout;
+    BorderSize       border      = BorderSize(static_cast<size_t>(data_layout == DataLayout::NCHW));
     ARM_COMPUTE_RETURN_ON_ERROR(validate_and_configure_window(input->clone().get(), output->clone().get(), info, border).first);
 
     return Status{};
@@ -170,12 +169,12 @@ void CLScaleKernel::configure(const CLCompileContext &compile_context, const ICL
     _input                = input;
     _output               = output;
     _interpolation_policy = info.interpolation_policy;
-    _data_layout          = input->info()->data_layout();
+    _data_layout          = info.data_layout == DataLayout::UNKNOWN ? input->info()->data_layout() : info.data_layout;
     _align_corners        = info.align_corners;
 
     float wr = 0.f;
     float hr = 0.f;
-    std::tie(wr, hr) = calculate_scale_factors(*input->info(), *output->info(), _align_corners);
+    std::tie(wr, hr) = calculate_scale_factors(*input->info(), *output->info(), info);
 
     const bool call_quantized_kernel = is_data_type_quantized_asymmetric(input->info()->data_type()) && _interpolation_policy == InterpolationPolicy::BILINEAR;
 
@@ -284,3 +283,4 @@ void CLScaleKernel::run(const Window &window, cl::CommandQueue &queue)
             ARM_COMPUTE_ERROR("Data layout not supported");
     }
 }
+} // namespace arm_compute
diff --git a/src/core/NEON/kernels/NEScaleKernel.cpp b/src/core/NEON/kernels/NEScaleKernel.cpp
index 39ed6317a1..f2c11b203c 100644
--- a/src/core/NEON/kernels/NEScaleKernel.cpp
+++ b/src/core/NEON/kernels/NEScaleKernel.cpp
@@ -158,7 +158,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *dx, const
     ARM_COMPUTE_UNUSED(info.constant_border_value);
     ARM_COMPUTE_RETURN_ERROR_ON_MSG(info.use_padding, "Padding is not supported");
 
-    const DataLayout data_layout   = input->data_layout();
+    const DataLayout data_layout   = info.data_layout == DataLayout::UNKNOWN ? input->data_layout() : info.data_layout;
     const auto       width_index   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
     const auto       height_index  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
     const auto       output_width  = output->dimension(width_index);
@@ -192,7 +192,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *dx, const
 
 NEScaleKernel::NEScaleKernel()
     : _func(nullptr), _offsets(nullptr), _dx(nullptr), _dy(nullptr), _input(nullptr), _output(nullptr), _policy(), _border_mode(), _constant_border_value(PixelValue()), _sampling_offset(0),
-      _align_corners(false)
+      _align_corners(false), _data_layout(DataLayout::UNKNOWN)
 {
 }
 
@@ -209,9 +209,9 @@ void NEScaleKernel::configure(const ITensor *input, const ITensor *dx, const ITe
                                                   info));
 
     // Get data layout and width/height indices
-    const DataLayout data_layout = input->info()->data_layout();
-    const int        idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
-    const int        idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+    _data_layout         = info.data_layout == DataLayout::UNKNOWN ? input->info()->data_layout() : info.data_layout;
+    const int idx_width  = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::WIDTH);
+    const int idx_height = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::HEIGHT);
 
     _input                 = input;
     _output                = output;
@@ -242,11 +242,11 @@ void NEScaleKernel::configure(const ITensor *input, const ITensor *dx, const ITe
     }
 
     // Configure scale function to run
-    if(_input->info()->data_layout() == DataLayout::NCHW)
+    if(_data_layout == DataLayout::NCHW)
     {
         std::string function_to_call("scale_");
         function_to_call += string_from_data_type(_input->info()->data_type()) + "_";
-        function_to_call += string_from_data_layout(_input->info()->data_layout()) + "_";
+        function_to_call += string_from_data_layout(_data_layout) + "_";
         function_to_call += string_from_interpolation_policy(_policy);
 
         static std::map<std::string, ScaleFunctionPtr> map_function =
@@ -471,9 +471,8 @@ template <typename T>
 void NEScaleKernel::scale_bilinear_qasymm(const Window &window)
 {
     // Get data layout and width/height indices
-    const DataLayout data_layout = _input->info()->data_layout();
-    const int        idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
-    const int        idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+    const int idx_width  = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::WIDTH);
+    const int idx_height = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::HEIGHT);
 
     // Compute the ratio between source height and destination height
     const auto hr = scale_utils::calculate_resize_ratio(_input->info()->dimension(idx_height), _output->info()->dimension(idx_height), _align_corners);
@@ -586,9 +585,9 @@ void NEScaleKernel::run(const Window &window, const ThreadInfo &info)
     ARM_COMPUTE_UNUSED(info);
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
     ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(INEKernel::window(), window);
-    ARM_COMPUTE_ERROR_ON(_func == nullptr && _input->info()->data_layout() == DataLayout::NCHW);
+    ARM_COMPUTE_ERROR_ON(_func == nullptr && _data_layout == DataLayout::NCHW);
 
-    if(_input->info()->data_layout() == DataLayout::NCHW)
+    if(_data_layout == DataLayout::NCHW)
     {
         (this->*_func)(window);
     }
diff --git a/src/core/NEON/kernels/NEScaleKernel.h b/src/core/NEON/kernels/NEScaleKernel.h
index b93a213e99..f6ee3fa4c5 100644
--- a/src/core/NEON/kernels/NEScaleKernel.h
+++ b/src/core/NEON/kernels/NEScaleKernel.h
@@ -116,6 +116,7 @@ private:
     PixelValue          _constant_border_value;
     float               _sampling_offset;
     bool                _align_corners;
+    DataLayout          _data_layout;
 };
 } // namespace arm_compute
 #endif /*ARM_COMPUTE_NESCALEKERNEL_H */
diff --git a/src/runtime/CL/functions/CLScale.cpp b/src/runtime/CL/functions/CLScale.cpp
index aab5d9ba73..9862d0a1b3 100644
--- a/src/runtime/CL/functions/CLScale.cpp
+++ b/src/runtime/CL/functions/CLScale.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -50,7 +50,8 @@ void CLScale::configure(const CLCompileContext &compile_context, ICLTensor *inpu
     // Tune kernels
     CLScheduler::get().tune_kernel_static(*_kernel);
 
-    if(input->info()->data_layout() == DataLayout::NCHW && !_kernel->border_size().empty())
+    const DataLayout data_layout = info.data_layout == DataLayout::UNKNOWN ? input->info()->data_layout() : info.data_layout;
+    if(data_layout == DataLayout::NCHW && !_kernel->border_size().empty())
     {
         _border_handler->configure(compile_context, input, _kernel->border_size(), info.border_mode, info.constant_border_value);
     }
diff --git a/src/runtime/NEON/functions/NEScale.cpp b/src/runtime/NEON/functions/NEScale.cpp
index 9d6e2ca754..f91de32191 100644
--- a/src/runtime/NEON/functions/NEScale.cpp
+++ b/src/runtime/NEON/functions/NEScale.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -109,7 +109,7 @@ void NEScale::configure(ITensor *input, ITensor *output, const ScaleKernelInfo &
     const bool is_align_corners_used = info.align_corners && arm_compute::scale_utils::is_align_corners_allowed_sampling_policy(info.sampling_policy);
 
     // Get data layout and width/height indices
-    const DataLayout data_layout = input->info()->data_layout();
+    const DataLayout data_layout = info.data_layout == DataLayout::UNKNOWN ? input->info()->data_layout() : info.data_layout;
     const int        idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
     const int        idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
 
@@ -182,7 +182,7 @@ Status NEScale::validate(const ITensorInfo *input, const ITensorInfo *output, co
     ITensorInfo *dy      = nullptr;
 
     // Get data layout and width/height indices
-    const DataLayout data_layout = input->data_layout();
+    const DataLayout data_layout = info.data_layout == DataLayout::UNKNOWN ? input->data_layout() : info.data_layout;
     const int        idx_width   = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
     const int        idx_height  = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);