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

8 files changed, 166 insertions, 17 deletions

diff --git a/src/armnnUtils/TensorUtils.cpp b/src/armnnUtils/TensorUtils.cpp
index 8baea78ab5..b4e8d5acda 100644
--- a/src/armnnUtils/TensorUtils.cpp
+++ b/src/armnnUtils/TensorUtils.cpp
@@ -107,4 +107,32 @@ armnn::TensorShape ExpandDims(const armnn::TensorShape& tensorShape, int axis)
     return armnn::TensorShape(outputDim, outputShape.data());
 }
 
+unsigned int GetNumElementsBetween(const armnn::TensorShape& shape,
+                                   const unsigned int firstAxisInclusive,
+                                   const 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;
+}
+
+unsigned int GetUnsignedAxis(const unsigned int inputDimension, const int axis)
+{
+    BOOST_ASSERT_MSG(axis < boost::numeric_cast<int>(inputDimension),
+                     "Required axis index greater than number of dimensions.");
+    BOOST_ASSERT_MSG(axis >= -boost::numeric_cast<int>(inputDimension),
+                     "Required axis index lower than negative of the number of dimensions");
+
+    unsigned int uAxis = axis < 0  ?
+                         inputDimension - boost::numeric_cast<unsigned int>(abs(axis))
+                         : boost::numeric_cast<unsigned int>(axis);
+    return uAxis;
+}
+
 }
diff --git a/src/armnnUtils/TensorUtils.hpp b/src/armnnUtils/TensorUtils.hpp
index 03b1c8a2df..2b1f6a24f3 100644
--- a/src/armnnUtils/TensorUtils.hpp
+++ b/src/armnnUtils/TensorUtils.hpp
@@ -26,4 +26,10 @@ std::pair<float, float> FindMinMax(armnn::ITensorHandle* tensorHandle);
 
 armnn::TensorShape ExpandDims(const armnn::TensorShape& tensorShape, int axis);
 
+unsigned int GetNumElementsBetween(const armnn::TensorShape& shape,
+                                   unsigned int firstAxisInclusive,
+                                   unsigned int lastAxisExclusive);
+
+unsigned int GetUnsignedAxis(const unsigned int inputDimension, const int axis);
+
 } // namespace armnnUtils
