IVGCVSW-3723 Adding reference workload support for ArgMinMax

Change-Id: I65209ecec4e3abf808163239748d6e830568c2e3
Signed-off-by: Nikhil Raj <nikhil.raj@arm.com>

10 files changed, 120 insertions, 0 deletions

diff --git a/src/backends/reference/RefLayerSupport.cpp b/src/backends/reference/RefLayerSupport.cpp
index 572f617636..14183a7b68 100644
--- a/src/backends/reference/RefLayerSupport.cpp
+++ b/src/backends/reference/RefLayerSupport.cpp
@@ -192,6 +192,29 @@ bool RefLayerSupport::IsAdditionSupported(const TensorInfo& input0,
     return supported;
 }
 
+bool RefLayerSupport::IsArgMinMaxSupported(const armnn::TensorInfo &input, const armnn::TensorInfo &output,
+                                           const armnn::ArgMinMaxDescriptor &descriptor,
+                                           armnn::Optional<std::string &> reasonIfUnsupported) const
+{
+    ignore_unused(descriptor);
+
+    std::array<DataType, 3> supportedTypes =
+    {
+        DataType::Float32,
+        DataType::QuantisedAsymm8,
+        DataType::QuantisedSymm16
+    };
+
+    bool supported = true;
+
+    supported &= CheckSupportRule(TypeAnyOf(input, supportedTypes), reasonIfUnsupported,
+                                  "Reference ArgMinMax: input is not a supported type.");
+    supported &= CheckSupportRule(TypeIs(output, DataType::Signed32), reasonIfUnsupported,
+                                  "Reference ArgMinMax: output type not supported");
+
+    return supported;
+}
+
 bool RefLayerSupport::IsBatchNormalizationSupported(const TensorInfo& input,
                                                     const TensorInfo& output,
                                                     const TensorInfo& mean,
diff --git a/src/backends/reference/RefLayerSupport.hpp b/src/backends/reference/RefLayerSupport.hpp
index 8200058633..9e8c914512 100644
--- a/src/backends/reference/RefLayerSupport.hpp
+++ b/src/backends/reference/RefLayerSupport.hpp
@@ -26,6 +26,11 @@ public:
                              const TensorInfo& output,
                              Optional<std::string&> reasonIfUnsupported = EmptyOptional()) const override;
 
