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, 45 insertions, 0 deletions

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