[COMPMID-386] Github: Support SoftmaxLayer on different number of dimensions?

Change-Id: I7422b977538ff29930a90f078badc2edee78af93
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/146638
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>

4 files changed, 67 insertions, 11 deletions

diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h
index d72547ed07..cb04182c21 100644
--- a/arm_compute/core/utils/misc/ShapeCalculator.h
+++ b/arm_compute/core/utils/misc/ShapeCalculator.h
@@ -275,6 +275,38 @@ inline TensorShape compute_flatten_shape(const ITensorInfo *input)
     return output_shape;
 }
 
+inline TensorShape compute_softmax_shape(const ITensorInfo *input, size_t axis = 1)
+{
+    // The output shape will be a 2D version of the input. For instance:
+    // - [x,y,z] and axis 1 will return [x, y*z]
+    // - [x,y,z,w] and axis 2 will return [x*y, w*z]
+    // - [x,y,z,w] and axis 3 will return [x*y*z, w]
+    TensorShape shape2D = input->tensor_shape();
+
+    if(axis < input->num_dimensions())
+    {
+        // Collapse from axis onward (this changes the shape)
+        shape2D.collapse_from(axis);
+
+        // Collapse the rest (collapse is inclusive)
+        shape2D.collapse(shape2D.num_dimensions() - 1);
+    }
+    else
+    {
+        // Collapse everything
+        shape2D.collapse(shape2D.num_dimensions());
+    }
+
+    if(axis == 0)
+    {
+        // If axis is zero the first dim should be one. Since
+        // collapse is an inclusive operation we need to shift
+        shape2D.shift_right(1);
+    }
+
+    return shape2D;
+}
+
 inline TensorShape compute_interleave_custom_shape(const TensorShape &input, const int x_interleave, const int y_interleave)
 {
     TensorShape output_shape{ input };
diff --git a/arm_compute/runtime/CL/functions/CLSoftmaxLayer.h b/arm_compute/runtime/CL/functions/CLSoftmaxLayer.h
index 90c99d6569..8d2c03f930 100644
--- a/arm_compute/runtime/CL/functions/CLSoftmaxLayer.h
+++ b/arm_compute/runtime/CL/functions/CLSoftmaxLayer.h
@@ -58,16 +58,22 @@ public:
      * @param[in]  input  Source tensor. Data types supported: QASYMM8/F16/F32
      * @param[out] output Destination tensor. Data types supported: same as @p input
      * @param[in]  beta   (Optional) A scaling factor for the exponent. Defaults to 1.f
+     * @param[in]  axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                    dimensions together. For instance, given a [4x4x4x4] image,
+     *                    when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
      */
-    void configure(const ICLTensor *input, ICLTensor *output, float beta = 1.0f);
+    void configure(const ICLTensor *input, ICLTensor *output, float beta = 1.0f, size_t axis = 1);
     /** Static function to check if given info will lead to a valid configuration of @ref CLSoftmaxLayer
      *
      * @param[in] input  Source tensor. Data types supported: QASYMM8/F16/F32
      * @param[in] output Destination tensor. Data types supported: same as @p input
-     *
+     * @param[in] beta   (Optional) A scaling factor for the exponent. Defaults to 1.f
+     * @param[in] axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                    dimensions together. For instance, given a [4x4x4x4] image,
+     *                    when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
      * @return a status
      */
-    static Status validate(const ITensorInfo *input, const ITensorInfo *output);
+    static Status validate(const ITensorInfo *input, const ITensorInfo *output, float beta = 1.0f, size_t axis = 1);
 
     // Inherited methods overridden:
     void run() override;
@@ -82,19 +88,22 @@ private:
      *
      * @param[in] input  Original source tensor.
      * @param[in] output Original destination tensor.
+     * @param[in] axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                    dimensions together. For instance, given a [4x4x4x4] image,
+     *                    when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
      */
-    void configure_flatten_kernel(const ICLTensor *input, const ICLTensor *output);
+    void configure_reshape_input_kernel(const ICLTensor *input, const ICLTensor *output, size_t axis);
 
     CLMemoryGroup                  _memory_group;
     CLLogits1DMaxShiftExpSumKernel _max_shift_exp_sum_kernel;
     CLLogits1DNormKernel           _norm_kernel;
-    CLFlattenLayerKernel           _flatten_kernel;
+    std::unique_ptr<ICLKernel>     _flatten_kernel_ptr;
     CLReshapeLayerKernel           _reshape_kernel;
     CLTensor                       _max;
     CLTensor                       _sum;
     CLTensor                       _tmp;
-    CLTensor                       _input_flat;
-    CLTensor                       _output_flat;
+    CLTensor                       _input_flattened;
+    CLTensor                       _output_flattened;
     bool                           _needs_flattening;
 };
 }