+    bool IsArgMinMaxSupported(const TensorInfo& input,
+                              const TensorInfo& output,
+                              const ArgMinMaxDescriptor& descriptor,
+                              Optional<std::string&> reasonIfUnsupported = EmptyOptional()) const override;
+
     bool IsBatchNormalizationSupported(const TensorInfo& input,
                                        const TensorInfo& output,
                                        const TensorInfo& mean,
diff --git a/src/backends/reference/RefWorkloadFactory.cpp b/src/backends/reference/RefWorkloadFactory.cpp
index 055c8da600..346fd691f2 100644
--- a/src/backends/reference/RefWorkloadFactory.cpp
+++ b/src/backends/reference/RefWorkloadFactory.cpp
@@ -134,6 +134,12 @@ std::unique_ptr<IWorkload> RefWorkloadFactory::CreateActivation(const Activation
     return std::make_unique<RefActivationWorkload>(descriptor, info);
 }
 
+std::unique_ptr<IWorkload> RefWorkloadFactory::CreateArgMinMax(const ArgMinMaxQueueDescriptor& descriptor,
+                                                               const WorkloadInfo&              info) const
+{
+    return std::make_unique<RefArgMinMaxWorkload>(descriptor, info);
+}
+
 std::unique_ptr<IWorkload> RefWorkloadFactory::CreateSoftmax(const SoftmaxQueueDescriptor& descriptor,
                                                              const WorkloadInfo&           info) const
 {
diff --git a/src/backends/reference/RefWorkloadFactory.hpp b/src/backends/reference/RefWorkloadFactory.hpp
index 2c40053f73..606da82d32 100644
--- a/src/backends/reference/RefWorkloadFactory.hpp
+++ b/src/backends/reference/RefWorkloadFactory.hpp
@@ -67,6 +67,9 @@ public:
     std::unique_ptr<IWorkload> CreateActivation(const ActivationQueueDescriptor& descriptor,
                                                 const WorkloadInfo& info) const override;
 
+    std::unique_ptr<IWorkload> CreateArgMinMax(const ArgMinMaxQueueDescriptor& descriptor,
+                                               const WorkloadInfo& info) const override;
+
     std::unique_ptr<IWorkload> CreateSoftmax(const SoftmaxQueueDescriptor& descriptor,
                                              const WorkloadInfo& info) const override;
 
diff --git a/src/backends/reference/backend.mk b/src/backends/reference/backend.mk
index b2ec7488e2..fd0df27907 100644
--- a/src/backends/reference/backend.mk
+++ b/src/backends/reference/backend.mk
@@ -23,6 +23,7 @@ BACKEND_SOURCES := \
         RefTensorHandleFactory.cpp \
         workloads/Abs.cpp \
         workloads/Activation.cpp \
+        workloads/ArgMinMax.cpp \
         workloads/BatchNormImpl.cpp \
         workloads/BatchToSpaceNd.cpp \
         workloads/Broadcast.cpp \
@@ -40,6 +41,7 @@ BACKEND_SOURCES := \
         workloads/PreluImpl.cpp \
         workloads/RefAbsWorkload.cpp \
         workloads/RefActivationWorkload.cpp \
+        workloads/RefArgMinMaxWorkload.cpp \
         workloads/RefBatchNormalizationWorkload.cpp \
         workloads/RefBatchToSpaceNdWorkload.cpp \
         workloads/RefConcatWorkload.cpp \
diff --git a/src/backends/reference/test/RefLayerTests.cpp b/src/backends/reference/test/RefLayerTests.cpp
index afeb8a458a..a5164f072b 100644
--- a/src/backends/reference/test/RefLayerTests.cpp
+++ b/src/backends/reference/test/RefLayerTests.cpp
@@ -1022,6 +1022,22 @@ ARMNN_AUTO_TEST_CASE(MeanVts3QuantisedSymm16, MeanVts3Test<DataType::QuantisedSy
 
 ARMNN_AUTO_TEST_CASE(AdditionAfterMaxPool, AdditionAfterMaxPoolTest)
 
+// ArgMinMax
+ARMNN_AUTO_TEST_CASE(ArgMaxFloat32, ArgMaxSimpleTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMinFloat32, ArgMinSimpleTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMinChannel4dFloat32, ArgMinChannel4dTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMaxChannel4dFloat32, ArgMaxChannel4dTest<DataType::Float32>)
+
+ARMNN_AUTO_TEST_CASE(ArgMaxQuantisedAsymm8, ArgMaxSimpleTest<DataType::QuantisedAsymm8>)
+ARMNN_AUTO_TEST_CASE(ArgMinQuantisedAsymm8, ArgMinSimpleTest<DataType::QuantisedAsymm8>)
+ARMNN_AUTO_TEST_CASE(ArgMinChannel4dQuantisedAsymm8, ArgMinChannel4dTest<DataType::QuantisedAsymm8>)
+ARMNN_AUTO_TEST_CASE(ArgMaxChannel4dQuantisedAsymm8, ArgMaxChannel4dTest<DataType::QuantisedAsymm8>)
+
+ARMNN_AUTO_TEST_CASE(ArgMaxQuantisedSymm16, ArgMaxSimpleTest<DataType::QuantisedSymm16>)
+ARMNN_AUTO_TEST_CASE(ArgMinQuantisedSymm16, ArgMinSimpleTest<DataType::QuantisedSymm16>)
+ARMNN_AUTO_TEST_CASE(ArgMinChannel4dQuantisedSymm16, ArgMinChannel4dTest<DataType::QuantisedSymm16>)
+ARMNN_AUTO_TEST_CASE(ArgMaxChannel4dQuantisedSymm16, ArgMaxChannel4dTest<DataType::QuantisedSymm16>)
+
 // Space To Batch Nd
 ARMNN_AUTO_TEST_CASE(SpaceToBatchNdSimpleFloat32, SpaceToBatchNdSimpleFloat32Test)
 ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiChannelsFloat32, SpaceToBatchNdMultiChannelsFloat32Test)
diff --git a/src/backends/reference/workloads/CMakeLists.txt b/src/backends/reference/workloads/CMakeLists.txt
index 30770956ba..83444eda1f 100644
--- a/src/backends/reference/workloads/CMakeLists.txt
+++ b/src/backends/reference/workloads/CMakeLists.txt
@@ -6,6 +6,8 @@
 list(APPEND armnnRefBackendWorkloads_sources
     Abs.cpp
     Abs.hpp
+    ArgMinMax.cpp
+    ArgMinMax.hpp
     Activation.cpp
     Activation.hpp
     ArgMinMax.cpp
@@ -49,6 +51,8 @@ list(APPEND armnnRefBackendWorkloads_sources
     RefAbsWorkload.hpp
     RefActivationWorkload.cpp
     RefActivationWorkload.hpp
+    RefArgMinMaxWorkload.cpp
+    RefArgMinMaxWorkload.hpp
     RefBatchNormalizationWorkload.cpp
     RefBatchNormalizationWorkload.hpp
     RefBatchToSpaceNdWorkload.cpp
diff --git a/src/backends/reference/workloads/RefArgMinMaxWorkload.cpp b/src/backends/reference/workloads/RefArgMinMaxWorkload.cpp
new file mode 100644
index 0000000000..5f1eb73b61
--- /dev/null
+++ b/src/backends/reference/workloads/RefArgMinMaxWorkload.cpp
@@ -0,0 +1,38 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "RefArgMinMaxWorkload.hpp"
+
+#include "ArgMinMax.hpp"
+#include "RefWorkloadUtils.hpp"
+#include "Decoders.hpp"
+#include "Encoders.hpp"
+#include "Profiling.hpp"
+
+namespace armnn
+{
+RefArgMinMaxWorkload::RefArgMinMaxWorkload(
+        const ArgMinMaxQueueDescriptor& descriptor,
+        const WorkloadInfo& info)
+        : BaseWorkload<ArgMinMaxQueueDescriptor>(descriptor, info) {}
+
+void RefArgMinMaxWorkload::Execute() const
+{
+    ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, "RefArgMinMaxWorkload_Execute");
+
+    const TensorInfo &inputTensorInfo = GetTensorInfo(m_Data.m_Inputs[0]);
+
+    std::unique_ptr<Decoder<float>> decoderPtr = MakeDecoder<float>(inputTensorInfo, m_Data.m_Inputs[0]->Map());
+    Decoder<float> &decoder = *decoderPtr;
+
+    const TensorInfo &outputTensorInfo = GetTensorInfo(m_Data.m_Outputs[0]);
+
+    int32_t* output = GetOutputTensorData<int32_t>(0, m_Data);
+
+    ArgMinMax(decoder, output, inputTensorInfo, outputTensorInfo, m_Data.m_Parameters.m_Function,
+              m_Data.m_Parameters.m_Axis);
+}
+
+} //namespace armnn
\ No newline at end of file
diff --git a/src/backends/reference/workloads/RefArgMinMaxWorkload.hpp b/src/backends/reference/workloads/RefArgMinMaxWorkload.hpp
new file mode 100644
index 0000000000..97b70772d1
--- /dev/null
+++ b/src/backends/reference/workloads/RefArgMinMaxWorkload.hpp
@@ -0,0 +1,21 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include <backendsCommon/Workload.hpp>
+#include <backendsCommon/WorkloadData.hpp>
+
+namespace armnn
+{
+class RefArgMinMaxWorkload : public BaseWorkload<ArgMinMaxQueueDescriptor>
+{
+public:
+    explicit RefArgMinMaxWorkload(const ArgMinMaxQueueDescriptor& descriptor,
+                                  const WorkloadInfo& info);
+
+    virtual void Execute() const override;
+};
+} //namespace armnn
\ No newline at end of file
diff --git a/src/backends/reference/workloads/RefWorkloads.hpp b/src/backends/reference/workloads/RefWorkloads.hpp
index 959226adf6..b4721b11aa 100644
--- a/src/backends/reference/workloads/RefWorkloads.hpp
+++ b/src/backends/reference/workloads/RefWorkloads.hpp
@@ -7,6 +7,7 @@
 
 #include "Abs.hpp"
 #include "Activation.hpp"
+#include "ArgMinMax.hpp"
 #include "BatchNormImpl.hpp"
 #include "ConvImpl.hpp"
 #include "Concatenate.hpp"
@@ -16,6 +17,7 @@
 #include "Pooling2d.hpp"
 #include "RefAbsWorkload.hpp"
 #include "RefActivationWorkload.hpp"
+#include "RefArgMinMaxWorkload.hpp"
 #include "RefBatchNormalizationWorkload.hpp"
 #include "RefBatchToSpaceNdWorkload.hpp"
 #include "RefConvolution2dWorkload.hpp"