IVGCVSW-2682 Add Serializer and Deserializer for BatchToSpaceNd

Change-Id: I8b8ac99b4ecf282b796a8af4b30594ade4e061e6
Signed-off-by: Nattapat Chaimanowong <nattapat.chaimanowong@arm.com>

1 files changed, 29 insertions, 0 deletions

diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index 1c1cefd455..dc25bab0c6 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -141,6 +141,35 @@ void SerializerVisitor::VisitAdditionLayer(const armnn::IConnectableLayer* layer
     CreateAnyLayer(flatBufferAdditionLayer.o, serializer::Layer::Layer_AdditionLayer);
 }
 
+// Build FlatBuffer for BatchToSpaceNd Layer
+void SerializerVisitor::VisitBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer,
+                                                 const armnn::BatchToSpaceNdDescriptor& descriptor,
+                                                 const char* name)
+{
+    // Create FlatBuffer BaseLayer
+    auto flatBufferBaseLayer = CreateLayerBase(layer, serializer::LayerType::LayerType_BatchToSpaceNd);
+
+    std::vector<unsigned int> crops;
+    crops.reserve(descriptor.m_Crops.size() * 2);
+    for (auto& crop : descriptor.m_Crops)
+    {
+        crops.push_back(crop.first);
+        crops.push_back(crop.second);
+    }
+
+    auto flatBufferDescriptor =
+        CreateBatchToSpaceNdDescriptor(m_flatBufferBuilder,
+                                       m_flatBufferBuilder.CreateVector(descriptor.m_BlockShape),
+                                       m_flatBufferBuilder.CreateVector(crops),
+                                       GetFlatBufferDataLayout(descriptor.m_DataLayout));
+
+    auto flatBufferLayer = serializer::CreateBatchToSpaceNdLayer(m_flatBufferBuilder,
+                                                                 flatBufferBaseLayer,
+                                                                 flatBufferDescriptor);
+
+    CreateAnyLayer(flatBufferLayer.o, serializer::Layer::Layer_BatchToSpaceNdLayer);
+}
+
 // Build FlatBuffer for Constant Layer
 void SerializerVisitor::VisitConstantLayer(const armnn::IConnectableLayer* layer,
                                            const armnn::ConstTensor& input,