IVGCVSW-4374 Add NEON backend support for SpaceToBatchNd

 * Added NeonSpaceToBatchNdWorkload class.
 * Added CreateSpaceToBatchNd implementation to NeonWorkloadFactory.
 * Added IsSpaceToBatchNdSupported implementation to NeonLayerSupport.
 * Enabled SpaceToBatchNd tests on Neon backend.

Signed-off-by: Mike Kelly <mike.kelly@arm.com>
Change-Id: Ic7d469a380482ec9febd44450f72dd7e53cda03f

10 files changed, 175 insertions, 5 deletions

diff --git a/src/backends/neon/NeonLayerSupport.cpp b/src/backends/neon/NeonLayerSupport.cpp
index fe9bffbee9..47bcc2e795 100644
--- a/src/backends/neon/NeonLayerSupport.cpp
+++ b/src/backends/neon/NeonLayerSupport.cpp
@@ -50,6 +50,7 @@
 #include "workloads/NeonRsqrtWorkload.hpp"
 #include "workloads/NeonSliceWorkload.hpp"
 #include "workloads/NeonSoftmaxBaseWorkload.hpp"
+#include "workloads/NeonSpaceToBatchNdWorkload.hpp"
 #include "workloads/NeonSpaceToDepthWorkload.hpp"
 #include "workloads/NeonSplitterWorkload.hpp"
 #include "workloads/NeonStackWorkload.hpp"
@@ -701,10 +702,22 @@ bool NeonLayerSupport::IsSoftmaxSupported(const TensorInfo& input,
     FORWARD_WORKLOAD_VALIDATE_FUNC(NeonSoftmaxWorkloadValidate, reasonIfUnsupported, input, output, descriptor);
 }
 