diff --git a/src/backends/reference/test/ArgMinMaxTests.cpp b/src/backends/reference/test/ArgMinMaxTests.cpp
new file mode 100644
index 0000000000..201a2c0c2e
--- /dev/null
+++ b/src/backends/reference/test/ArgMinMaxTests.cpp
@@ -0,0 +1,58 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include <reference/workloads/ArgMinMax.hpp>
+
+#include <boost/test/unit_test.hpp>
+
+BOOST_AUTO_TEST_SUITE(RefArgMinMax)
+
+BOOST_AUTO_TEST_CASE(ArgMinTest)
+{
+    const armnn::TensorInfo inputInfo({ 1, 2, 3 } , armnn::DataType::Float32);
+    const armnn::TensorInfo outputInfo({ 1, 3 }, armnn::DataType::Float32);
+
+    std::vector<float> inputValues({ 1.0f, 5.0f, 3.0f, 4.0f, 2.0f, 6.0f});
+    std::vector<int32_t> outputValues(outputInfo.GetNumElements());
+    std::vector<int32_t> expectedValues({ 0, 1, 0 });
+
+    ArgMinMax(*armnn::MakeDecoder<float>(inputInfo, inputValues.data()),
+               outputValues.data(),
+               inputInfo,
+               outputInfo,
+               armnn::ArgMinMaxFunction::Min,
+               -2);
+
+    BOOST_CHECK_EQUAL_COLLECTIONS(outputValues.begin(),
+                                  outputValues.end(),
+                                  expectedValues.begin(),
+                                  expectedValues.end());
+
+}
+
+BOOST_AUTO_TEST_CASE(ArgMaxTest)
+{
+    const armnn::TensorInfo inputInfo({ 1, 2, 3 } , armnn::DataType::Float32);
+    const armnn::TensorInfo outputInfo({ 1, 3 }, armnn::DataType::Float32);
+
+    std::vector<float> inputValues({ 1.0f, 5.0f, 3.0f, 4.0f, 2.0f, 6.0f });
+    std::vector<int32_t> outputValues(outputInfo.GetNumElements());
+    std::vector<int32_t> expectedValues({ 1, 0, 1 });
+
+    ArgMinMax(*armnn::MakeDecoder<float>(inputInfo, inputValues.data()),
+               outputValues.data(),
+               inputInfo,
+               outputInfo,
+               armnn::ArgMinMaxFunction::Max,
+               -2);
+
+    BOOST_CHECK_EQUAL_COLLECTIONS(outputValues.begin(),
+                                  outputValues.end(),
+                                  expectedValues.begin(),
+                                  expectedValues.end());
+
+}
+
+BOOST_AUTO_TEST_SUITE_END()
\ No newline at end of file
diff --git a/src/backends/reference/test/CMakeLists.txt b/src/backends/reference/test/CMakeLists.txt
index b56b35354e..2a7aa1501a 100644
--- a/src/backends/reference/test/CMakeLists.txt
+++ b/src/backends/reference/test/CMakeLists.txt
@@ -4,6 +4,7 @@
 #
 
 list(APPEND armnnRefBackendUnitTests_sources
+    ArgMinMaxTests.cpp
     RefCreateWorkloadTests.cpp
     RefDetectionPostProcessTests.cpp
     RefEndToEndTests.cpp
diff --git a/src/backends/reference/workloads/ArgMinMax.cpp b/src/backends/reference/workloads/ArgMinMax.cpp
new file mode 100644
index 0000000000..2687a4e8ac
--- /dev/null
+++ b/src/backends/reference/workloads/ArgMinMax.cpp
@@ -0,0 +1,45 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "ArgMinMax.hpp"
+
+#include <TensorUtils.hpp>
+
+#include <boost/numeric/conversion/cast.hpp>
+
+namespace armnn
+{
+
+void ArgMinMax(Decoder<float>& in, int32_t* out, const TensorInfo& inputTensorInfo,
+               const TensorInfo& outputTensorInfo, ArgMinMaxFunction function, int axis)
+{
+    unsigned int uAxis = armnnUtils::GetUnsignedAxis(inputTensorInfo.GetNumDimensions(), axis);
+
+    const unsigned int outerElements = armnnUtils::GetNumElementsBetween(inputTensorInfo.GetShape(), 0, uAxis);
+    const unsigned int axisSize = inputTensorInfo.GetShape()[uAxis];
+    const unsigned int innerElements = armnnUtils::GetNumElementsBetween(inputTensorInfo.GetShape(),
+                                                                         uAxis + 1,
+                                                                         inputTensorInfo.GetNumDimensions());
+
+    for (unsigned int outer = 0; outer < outerElements; ++outer) {
+        for (unsigned int inner = 0; inner < innerElements; ++inner) {
+            in[outer * axisSize * innerElements + inner];
+            auto tmpValue = in.Get();
+            unsigned int tmpIndex = 0;
+            for (unsigned int i = 1; i < axisSize; ++i) {
+                in[(outer * axisSize * innerElements) + (i * innerElements) + inner];
+                const auto& value = in.Get();
+                if ((function == armnn::ArgMinMaxFunction::Min && value < tmpValue) ||
+                    (function == armnn::ArgMinMaxFunction::Max &&  value > tmpValue)) {
+                    tmpValue = value;
+                    tmpIndex = i;
+                }
+            }
+            out[outer * innerElements + inner] = boost::numeric_cast<int32_t>(tmpIndex);
+        }
+    }
+}
+
+} //namespace armnn
diff --git a/src/backends/reference/workloads/ArgMinMax.hpp b/src/backends/reference/workloads/ArgMinMax.hpp
new file mode 100644
index 0000000000..5a9c6a8a2a
--- /dev/null
+++ b/src/backends/reference/workloads/ArgMinMax.hpp
@@ -0,0 +1,20 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include "armnn/Tensor.hpp"
+#include "armnn/Descriptors.hpp"
+
+#include "Decoders.hpp"
+
+namespace armnn
+{
+
+void ArgMinMax(Decoder<float>& in, int32_t* out, const TensorInfo& inputTensorInfo,
+               const TensorInfo& outputTensorInfo, ArgMinMaxFunction function, int axis);
+
+} //namespace armnn
+
diff --git a/src/backends/reference/workloads/CMakeLists.txt b/src/backends/reference/workloads/CMakeLists.txt
index 7f49e800e5..23d6024530 100644
--- a/src/backends/reference/workloads/CMakeLists.txt
+++ b/src/backends/reference/workloads/CMakeLists.txt
@@ -8,6 +8,8 @@ list(APPEND armnnRefBackendWorkloads_sources
     Abs.hpp
     Activation.cpp
     Activation.hpp
+    ArgMinMax.cpp
+    ArgMinMax.hpp
     BaseIterator.hpp
     BatchNormImpl.cpp
     BatchNormImpl.hpp
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)
     {