diff --git a/arm_compute/runtime/GLES_COMPUTE/functions/GCSoftmaxLayer.h b/arm_compute/runtime/GLES_COMPUTE/functions/GCSoftmaxLayer.h
index 1011c9a2ef..f6c6edb6a1 100644
--- a/arm_compute/runtime/GLES_COMPUTE/functions/GCSoftmaxLayer.h
+++ b/arm_compute/runtime/GLES_COMPUTE/functions/GCSoftmaxLayer.h
@@ -52,9 +52,14 @@ public:
      *
      * @param[in]  input  Source tensor. Data types supported: F16/F32
      * @param[out] output Destination tensor. Data types supported: same as @p input
-     * @param[in]  beta   (Optional) A scaling factor for the exponent. Only beta = 1 is supported.
+     * @param[in]  beta   (Optional) A scaling factor for the exponent. Only beta = 1 is supported
+     * @param[in]  axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                    dimensions together. For instance, given a [4x4x4x4] image,
+     *                    when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
+     *
+     * @note The value of @p axis must be always 1 for GLES
      */
-    void configure(const IGCTensor *input, IGCTensor *output, float beta = 1.0f);
+    void configure(const IGCTensor *input, IGCTensor *output, float beta = 1.0f, size_t axis = 1);
 
     // Inherited methods overridden:
     void run() override;
diff --git a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
index 61f46004d6..3f5ec8e820 100644
--- a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
+++ b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
@@ -56,17 +56,27 @@ public:
      *                       last value of each row to the nearest multiple.
      * @param[out]    output Destination tensor. Data types supported: same as @p input.
      * @param[in]     beta   (Optional) A scaling factor for the exponent.
+     * @param[in]     axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                       dimensions together. For instance, given a [4x4x4x4] image,
+     *                       when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
+     *
+     * @note The value of @p axis must be always 1 for NEON
      */
-    void configure(ITensor *input, ITensor *output, float beta = 1.0f);
+    void configure(ITensor *input, ITensor *output, float beta = 1.0f, size_t axis = 1);
     /** Static function to check if given info will lead to a valid configuration of @ref NESoftmaxLayer
      *
      * @param[in] input  Source tensor. Data types supported: QASYMM8/F16/F32.
      * @param[in] output Destination tensor. Data types supported: same as @p input
      * @param[in] beta   (Optional) A scaling factor for the exponent.
+     * @param[in] axis   (Optional) Reduction axis. It has the purpose of squashing the first @p axis
+     *                   dimensions together. For instance, given a [4x4x4x4] image,
+     *                   when @p axis is 2, the Softmax reduction will be applied on each of the [4x4] planes of the input image.
+     *
+     * @note The value of @p axis must be always 1 for NEON
      *
      * @return a status
      */
-    static Status validate(const ITensorInfo *input, const ITensorInfo *output, float beta = 1.0f);
+    static Status validate(const ITensorInfo *input, const ITensorInfo *output, float beta = 1.0f, size_t axis = 1);
 
     // Inherited methods overridden:
     void run() override;