+bool NeonLayerSupport::IsSpaceToBatchNdSupported(const TensorInfo& input,
+                                                 const TensorInfo& output,
+                                                 const SpaceToBatchNdDescriptor& descriptor,
+                                                 Optional<std::string&> reasonIfUnsupported) const
+{
+    FORWARD_WORKLOAD_VALIDATE_FUNC(NeonSpaceToBatchNdWorkloadValidate,
+                                   reasonIfUnsupported,
+                                   input,
+                                   output,
+                                   descriptor);
+}
+
 bool NeonLayerSupport::IsSpaceToDepthSupported(const TensorInfo& input,
-                                              const TensorInfo& output,
-                                              const SpaceToDepthDescriptor& descriptor,
-                                              Optional<std::string&> reasonIfUnsupported) const
+                                               const TensorInfo& output,
+                                               const SpaceToDepthDescriptor& descriptor,
+                                               Optional<std::string&> reasonIfUnsupported) const
 {
     FORWARD_WORKLOAD_VALIDATE_FUNC(NeonSpaceToDepthWorkloadValidate,
                                    reasonIfUnsupported,
diff --git a/src/backends/neon/NeonLayerSupport.hpp b/src/backends/neon/NeonLayerSupport.hpp
index d429aeceec..e49e78b3d2 100644
--- a/src/backends/neon/NeonLayerSupport.hpp
+++ b/src/backends/neon/NeonLayerSupport.hpp
@@ -245,6 +245,11 @@ public:
                             const SoftmaxDescriptor& descriptor,
                             Optional<std::string&> reasonIfUnsupported = EmptyOptional()) const override;
 
+    bool IsSpaceToBatchNdSupported(const TensorInfo& input,
+                                   const TensorInfo& output,
+                                   const SpaceToBatchNdDescriptor& descriptor,
+                                   Optional<std::string&> reasonIfUnsupported = EmptyOptional()) const override;
+
     bool IsSpaceToDepthSupported(const TensorInfo& input,
                                  const TensorInfo& output,
                                  const SpaceToDepthDescriptor& descriptor,
diff --git a/src/backends/neon/NeonWorkloadFactory.cpp b/src/backends/neon/NeonWorkloadFactory.cpp
index 82f9bdb924..649cb9f699 100644
--- a/src/backends/neon/NeonWorkloadFactory.cpp
+++ b/src/backends/neon/NeonWorkloadFactory.cpp
@@ -467,8 +467,14 @@ std::unique_ptr<IWorkload> NeonWorkloadFactory::CreateSoftmax(const SoftmaxQueue
         descriptor, info, m_MemoryManager->GetIntraLayerManager());
 }
 
-std::unique_ptr<IWorkload> NeonWorkloadFactory::CreateSpaceToDepth(
-    const armnn::SpaceToDepthQueueDescriptor& descriptor, const armnn::WorkloadInfo& info) const
+std::unique_ptr<IWorkload> NeonWorkloadFactory::CreateSpaceToBatchNd(const SpaceToBatchNdQueueDescriptor& descriptor,
+                                                                     const WorkloadInfo& info) const
+{
+    return std::make_unique<NeonSpaceToBatchNdWorkload>(descriptor, info);
+}
+
+std::unique_ptr<IWorkload> NeonWorkloadFactory::CreateSpaceToDepth(const SpaceToDepthQueueDescriptor& descriptor,
+                                                                   const WorkloadInfo& info) const
 {
     return std::make_unique<NeonSpaceToDepthWorkload>(descriptor, info);
 }
diff --git a/src/backends/neon/NeonWorkloadFactory.hpp b/src/backends/neon/NeonWorkloadFactory.hpp
index 44c0629ece..95271e200d 100644
--- a/src/backends/neon/NeonWorkloadFactory.hpp
+++ b/src/backends/neon/NeonWorkloadFactory.hpp
@@ -195,6 +195,9 @@ public:
     std::unique_ptr<IWorkload> CreateSoftmax(const SoftmaxQueueDescriptor& descriptor,
                                              const WorkloadInfo& info) const override;
 
+    std::unique_ptr<IWorkload> CreateSpaceToBatchNd(const SpaceToBatchNdQueueDescriptor& descriptor,
+                                                    const WorkloadInfo& info) const override;
+
     std::unique_ptr<IWorkload> CreateSpaceToDepth(const SpaceToDepthQueueDescriptor& descriptor,
                                                   const WorkloadInfo& info) const override;
 
diff --git a/src/backends/neon/backend.mk b/src/backends/neon/backend.mk
index 3ddc79ada3..05dcd02c9a 100644
--- a/src/backends/neon/backend.mk
+++ b/src/backends/neon/backend.mk
@@ -59,6 +59,7 @@ BACKEND_SOURCES := \
         workloads/NeonSoftmaxBaseWorkload.cpp \
         workloads/NeonSoftmaxFloatWorkload.cpp \
         workloads/NeonSoftmaxUint8Workload.cpp \
+        workloads/NeonSpaceToBatchNdWorkload.cpp \
         workloads/NeonSpaceToDepthWorkload.cpp \
         workloads/NeonSplitterWorkload.cpp \
         workloads/NeonStackWorkload.cpp \
diff --git a/src/backends/neon/test/NeonLayerTests.cpp b/src/backends/neon/test/NeonLayerTests.cpp
index cde1435045..f14b2a446b 100644
--- a/src/backends/neon/test/NeonLayerTests.cpp
+++ b/src/backends/neon/test/NeonLayerTests.cpp
@@ -421,6 +421,27 @@ ARMNN_AUTO_TEST_CASE(Simple3dSoftmaxBeta1Uint8, Simple3dSoftmaxUint8Test, 1.0f)
 ARMNN_AUTO_TEST_CASE(Simple4dSoftmaxBeta1, Simple4dSoftmaxTest, 1.0f)
 ARMNN_AUTO_TEST_CASE(Simple4dSoftmaxBeta1Uint8, Simple4dSoftmaxUint8Test, 1.0f)
 
+// Space To Batch Nd
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdSimpleFloat32, SpaceToBatchNdSimpleFloat32Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiChannelsFloat32, SpaceToBatchNdMultiChannelsFloat32Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiBlockFloat32, SpaceToBatchNdMultiBlockFloat32Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdPaddingFloat32, SpaceToBatchNdPaddingFloat32Test)
+
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdSimpleUint8, SpaceToBatchNdSimpleUint8Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiChannelsUint8, SpaceToBatchNdMultiChannelsUint8Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiBlockUint8, SpaceToBatchNdMultiBlockUint8Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdPaddingUint8, SpaceToBatchNdPaddingUint8Test)
+
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdSimpleNhwcFloat32, SpaceToBatchNdSimpleNhwcFloat32Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiChannelsNhwcFloat32, SpaceToBatchNdMultiChannelsNhwcFloat32Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiBlockNhwcFloat32, SpaceToBatchNdMultiBlockNhwcFloat32Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdPaddingNhwcFloat32, SpaceToBatchNdPaddingNhwcFloat32Test)
+
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdSimpleNhwcUint8, SpaceToBatchNdSimpleNhwcUint8Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiChannelsNhwcUint8, SpaceToBatchNdMultiChannelsNhwcUint8Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdMultiBlockNhwcUint8, SpaceToBatchNdMultiBlockNhwcUint8Test)
+ARMNN_AUTO_TEST_CASE(SpaceToBatchNdPaddingNhwcUint8, SpaceToBatchNdPaddingNhwcUint8Test)
+
 // SpaceToDepth
 ARMNN_AUTO_TEST_CASE(SpaceToDepthNchwAsymmQ8, SpaceToDepthNchwAsymmQ8Test)
 ARMNN_AUTO_TEST_CASE(SpaceToDepthNhwcsymmQ8, SpaceToDepthNhwcAsymmQ8Test)
