4 files changed, 73 insertions, 17 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
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)
     {