From 262553e0d4e2571c8d7b3c57a7271c23a783e04b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89anna=20=C3=93=20Cath=C3=A1in?=
 <eanna.ocathain@arm.com>
Date: Wed, 14 Nov 2018 11:26:23 +0000
Subject: IVGCVSW-2054: Fixing issue with InferOutputShape implementation in
 BatchToSpaceNdLayer.     * added Unit test

Change-Id: I80f55e8c7afb39d96006c8dd027fc9683ea8182e
---
 src/backends/backendsCommon/test/LayerTests.cpp          | 16 ++++++++++++++++
 src/backends/backendsCommon/test/LayerTests.hpp          |  2 ++
 src/backends/reference/test/RefLayerTests.cpp            |  2 ++
 .../workloads/RefBatchToSpaceNdUint8Workload.cpp         |  3 +--
 4 files changed, 21 insertions(+), 2 deletions(-)

(limited to 'src/backends')

diff --git a/src/backends/backendsCommon/test/LayerTests.cpp b/src/backends/backendsCommon/test/LayerTests.cpp
index 2556a3e1f5..1ce69142df 100755
--- a/src/backends/backendsCommon/test/LayerTests.cpp
+++ b/src/backends/backendsCommon/test/LayerTests.cpp
@@ -6113,3 +6113,19 @@ LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(armnn::IWorkloadFactory
     return BatchToSpaceNdHelper<float, 4, 4>(workloadFactory, armnn::DataLayout::NCHW, inputShape, input, blockShape,
                                              crops, outputShape, expectedOutput);
 }
+
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest1(armnn::IWorkloadFactory& workloadFactory)
+{
+    const unsigned int inputShape[] = {4, 2, 2, 1};
+    const unsigned int outputShape[] = {1, 4, 4, 1};
+
+    std::vector<uint8_t> input({ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 });
+    std::vector<uint8_t> expectedOutput({ 1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16});
+
+    std::vector<unsigned int> blockShape({2, 2});
+    std::vector<std::pair<unsigned int, unsigned int>> crops = {{0, 0}, {0, 0}};
+
+    return BatchToSpaceNdHelper<uint8_t, 4, 4>(workloadFactory, armnn::DataLayout::NHWC, inputShape, input, blockShape,
+                                               crops, outputShape, expectedOutput);
+}
\ No newline at end of file
diff --git a/src/backends/backendsCommon/test/LayerTests.hpp b/src/backends/backendsCommon/test/LayerTests.hpp
index 34d55832da..c38d84fac2 100644
--- a/src/backends/backendsCommon/test/LayerTests.hpp
+++ b/src/backends/backendsCommon/test/LayerTests.hpp
@@ -439,3 +439,5 @@ LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test2(armnn::IWorkloadFactory
 LayerTestResult<float, 4> BatchToSpaceNdNhwcFloat32Test3(armnn::IWorkloadFactory& workloadFactory);
 
 LayerTestResult<float, 4> BatchToSpaceNdNchwFloat32Test1(armnn::IWorkloadFactory &workloadFactory);
+
+LayerTestResult<uint8_t, 4> BatchToSpaceNdNhwcUintTest1(armnn::IWorkloadFactory &workloadFactory);
\ No newline at end of file
diff --git a/src/backends/reference/test/RefLayerTests.cpp b/src/backends/reference/test/RefLayerTests.cpp
index 8bbf11bfad..4ff5cf2a2e 100644
--- a/src/backends/reference/test/RefLayerTests.cpp
+++ b/src/backends/reference/test/RefLayerTests.cpp
@@ -385,4 +385,6 @@ ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcFloat323, BatchToSpaceNdNhwcFloat32Test3)
 
 ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNchwFloat321, BatchToSpaceNdNchwFloat32Test1)
 
+ARMNN_AUTO_TEST_CASE(BatchToSpaceNdNhwcUint1, BatchToSpaceNdNhwcUintTest1)
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/backends/reference/workloads/RefBatchToSpaceNdUint8Workload.cpp b/src/backends/reference/workloads/RefBatchToSpaceNdUint8Workload.cpp
index a66bcd42de..e4e108edb1 100644
--- a/src/backends/reference/workloads/RefBatchToSpaceNdUint8Workload.cpp
+++ b/src/backends/reference/workloads/RefBatchToSpaceNdUint8Workload.cpp
@@ -18,11 +18,10 @@ void RefBatchToSpaceNdUint8Workload::Execute() const
     const TensorInfo& inputInfo = GetTensorInfo(m_Data.m_Inputs[0]);
     const TensorInfo& outputInfo = GetTensorInfo(m_Data.m_Outputs[0]);
     auto dequantizedInputData = Dequantize(GetInputTensorDataU8(0, m_Data), inputInfo);
-    float* outputData = GetOutputTensorDataFloat(0, m_Data);
 
     std::vector<float> results(outputInfo.GetNumElements());
     BatchToSpaceNd(m_Data.m_Parameters.m_DataLayout, inputInfo, outputInfo, m_Data.m_Parameters.m_BlockShape,
-                   m_Data.m_Parameters.m_Crops, dequantizedInputData.data(), outputData);
+                   m_Data.m_Parameters.m_Crops, dequantizedInputData.data(), results.data());
 
     Quantize(GetOutputTensorDataU8(0, m_Data), results.data(), outputInfo);
 }
-- 
cgit v1.2.1