diff --git a/src/backends/neon/workloads/CMakeLists.txt b/src/backends/neon/workloads/CMakeLists.txt
index 44db6d263a..d2c549ee25 100644
--- a/src/backends/neon/workloads/CMakeLists.txt
+++ b/src/backends/neon/workloads/CMakeLists.txt
@@ -80,6 +80,8 @@ list(APPEND armnnNeonBackendWorkloads_sources
     NeonSoftmaxFloatWorkload.hpp
     NeonSoftmaxUint8Workload.cpp
     NeonSoftmaxUint8Workload.hpp
+    NeonSpaceToBatchNdWorkload.cpp
+    NeonSpaceToBatchNdWorkload.hpp
     NeonSpaceToDepthWorkload.cpp
     NeonSpaceToDepthWorkload.hpp
     NeonSplitterWorkload.cpp
diff --git a/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.cpp b/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.cpp
new file mode 100644
index 0000000000..199e926142
--- /dev/null
+++ b/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.cpp
@@ -0,0 +1,83 @@
+//
+// Copyright © 2020 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "NeonSpaceToBatchNdWorkload.hpp"
+
+#include "NeonWorkloadUtils.hpp"
+#include <ResolveType.hpp>
+
+namespace armnn
+{
+
+using namespace armcomputetensorutils;
+
+arm_compute::Status NeonSpaceToBatchNdWorkloadValidate(const TensorInfo& input,
+                                                       const TensorInfo& output,
+                                                       const SpaceToBatchNdDescriptor& descriptor)
+{
+    const arm_compute::TensorInfo aclInputInfo  = BuildArmComputeTensorInfo(input, descriptor.m_DataLayout);
+    const arm_compute::TensorInfo aclOutputInfo = BuildArmComputeTensorInfo(output, descriptor.m_DataLayout);
+
+    // ArmNN blockShape is [H, W] Cl asks for W, H
+    int32_t blockHeight = boost::numeric_cast<int32_t>(descriptor.m_BlockShape[0]);
+    int32_t blockWidth  = boost::numeric_cast<int32_t>(descriptor.m_BlockShape[1]);
+
+    arm_compute::Size2D paddingLeftTop = BuildArmComputeSize2D(
+            descriptor.m_PadList[1].first, descriptor.m_PadList[0].first);
+    arm_compute::Size2D paddingRightBottom  = BuildArmComputeSize2D(
+            descriptor.m_PadList[1].second, descriptor.m_PadList[0].second);
+
+    return arm_compute::NESpaceToBatchLayer::validate(&aclInputInfo,
+                                                      blockWidth,
+                                                      blockHeight,
+                                                      paddingLeftTop,
+                                                      paddingRightBottom,
+                                                      &aclOutputInfo);
+}
+
+NeonSpaceToBatchNdWorkload::NeonSpaceToBatchNdWorkload(const SpaceToBatchNdQueueDescriptor& desc,
+                                                       const WorkloadInfo& info)
+        : BaseWorkload<SpaceToBatchNdQueueDescriptor>(desc, info)
+{
+    m_Data.ValidateInputsOutputs("NESpaceToBatchNdWorkload", 1, 1);
+
+    arm_compute::ITensor& input  =
+            boost::polymorphic_pointer_downcast<IAclTensorHandle>(m_Data.m_Inputs[0])->GetTensor();
+    arm_compute::ITensor& output =
+            boost::polymorphic_pointer_downcast<IAclTensorHandle>(m_Data.m_Outputs[0])->GetTensor();
+
+    // ArmNN blockShape is [H, W] Cl asks for W, H
+    int32_t blockHeight = boost::numeric_cast<int32_t>(m_Data.m_Parameters.m_BlockShape[0]);
+    int32_t blockWidth  = boost::numeric_cast<int32_t>(m_Data.m_Parameters.m_BlockShape[1]);
+
+    arm_compute::Size2D paddingLeftTop = BuildArmComputeSize2D(
+            m_Data.m_Parameters.m_PadList[1].first, m_Data.m_Parameters.m_PadList[0].first);
+    arm_compute::Size2D paddingRightBottom  = BuildArmComputeSize2D(
+            m_Data.m_Parameters.m_PadList[1].second, m_Data.m_Parameters.m_PadList[0].second);
+
+    arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout);
+    input.info()->set_data_layout(aclDataLayout);
+    output.info()->set_data_layout(aclDataLayout);
+
+    m_Layer.reset(new arm_compute::NESpaceToBatchLayer());
+    m_Layer->configure(&input,
+                       blockWidth,
+                       blockHeight,
+                       paddingLeftTop,
+                       paddingRightBottom,
+                       &output);
+    m_Layer->prepare();
+}
+
+void NeonSpaceToBatchNdWorkload::Execute() const
+{
+    if (m_Layer)
+    {
+        ARMNN_SCOPED_PROFILING_EVENT_NEON("NeonSpaceToBatchNdWorkload_Execute");
+        m_Layer->run();
+    }
+}
+
+} //namespace armnn
\ No newline at end of file
diff --git a/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp b/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp
new file mode 100644
index 0000000000..feb8dba16f
--- /dev/null
+++ b/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp
@@ -0,0 +1,35 @@
+//
+// Copyright © 2020 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include <armnn/Tensor.hpp>
+#include <armnn/Descriptors.hpp>
+
+#include <backendsCommon/Workload.hpp>
+
+#include <arm_compute/runtime/NEON/functions/NESpaceToBatchLayer.h>
+
+namespace armnn
+{
+
+arm_compute::Status NeonSpaceToBatchNdWorkloadValidate(const TensorInfo& input,
+                                                       const TensorInfo& output,
+                                                       const SpaceToBatchNdDescriptor& descriptor);
+
+class NeonSpaceToBatchNdWorkload : public BaseWorkload<SpaceToBatchNdQueueDescriptor>
+{
+public:
+    using BaseWorkload<SpaceToBatchNdQueueDescriptor>::BaseWorkload;
+
+    NeonSpaceToBatchNdWorkload(const SpaceToBatchNdQueueDescriptor& descriptor, const WorkloadInfo& info);
+
+    virtual void Execute() const override;
+
+private:
+    mutable std::unique_ptr<arm_compute::NESpaceToBatchLayer> m_Layer;
+};
+
+} //namespace armnn
\ No newline at end of file
diff --git a/src/backends/neon/workloads/NeonWorkloads.hpp b/src/backends/neon/workloads/NeonWorkloads.hpp
index dc9bef31b2..104504e097 100644
--- a/src/backends/neon/workloads/NeonWorkloads.hpp
+++ b/src/backends/neon/workloads/NeonWorkloads.hpp
@@ -41,6 +41,7 @@
 #include "NeonSliceWorkload.hpp"
 #include "NeonSoftmaxFloatWorkload.hpp"
 #include "NeonSoftmaxUint8Workload.hpp"
+#include "NeonSpaceToBatchNdWorkload.hpp"
 #include "NeonSpaceToDepthWorkload.hpp"
 #include "NeonSplitterWorkload.hpp"
 #include "NeonStackWorkload.hpp"