IVGCVSW-3694 Add ArgMinMax implementation for Ref

 * Add ArgMinMax implementation
 * Add utility function to get number of elements between axis
 * Add utility function to get unsigned axis
 * Unit tests for ArgMinMax function

Signed-off-by: Narumol Prangnawarat <narumol.prangnawarat@arm.com>
Change-Id: I7bc3d610dda9526190187eb87394a8ed7a4b5cdd

1 files changed, 6 insertions, 17 deletions

diff --git a/src/backends/reference/workloads/Softmax.cpp b/src/backends/reference/workloads/Softmax.cpp
index ec4fdb8839..f745d816c2 100644
--- a/src/backends/reference/workloads/Softmax.cpp
+++ b/src/backends/reference/workloads/Softmax.cpp
@@ -5,27 +5,14 @@
 
 #include "Softmax.hpp"
 
+#include <TensorUtils.hpp>
+
 #include <cmath>
 #include <vector>
 
 namespace armnn
 {
 
-unsigned int GetNumElementsBetween(const TensorShape& shape,
-                                   unsigned int firstAxisInclusive,
-                                   unsigned int lastAxisExclusive)
-{
-    BOOST_ASSERT(0 <= firstAxisInclusive);
-    BOOST_ASSERT(firstAxisInclusive <= lastAxisExclusive);
-    BOOST_ASSERT(lastAxisExclusive <= shape.GetNumDimensions());
-    unsigned int count = 1;
-    for (unsigned int i = firstAxisInclusive; i < lastAxisExclusive; i++)
-    {
-        count *= shape[i];
-    }
-    return count;
-}
-
 /// Computes the softmax function on some inputs, into outputs, with a shape given by tensorInfo.
 void Softmax(Decoder<float>& in, Encoder<float>& out, const TensorInfo& inputTensorInfo, float beta, int axis)
 {
@@ -39,9 +26,11 @@ void Softmax(Decoder<float>& in, Encoder<float>& out, const TensorInfo& inputTen
                          : static_cast<unsigned int>(axis);
 
     const TensorShape& inputShape = inputTensorInfo.GetShape();
-    const unsigned int outerSize  = GetNumElementsBetween(inputShape, 0, uAxis);
+    const unsigned int outerSize  = armnnUtils::GetNumElementsBetween(inputShape, 0, uAxis);
     const unsigned int axisSize   = inputShape[uAxis];
-    const unsigned int innerSize  = GetNumElementsBetween(inputShape, uAxis + 1, inputShape.GetNumDimensions());
+    const unsigned int innerSize  = armnnUtils::GetNumElementsBetween(inputShape,
+                                                                      uAxis + 1,
+                                                                      inputShape.GetNumDimensions());
 
     for (unsigned int outer = 0; outer < outerSize; ++outer)
     {