IVGCVSW-7104: BatchMatMul Serializer/Deserializer Support

  * Updated FlatBuffers schema for BatchMatMul layer type
  * Added Serializer and Deserializer implementations for BatchMatMul
  * Added unit tests for BatchMatMul serialization and deserialization
  * Updated CMakeLists and docs

Signed-off-by: Samuel Yap <samuel.yap@arm.com>
Change-Id: Iad63afbd036a3eb648683eb7416a475561aa20cb

3 files changed, 248 insertions, 0 deletions

diff --git a/src/armnnDeserializer/Deserializer.cpp b/src/armnnDeserializer/Deserializer.cpp
index a405cb92a5..702b060512 100644
--- a/src/armnnDeserializer/Deserializer.cpp
+++ b/src/armnnDeserializer/Deserializer.cpp
@@ -214,6 +214,7 @@ m_ParserFunctions(Layer_MAX+1, &IDeserializer::DeserializerImpl::ParseUnsupporte
     m_ParserFunctions[Layer_ActivationLayer]             = &DeserializerImpl::ParseActivation;
     m_ParserFunctions[Layer_AdditionLayer]               = &DeserializerImpl::ParseAdd;
     m_ParserFunctions[Layer_ArgMinMaxLayer]              = &DeserializerImpl::ParseArgMinMax;
+    m_ParserFunctions[Layer_BatchMatMulLayer]            = &DeserializerImpl::ParseBatchMatMul;
     m_ParserFunctions[Layer_BatchToSpaceNdLayer]         = &DeserializerImpl::ParseBatchToSpaceNd;
     m_ParserFunctions[Layer_BatchNormalizationLayer]     = &DeserializerImpl::ParseBatchNormalization;
     m_ParserFunctions[Layer_CastLayer]                   = &DeserializerImpl::ParseCast;
@@ -292,6 +293,8 @@ LayerBaseRawPtr IDeserializer::DeserializerImpl::GetBaseLayer(const GraphPtr& gr
             return graphPtr->layers()->Get(layerIndex)->layer_as_AdditionLayer()->base();
         case Layer::Layer_ArgMinMaxLayer:
             return graphPtr->layers()->Get(layerIndex)->layer_as_ArgMinMaxLayer()->base();
+        case Layer::Layer_BatchMatMulLayer:
+            return graphPtr->layers()->Get(layerIndex)->layer_as_BatchMatMulLayer()->base();
         case Layer::Layer_BatchToSpaceNdLayer:
             return graphPtr->layers()->Get(layerIndex)->layer_as_BatchToSpaceNdLayer()->base();
         case Layer::Layer_BatchNormalizationLayer:
@@ -1258,6 +1261,37 @@ void IDeserializer::DeserializerImpl::ParseArgMinMax(GraphPtr graph, unsigned in
     RegisterOutputSlots(graph, layerIndex, layer);
 }
 
+void IDeserializer::DeserializerImpl::ParseBatchMatMul(GraphPtr graph, unsigned int layerIndex)
+{
+    CHECK_LAYERS(graph, 0, layerIndex);
+
+    auto inputs = GetInputs(graph, layerIndex);
+    CHECK_LOCATION();
+    CHECK_VALID_SIZE(inputs.size(), 2);
+
+    auto outputs = GetOutputs(graph, layerIndex);
+    CHECK_VALID_SIZE(outputs.size(), 1);
+
+    auto serializerLayer = graph->layers()->Get(layerIndex)->layer_as_BatchMatMulLayer();
+    auto serializerDescriptor = serializerLayer->descriptor();
+
+    armnn::BatchMatMulDescriptor descriptor(serializerDescriptor->transposeX(),
+                                            serializerDescriptor->transposeY(),
+                                            serializerDescriptor->adjointX(),
+                                            serializerDescriptor->adjointY(),
+                                            ToDataLayout(serializerDescriptor->dataLayoutX()),
+                                            ToDataLayout(serializerDescriptor->dataLayoutY()));
+
+    auto layerName = GetLayerName(graph, layerIndex);
+    IConnectableLayer* layer = m_Network->AddBatchMatMulLayer(descriptor, layerName.c_str());
+
+    armnn::TensorInfo outputTensorInfo = ToTensorInfo(outputs[0]);
+    layer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo);
+
+    RegisterInputSlots(graph, layerIndex, layer);
+    RegisterOutputSlots(graph, layerIndex, layer);
+}
+
 void IDeserializer::DeserializerImpl::ParseBatchToSpaceNd(GraphPtr graph, unsigned int layerIndex)
 {
     CHECK_LAYERS(graph, 0, layerIndex);
diff --git a/src/armnnDeserializer/Deserializer.hpp b/src/armnnDeserializer/Deserializer.hpp
index 277c09ae48..bd01a35431 100644
--- a/src/armnnDeserializer/Deserializer.hpp
+++ b/src/armnnDeserializer/Deserializer.hpp
@@ -88,6 +88,7 @@ private:
     void ParseActivation(GraphPtr graph, unsigned int layerIndex);
     void ParseAdd(GraphPtr graph, unsigned int layerIndex);
     void ParseArgMinMax(GraphPtr graph, unsigned int layerIndex);
+    void ParseBatchMatMul(GraphPtr graph, unsigned int layerIndex);
     void ParseBatchToSpaceNd(GraphPtr graph, unsigned int layerIndex);
     void ParseBatchNormalization(GraphPtr graph, unsigned int layerIndex);
     void ParseCast(GraphPtr graph, unsigned int layerIndex);
diff --git a/src/armnnDeserializer/test/DeserializeBatchMatMul.cpp b/src/armnnDeserializer/test/DeserializeBatchMatMul.cpp
new file mode 100644
index 0000000000..40f93ce420
--- /dev/null
+++ b/src/armnnDeserializer/test/DeserializeBatchMatMul.cpp
@@ -0,0 +1,213 @@
+//
+// Copyright © 2022 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "ParserFlatbuffersSerializeFixture.hpp"
+#include <armnnDeserializer/IDeserializer.hpp>
+
+#include <doctest/doctest.h>
+
+#include <string>
+
+TEST_SUITE("Deserializer_BatchMatMul")
+{
+struct BatchMatMulFixture : public ParserFlatbuffersSerializeFixture
+{
+    explicit BatchMatMulFixture(const std::string& inputXShape,
+                                const std::string& inputYShape,
+                                const std::string& outputShape,
+                                const std::string& dataType)
+    {
+        m_JsonString = R"(
+            {
+                inputIds:[
+                    0,
+                    1
+                ],
+                outputIds:[
+                    3
+                ],
+                layers:[
+                    {
+                        layer_type:"InputLayer",
+                        layer:{
+                            base:{
+                                layerBindingId:0,
+                                base:{
+                                    index:0,
+                                    layerName:"InputXLayer",
+                                    layerType:"Input",
+                                    inputSlots:[
+                                        {
+                                            index:0,
+                                            connection:{
+                                                sourceLayerIndex:0,
+                                                outputSlotIndex:0
+                                            },
+
+                                        }
+                                    ],
+                                    outputSlots:[
+                                        {
+                                            index:0,
+                                            tensorInfo:{
+                                                dimensions:)" + inputXShape + R"(,
+                                                dataType:)" + dataType + R"(
+                                            },
+
+                                        }
+                                    ],
+
+                                },
+
+                            }
+                        },
+
+                    },
+                    {
+                        layer_type:"InputLayer",
+                        layer:{
+                            base:{
+                                layerBindingId:1,
+                                base:{
+                                    index:1,
+                                    layerName:"InputYLayer",
+                                    layerType:"Input",
+                                    inputSlots:[
+                                        {
+                                            index:0,
+                                            connection:{
+                                                sourceLayerIndex:0,
+                                                outputSlotIndex:0
+                                            },
+
+                                        }
+                                    ],
+                                    outputSlots:[
+                                        {
+                                            index:0,
+                                            tensorInfo:{
+                                                dimensions:)" + inputYShape + R"(,
+                                                dataType:)" + dataType + R"(
+                                            },
+
+                                        }
+                                    ],
+
+                                },
+
+                            }
+                        },
+
+                    },
+                    {
+                        layer_type:"BatchMatMulLayer",
+                        layer:{
+                            base:{
+                                index:2,
+                                layerName:"BatchMatMulLayer",
+                                layerType:"BatchMatMul",
+                                inputSlots:[
+                                    {
+                                        index:0,
+                                        connection:{
+                                            sourceLayerIndex:0,
+                                            outputSlotIndex:0
+                                        },
+
+                                    },
+                                    {
+                                        index:1,
+                                        connection:{
+                                            sourceLayerIndex:1,
+                                            outputSlotIndex:0
+                                        },
+
+                                    }
+                                ],
+                                outputSlots:[
+                                    {
+                                        index:0,
+                                        tensorInfo:{
+                                            dimensions:)" + outputShape + R"(,
+                                            dataType:)" + dataType + R"(
+                                        },
+
+                                    }
+                                ],
+
+                            },
+                            descriptor:{
+                                transposeX:false,
+                                transposeY:false,
+                                adjointX:false,
+                                adjointY:false,
+                                dataLayoutX:NHWC,
+                                dataLayoutY:NHWC
+                            }
+                        },
+
+                    },
+                    {
+                        layer_type:"OutputLayer",
+                        layer:{
+                            base:{
+                                layerBindingId:0,
+                                base:{
+                                    index:3,
+                                    layerName:"OutputLayer",
+                                    layerType:"Output",
+                                    inputSlots:[
+                                        {
+                                            index:0,
+                                            connection:{
+                                                sourceLayerIndex:2,
+                                                outputSlotIndex:0
+                                            },
+
+                                        }
+                                    ],
+                                    outputSlots:[
+                                        {
+                                            index:0,
+                                            tensorInfo:{
+                                                dimensions:)" + outputShape + R"(,
+                                                dataType:)" + dataType + R"(
+                                            },
+
+                                        }
+                                    ],
+
+                                }
+                            }
+                        },
+
+                    }
+                ]
+            }
+        )";
+        Setup();
+    }
+};
+
+struct SimpleBatchMatMulFixture : BatchMatMulFixture
+{
+    SimpleBatchMatMulFixture()
+        : BatchMatMulFixture("[ 1, 2, 2, 1 ]",
+                             "[ 1, 2, 2, 1 ]",
+                             "[ 1, 2, 2, 1 ]",
+                             "Float32")
+    {}
+};
+
+TEST_CASE_FIXTURE(SimpleBatchMatMulFixture, "SimpleBatchMatMulTest")
+{
+    RunTest<4, armnn::DataType::Float32>(
+        0,
+        {{"InputXLayer", { 1.0f, 2.0f, 3.0f, 4.0f }},
+         {"InputYLayer", { 5.0f, 6.0f, 7.0f, 8.0f }}},
+        {{"OutputLayer", { 19.0f, 22.0f, 43.0f, 50.0f }}});
+}
+
+}
\ No newline at end of file