IVGCVSW-3194 Refactor SpaceToBatchNd workload

Change-Id: Iac2ded9b20c37299e3de51465dcbfb5a7bfc52d5
Signed-off-by: nikraj01 <nikhil.raj@arm.com>

5 files changed, 34 insertions, 46 deletions

diff --git a/src/backends/reference/RefWorkloadFactory.cpp b/src/backends/reference/RefWorkloadFactory.cpp
index 71b1cb034d..bf2339756c 100644
--- a/src/backends/reference/RefWorkloadFactory.cpp
+++ b/src/backends/reference/RefWorkloadFactory.cpp
@@ -278,7 +278,11 @@ std::unique_ptr<IWorkload> RefWorkloadFactory::CreateReshape(const ReshapeQueueD
 std::unique_ptr<IWorkload> RefWorkloadFactory::CreateSpaceToBatchNd(const SpaceToBatchNdQueueDescriptor& descriptor,
     const WorkloadInfo& info) const
 {
-    return MakeWorkload<RefSpaceToBatchNdFloat32Workload, RefSpaceToBatchNdUint8Workload>(descriptor, info);
+    if (IsFloat16(info))
+    {
+        return MakeWorkload<NullWorkload, NullWorkload>(descriptor, info);
+    }
+    return std::make_unique<RefSpaceToBatchNdWorkload>(descriptor, info);
 }
 
 std::unique_ptr<IWorkload> RefWorkloadFactory::CreateFloor(const FloorQueueDescriptor& descriptor,
diff --git a/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.cpp b/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.cpp
index 4d2bb05daa..1a721fe389 100644
--- a/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.cpp
+++ b/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.cpp
@@ -12,23 +12,17 @@
 namespace armnn
 {
 
-template<armnn::DataType DataType>
-void RefSpaceToBatchNdWorkload<DataType>::Execute() const
+void RefSpaceToBatchNdWorkload::Execute() const
 {
-    using T = ResolveType<DataType>;
-
-    ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, GetName() + "_Execute");
+    ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, "RefSpaceToBatchNdWorkload_Execute");
 
     const TensorInfo& inputInfo = GetTensorInfo(m_Data.m_Inputs[0]);
-    const TensorInfo& outputInfo = GetTensorInfo(m_Data.m_Outputs[0]);
+    std::unique_ptr<Decoder<float>> decoder = MakeDecoder<float>(inputInfo, m_Data.m_Inputs[0]->Map());
 
-    const T* inputData = GetInputTensorData<T>(0, m_Data);
-    T* outputData = GetOutputTensorData<T>(0, m_Data);
+    const TensorInfo& outputInfo = GetTensorInfo(m_Data.m_Outputs[0]);
+    std::unique_ptr<Encoder<float>> encoder = MakeEncoder<float>(outputInfo, m_Data.m_Outputs[0]->Map());
 
-    SpaceToBatchNd(inputInfo, outputInfo, m_Data.m_Parameters, inputData, outputData);
+    SpaceToBatchNd(inputInfo, outputInfo, m_Data.m_Parameters, *decoder, *encoder);
 }
 
-template class RefSpaceToBatchNdWorkload<DataType::Float32>;
-template class RefSpaceToBatchNdWorkload<DataType::QuantisedAsymm8>;
-
 } //namespace armnn
diff --git a/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.hpp b/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.hpp
index 3a08662ff7..0deb9e1194 100644
--- a/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.hpp
+++ b/src/backends/reference/workloads/RefSpaceToBatchNdWorkload.hpp
@@ -11,23 +11,11 @@
 namespace armnn
 {
 
-template <armnn::DataType DataType>
-class RefSpaceToBatchNdWorkload : public TypedWorkload<SpaceToBatchNdQueueDescriptor, DataType>
+class RefSpaceToBatchNdWorkload : public BaseWorkload<SpaceToBatchNdQueueDescriptor>
 {
 public:
-    static const std::string& GetName()
-    {
-        static const std::string name = std::string("RefSpaceToBatchNd") + GetDataTypeName(DataType) + "Workload";
-        return name;
-    }
-
-    using TypedWorkload<SpaceToBatchNdQueueDescriptor, DataType>::m_Data;
-    using TypedWorkload<SpaceToBatchNdQueueDescriptor, DataType>::TypedWorkload;
-
+    using BaseWorkload<SpaceToBatchNdQueueDescriptor>::BaseWorkload;
     void Execute() const override;
 };
 
-using RefSpaceToBatchNdFloat32Workload = RefSpaceToBatchNdWorkload<DataType::Float32>;
-using RefSpaceToBatchNdUint8Workload = RefSpaceToBatchNdWorkload<DataType::QuantisedAsymm8>;
-
 } //namespace armnn
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
diff --git a/src/backends/reference/workloads/SpaceToBatchNd.hpp b/src/backends/reference/workloads/SpaceToBatchNd.hpp
index e74e457a72..57c9b6bc25 100644
--- a/src/backends/reference/workloads/SpaceToBatchNd.hpp
+++ b/src/backends/reference/workloads/SpaceToBatchNd.hpp
@@ -5,17 +5,20 @@
 
 #pragma once
 
+#include "BaseIterator.hpp"
+#include "Decoders.hpp"
+#include "Encoders.hpp"
+
 #include <armnn/Descriptors.hpp>
 #include "armnn/Tensor.hpp"
 
 namespace armnn
 {
 
-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);
 
 } //namespace armnn