IVGCVSW-3696 Add NEON ArgMinMax workload and tests

 * Added layer tests and fixed WorkloadData validate.
 * Also enabled copy to/from NEON for Signed32.

Signed-off-by: James Conroy <james.conroy@arm.com>
Change-Id: I5e961f88434e18d5a8ebff956d20a1c2cf1b50bb

11 files changed, 163 insertions, 0 deletions

diff --git a/src/backends/neon/NeonLayerSupport.cpp b/src/backends/neon/NeonLayerSupport.cpp
index b713abaef3..2f3643f210 100644
--- a/src/backends/neon/NeonLayerSupport.cpp
+++ b/src/backends/neon/NeonLayerSupport.cpp
@@ -21,6 +21,7 @@
 #include "workloads/NeonAbsWorkload.hpp"
 #include "workloads/NeonAdditionWorkload.hpp"
 #include "workloads/NeonActivationWorkload.hpp"
+#include "workloads/NeonArgMinMaxWorkload.hpp"
 #include "workloads/NeonBatchNormalizationWorkload.hpp"
 #include "workloads/NeonConvolution2dWorkload.hpp"
 #include "workloads/NeonDepthToSpaceWorkload.hpp"
@@ -146,6 +147,18 @@ bool NeonLayerSupport::IsAdditionSupported(const TensorInfo& input0,
                                    output);
 }
 
