From 22f0f2bc3d8d8e460200ae353dcea42157d9ad39 Mon Sep 17 00:00:00 2001
From: nikraj01 <nikhil.raj@arm.com>
Date: Thu, 30 May 2019 17:29:32 +0100
Subject: IVGCVSW-3194 Refactor SpaceToBatchNd workload

Change-Id: Iac2ded9b20c37299e3de51465dcbfb5a7bfc52d5
Signed-off-by: nikraj01 <nikhil.raj@arm.com>
---
 .../reference/workloads/SpaceToBatchNd.cpp         | 31 +++++++++++-----------
 1 file changed, 15 insertions(+), 16 deletions(-)

(limited to 'src/backends/reference/workloads/SpaceToBatchNd.cpp')

diff --git a/src/backends/reference/workloads/SpaceToBatchNd.cpp b/src/backends/reference/workloads/SpaceToBatchNd.cpp
index 51e45a8c3f..0bc2396973 100644
--- a/src/backends/reference/workloads/SpaceToBatchNd.cpp
+++ b/src/backends/reference/workloads/SpaceToBatchNd.cpp
@@ -31,12 +31,11 @@ unsigned int GetOffset(const TensorShape& shape,
     }
 }
 
-template<typename T>
 void SpaceToBatchNd(const TensorInfo& inputInfo,
                     const TensorInfo& outputInfo,
                     const SpaceToBatchNdDescriptor& params,
-                    const T* inputData,
-                    T* outputData)
+                    Decoder<float>& inputData,
+                    Encoder<float>& outputData)
 {
     DataLayoutIndexed dataLayout = params.m_DataLayout;
 
@@ -83,7 +82,9 @@ void SpaceToBatchNd(const TensorInfo& inputInfo,
                                                            outW,
                                                            c,
                                                            dataLayout);
-                        outputData[outOffset] = 0;
+                        outputData += outOffset;
+                        outputData.Set(0);
+                        outputData -= outOffset;
                     }
                 }
                 else
@@ -104,7 +105,11 @@ void SpaceToBatchNd(const TensorInfo& inputInfo,
                                                            c,
                                                            dataLayout);
 
-                        outputData[outOffset] = inputData[inOffset];
+                        outputData += outOffset;
+                        inputData += inOffset;
+                        outputData.Set(inputData.Get());
+                        inputData -= inOffset;
+                        outputData -= outOffset;
                     }
                 }
             }
@@ -112,16 +117,10 @@ void SpaceToBatchNd(const TensorInfo& inputInfo,
     }
 }
 
-template void SpaceToBatchNd<float>(const TensorInfo& inputInfo,
-                                    const TensorInfo& outputInfo,
-                                    const SpaceToBatchNdDescriptor& params,
-                                    const float* inputData,
-                                    float* outData);
-
-template void SpaceToBatchNd<uint8_t>(const TensorInfo& inputInfo,
-                                      const TensorInfo& outputInfo,
-                                      const SpaceToBatchNdDescriptor& params,
-                                      const uint8_t* inputData,
-                                      uint8_t* outData);
+void SpaceToBatchNd(const TensorInfo& inputInfo,
+                    const TensorInfo& outputInfo,
+                    const SpaceToBatchNdDescriptor& params,
+                    Decoder<float>& inputData,
+                    Encoder<float>& outData);
 
 } //namespace armnn
-- 
cgit v1.2.1