From 2ea403d130db0d2853d5c43c29b5112893efc2bf Mon Sep 17 00:00:00 2001
From: Teresa Charlin <teresa.charlinreyes@arm.com>
Date: Mon, 19 Jun 2023 12:06:19 +0100
Subject: IVGCVSW-7785 3D tensors in BATCH_TO_SPACE and SPACE_TO_BATCH in
 CpuAcc & GpuAcc

* Add Reshape layers before and after to extend support for 3D tensors, as ACL only supports 4D tensors for those layers
* Add Unit Tests

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I4431185ce3a3b2f595d2a79bdda7095212d1c52d
---
 src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

(limited to 'src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp')

diff --git a/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp b/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp
index 4bd7d2d4a4..35d70d3fed 100644
--- a/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp
+++ b/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp
@@ -1,16 +1,15 @@
 //
-// Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
+// Copyright © 2020-2023 Arm Ltd and Contributors. All rights reserved.
 // SPDX-License-Identifier: MIT
 //
 
 #pragma once
 
-#include <armnn/Tensor.hpp>
-#include <armnn/Descriptors.hpp>
-
 #include "NeonBaseWorkload.hpp"
+#include "NeonWorkloadUtils.hpp"
 
 #include <arm_compute/runtime/NEON/functions/NESpaceToBatchLayer.h>
+#include <arm_compute/runtime/NEON/functions/NEReshapeLayer.h>
 
 namespace armnn
 {
@@ -30,6 +29,10 @@ public:
 
 private:
     mutable std::unique_ptr<arm_compute::NESpaceToBatchLayer> m_Layer;
+    mutable std::unique_ptr<arm_compute::NEReshapeLayer> m_LayerReshapeInput;
+    mutable std::unique_ptr<arm_compute::NEReshapeLayer> m_LayerReshapeOutput;
+    arm_compute::Tensor m_ReshapeInputTensor;
+    arm_compute::Tensor m_ReshapeOutputTensor;
 };
 
 } //namespace armnn
\ No newline at end of file
-- 
cgit v1.2.1