+bool NeonLayerSupport::IsArgMinMaxSupported(const TensorInfo& input,
+                                            const TensorInfo& output,
+                                            const ArgMinMaxDescriptor& descriptor,
+                                            Optional<std::string&> reasonIfUnsupported) const
+{
+    FORWARD_WORKLOAD_VALIDATE_FUNC(NeonArgMinMaxWorkloadValidate,
+                                   reasonIfUnsupported,
+                                   input,
+                                   output,
+                                   descriptor);
+}
+
 bool NeonLayerSupport::IsBatchNormalizationSupported(const TensorInfo& input,
                                                      const TensorInfo& output,
                                                      const TensorInfo& mean,
diff --git a/src/backends/neon/NeonLayerSupport.hpp b/src/backends/neon/NeonLayerSupport.hpp
index 31f9e576d7..76eb342eb9 100644
--- a/src/backends/neon/NeonLayerSupport.hpp
+++ b/src/backends/neon/NeonLayerSupport.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/neon/NeonTensorHandle.hpp b/src/backends/neon/NeonTensorHandle.hpp
index 37013eb92d..ca5bfb04b1 100644
--- a/src/backends/neon/NeonTensorHandle.hpp
+++ b/src/backends/neon/NeonTensorHandle.hpp
@@ -186,6 +186,10 @@ private:
                 armcomputetensorutils::CopyArmComputeITensorData(this->GetTensor(),
                                                                  static_cast<int16_t*>(memory));
                 break;
+            case arm_compute::DataType::S32:
+                armcomputetensorutils::CopyArmComputeITensorData(this->GetTensor(),
+                                                                 static_cast<int32_t*>(memory));
+                break;
             default:
             {
                 throw armnn::UnimplementedException();
@@ -216,6 +220,10 @@ private:
                 armcomputetensorutils::CopyArmComputeITensorData(static_cast<const int16_t*>(memory),
                                                                  this->GetTensor());
                 break;
+            case arm_compute::DataType::S32:
+                armcomputetensorutils::CopyArmComputeITensorData(static_cast<const int32_t*>(memory),
+                                                                 this->GetTensor());
+                break;
             default:
             {
                 throw armnn::UnimplementedException();
@@ -292,6 +300,10 @@ private:
                 armcomputetensorutils::CopyArmComputeITensorData(this->GetTensor(),
                                                                  static_cast<int16_t*>(memory));
                 break;
+            case arm_compute::DataType::S32:
+                armcomputetensorutils::CopyArmComputeITensorData(this->GetTensor(),
+                                                                 static_cast<int32_t*>(memory));
+                break;
             default:
             {
                 throw armnn::UnimplementedException();
@@ -318,6 +330,10 @@ private:
                 armcomputetensorutils::CopyArmComputeITensorData(static_cast<const int16_t*>(memory),
                                                                  this->GetTensor());
                 break;
+            case arm_compute::DataType::S32:
+                armcomputetensorutils::CopyArmComputeITensorData(static_cast<const int32_t*>(memory),
+                                                                 this->GetTensor());
+                break;
             default:
             {
                 throw armnn::UnimplementedException();
diff --git a/src/backends/neon/NeonWorkloadFactory.cpp b/src/backends/neon/NeonWorkloadFactory.cpp
index f0b738c350..3492923f72 100644
--- a/src/backends/neon/NeonWorkloadFactory.cpp
+++ b/src/backends/neon/NeonWorkloadFactory.cpp
@@ -442,4 +442,10 @@ std::unique_ptr<IWorkload> NeonWorkloadFactory::CreateStack(const StackQueueDesc
     return std::make_unique<NeonStackWorkload>(descriptor, info);
 }
 
+std::unique_ptr<IWorkload> NeonWorkloadFactory::CreateArgMinMax(const ArgMinMaxQueueDescriptor& descriptor,
+                                                                const WorkloadInfo& info) const
+{
+    return std::make_unique<NeonArgMinMaxWorkload>(descriptor, info);
+}
+
 } // namespace armnn
diff --git a/src/backends/neon/NeonWorkloadFactory.hpp b/src/backends/neon/NeonWorkloadFactory.hpp
index 4bdbc8e859..aad9cf949d 100644
--- a/src/backends/neon/NeonWorkloadFactory.hpp
+++ b/src/backends/neon/NeonWorkloadFactory.hpp
@@ -200,6 +200,9 @@ public:
     std::unique_ptr<IWorkload> CreateStack(const StackQueueDescriptor& descriptor,
                                            const WorkloadInfo& info) const override;
 
+    std::unique_ptr<IWorkload> CreateArgMinMax(const ArgMinMaxQueueDescriptor& descriptor,
+                                               const WorkloadInfo& info) const override;
+
 private:
     mutable std::shared_ptr<NeonMemoryManager> m_MemoryManager;
 };
diff --git a/src/backends/neon/backend.mk b/src/backends/neon/backend.mk
index 3e46387ae4..fb10a0dcef 100644
--- a/src/backends/neon/backend.mk
+++ b/src/backends/neon/backend.mk
@@ -24,6 +24,7 @@ BACKEND_SOURCES := \
         workloads/NeonAbsWorkload.cpp \
         workloads/NeonActivationWorkload.cpp \
         workloads/NeonAdditionWorkload.cpp \
+        workloads/NeonArgMinMaxWorkload.cpp \
         workloads/NeonBatchNormalizationWorkload.cpp \
         workloads/NeonConcatWorkload.cpp \
         workloads/NeonConstantWorkload.cpp \
diff --git a/src/backends/neon/test/NeonLayerTests.cpp b/src/backends/neon/test/NeonLayerTests.cpp
index 38771833bd..0d1faa9fc5 100644
--- a/src/backends/neon/test/NeonLayerTests.cpp
+++ b/src/backends/neon/test/NeonLayerTests.cpp
@@ -910,6 +910,14 @@ ARMNN_AUTO_TEST_CASE(Rsqrt3d, Rsqrt3dTest<DataType::Float32>)
 ARMNN_AUTO_TEST_CASE(RsqrtZero, RsqrtZeroTest<DataType::Float32>)
 ARMNN_AUTO_TEST_CASE(RsqrtNegative, RsqrtNegativeTest<DataType::Float32>)
 
+// ArgMinMax
+ARMNN_AUTO_TEST_CASE(ArgMinFloat32, ArgMinSimpleTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMaxFloat32, ArgMaxSimpleTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMinChannel, ArgMinChannelTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMaxChannel, ArgMaxChannelTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMaxHeight, ArgMaxHeightTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE(ArgMinWidth, ArgMinWidthTest<DataType::Float32>)
+
 #if defined(ARMNNREF_ENABLED)
 
 // The ARMNN_COMPARE_REF_AUTO_TEST_CASE and the ARMNN_COMPARE_REF_FIXTURE_TEST_CASE test units are not available
diff --git a/src/backends/neon/workloads/CMakeLists.txt b/src/backends/neon/workloads/CMakeLists.txt
index 42ac641369..f8d5922b70 100644
--- a/src/backends/neon/workloads/CMakeLists.txt
+++ b/src/backends/neon/workloads/CMakeLists.txt
@@ -10,6 +10,8 @@ list(APPEND armnnNeonBackendWorkloads_sources
     NeonActivationWorkload.hpp
     NeonAdditionWorkload.cpp
     NeonAdditionWorkload.hpp
+    NeonArgMinMaxWorkload.cpp
+    NeonArgMinMaxWorkload.hpp
     NeonBatchNormalizationWorkload.cpp
     NeonBatchNormalizationWorkload.hpp
     NeonConcatWorkload.cpp
diff --git a/src/backends/neon/workloads/NeonArgMinMaxWorkload.cpp b/src/backends/neon/workloads/NeonArgMinMaxWorkload.cpp
new file mode 100644
index 0000000000..e8d537f2ef
--- /dev/null
+++ b/src/backends/neon/workloads/NeonArgMinMaxWorkload.cpp
@@ -0,0 +1,79 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "NeonArgMinMaxWorkload.hpp"
+#include "NeonWorkloadUtils.hpp"
+
+#include <aclCommon/ArmComputeTensorUtils.hpp>
+#include <backendsCommon/CpuTensorHandle.hpp>
+#include <TensorUtils.hpp>
+
+#include <arm_compute/runtime/NEON/functions/NEArgMinMaxLayer.h>
+
+namespace
+{
+unsigned int CalcAclAxis(unsigned int numDimensions, unsigned int axisIndex)
+{
+    return (numDimensions - axisIndex) - 1;
+}
+
+} //namespace
+
+namespace armnn
+{
+
+arm_compute::Status NeonArgMinMaxWorkloadValidate(const TensorInfo& input,
+                                                  const TensorInfo& output,
+                                                  const ArgMinMaxDescriptor& descriptor)
+{
+    const arm_compute::TensorInfo aclInput = armcomputetensorutils::BuildArmComputeTensorInfo(input);
+    const arm_compute::TensorInfo aclOutput = armcomputetensorutils::BuildArmComputeTensorInfo(output);
+
+    auto numDims = input.GetNumDimensions();
+    auto unsignedAxis = armnnUtils::GetUnsignedAxis(numDims, descriptor.m_Axis);
+    int aclAxis = boost::numeric_cast<int>(CalcAclAxis(numDims, unsignedAxis));
+
+    if (descriptor.m_Function == ArgMinMaxFunction::Max)
+    {
+        return arm_compute::NEArgMinMaxLayer::validate(&aclInput, aclAxis, &aclOutput,
+                                                       arm_compute::ReductionOperation::ARG_IDX_MAX);
+    }
+    else
+    {
+        return arm_compute::NEArgMinMaxLayer::validate(&aclInput, aclAxis, &aclOutput,
+                                                       arm_compute::ReductionOperation::ARG_IDX_MIN);
+    }
+}
+
+
+NeonArgMinMaxWorkload::NeonArgMinMaxWorkload(const ArgMinMaxQueueDescriptor& descriptor,
+                                             const WorkloadInfo& info)
+        : BaseWorkload<ArgMinMaxQueueDescriptor>(descriptor, info)
+{
+    arm_compute::ITensor& input = boost::polymorphic_downcast<IAclTensorHandle*>(m_Data.m_Inputs[0])->GetTensor();
+    arm_compute::ITensor& output = boost::polymorphic_downcast<IAclTensorHandle*>(m_Data.m_Outputs[0])->GetTensor();
+
+    auto numDims = info.m_InputTensorInfos[0].GetNumDimensions();
+    auto unsignedAxis = armnnUtils::GetUnsignedAxis(numDims, m_Data.m_Parameters.m_Axis);
+    int aclAxis = boost::numeric_cast<int>(CalcAclAxis(numDims, unsignedAxis));
+
+    if (m_Data.m_Parameters.m_Function == ArgMinMaxFunction::Max)
+    {
+        m_ArgMinMaxLayer.configure(&input, aclAxis, &output, arm_compute::ReductionOperation::ARG_IDX_MAX);
+    }
+    else
+    {
+        m_ArgMinMaxLayer.configure(&input, aclAxis, &output, arm_compute::ReductionOperation::ARG_IDX_MIN);
+    }
+}
+
+void NeonArgMinMaxWorkload::Execute() const
+{
+    ARMNN_SCOPED_PROFILING_EVENT_NEON("NeonArgMinMaxWorkload_Execute");
+    m_ArgMinMaxLayer.run();
+}
+
+} //namespace armnn
+
diff --git a/src/backends/neon/workloads/NeonArgMinMaxWorkload.hpp b/src/backends/neon/workloads/NeonArgMinMaxWorkload.hpp
new file mode 100644
index 0000000000..6301b13718
--- /dev/null
+++ b/src/backends/neon/workloads/NeonArgMinMaxWorkload.hpp
@@ -0,0 +1,29 @@
+//
+// Copyright © 2019 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include <backendsCommon/Workload.hpp>
+
+#include <arm_compute/core/Error.h>
+#include <arm_compute/runtime/NEON/functions/NEArgMinMaxLayer.h>
+namespace armnn
+{
+
+arm_compute::Status NeonArgMinMaxWorkloadValidate(const TensorInfo& input,
+                                                  const TensorInfo& output,
+                                                  const ArgMinMaxDescriptor& descriptor);
+
+class NeonArgMinMaxWorkload : public BaseWorkload<ArgMinMaxQueueDescriptor>
+{
+public:
+    NeonArgMinMaxWorkload(const ArgMinMaxQueueDescriptor& descriptor, const WorkloadInfo& info);
+    virtual void Execute() const override;
+
+private:
+    mutable arm_compute::NEArgMinMaxLayer m_ArgMinMaxLayer;
+};
+
+} //namespace armnn
diff --git a/src/backends/neon/workloads/NeonWorkloads.hpp b/src/backends/neon/workloads/NeonWorkloads.hpp
index 9d35ed42fe..8044a4f004 100644
--- a/src/backends/neon/workloads/NeonWorkloads.hpp
+++ b/src/backends/neon/workloads/NeonWorkloads.hpp
@@ -7,6 +7,7 @@
 #include "NeonAbsWorkload.hpp"
 #include "NeonActivationWorkload.hpp"
 #include "NeonAdditionWorkload.hpp"
+#include "NeonArgMinMaxWorkload.hpp"
 #include "NeonBatchNormalizationWorkload.hpp"
 #include "NeonConstantWorkload.hpp"
 #include "NeonConvertFp16ToFp32Workload.hpp"