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

9 files changed, 348 insertions, 0 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c63d8fc0df..4e4818d232 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -775,6 +775,7 @@ if(BUILD_UNIT_TESTS)
             src/armnnDeserializer/test/DeserializeActivation.cpp
             src/armnnDeserializer/test/DeserializeAdd.cpp
             src/armnnDeserializer/test/DeserializeArgMinMax.cpp
+            src/armnnDeserializer/test/DeserializeBatchMatMul.cpp
             src/armnnDeserializer/test/DeserializeBatchToSpaceNd.cpp
             src/armnnDeserializer/test/DeserializeBatchNormalization.cpp
             src/armnnDeserializer/test/DeserializeCast.cpp
diff --git a/docs/05_02_deserializer_serializer.dox b/docs/05_02_deserializer_serializer.dox
index 6cfaf29968..c36e010323 100644
--- a/docs/05_02_deserializer_serializer.dox
+++ b/docs/05_02_deserializer_serializer.dox
@@ -22,6 +22,7 @@ The Arm NN SDK Serializer currently supports the following layers:
 - Activation
 - Addition
 - ArgMinMax
+- BatchMatMul
 - BatchToSpaceNd
 - BatchNormalization
 - Cast
@@ -114,6 +115,7 @@ The Arm NN SDK Deserialize parser currently supports the following layers:
 - Activation
 - Addition
 - ArgMinMax
+- BatchMatMul
 - BatchToSpaceNd
 - BatchNormalization
 - Cast
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
diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs
index f301fce818..2dbfd85b23 100644
--- a/src/armnnSerializer/ArmnnSchema.fbs
+++ b/src/armnnSerializer/ArmnnSchema.fbs
@@ -182,6 +182,7 @@ enum LayerType : uint {
     Convolution3d = 65,
     Pooling3d = 66,
     GatherNd = 67,
+    BatchMatMul = 68,
 }
 
 // Base layer table to be used as part of other layers
@@ -1009,6 +1010,20 @@ table UnidirectionalSequenceLstmLayer {
     inputParams:LstmInputParams;
 }
 
+table BatchMatMulDescriptor {
+    transposeX:bool = false;
+    transposeY:bool = false;
+    adjointX:bool = false;
+    adjointY:bool = false;
+    dataLayoutX:DataLayout = NCHW;
+    dataLayoutY:DataLayout = NCHW;
+}
+
+table BatchMatMulLayer {
+    base:LayerBase;
+    descriptor:BatchMatMulDescriptor;
+}
+
 union Layer {
     ActivationLayer,
     AdditionLayer,
@@ -1078,6 +1093,7 @@ union Layer {
     Convolution3dLayer,
     Pooling3dLayer,
     GatherNdLayer,
+    BatchMatMulLayer,
 }
 
 table AnyLayer {
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index 488dac6186..c9a3022b8d 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -218,6 +218,33 @@ void SerializerStrategy::SerializeArgMinMaxLayer(const armnn::IConnectableLayer
     CreateAnyLayer(flatBufferLayer.o, serializer::Layer::Layer_ArgMinMaxLayer);
 }
 
+void SerializerStrategy::SerializeBatchMatMulLayer(const armnn::IConnectableLayer* layer,
+                                                   const armnn::BatchMatMulDescriptor& descriptor,
+                                                   const char* name)
+{
+    IgnoreUnused(name);
+
+    // Create FlatBuffer BaseLayer
+    auto flatBufferBaseLayer = CreateLayerBase(layer, serializer::LayerType::LayerType_BatchMatMul);
+
+    // Create the FlatBuffer BatchMatMulDescriptor
+    auto flatBufferDescriptor = CreateBatchMatMulDescriptor(m_flatBufferBuilder,
+                                                            descriptor.m_TransposeX,
+                                                            descriptor.m_TransposeY,
+                                                            descriptor.m_AdjointX,
+                                                            descriptor.m_AdjointY,
+                                                            GetFlatBufferDataLayout(descriptor.m_DataLayoutX),
+                                                            GetFlatBufferDataLayout(descriptor.m_DataLayoutY));
+
+    // Create the FlatBuffer BatchMatMulLayer
+    auto flatBufferBatchMatMulLayer = CreateBatchMatMulLayer(m_flatBufferBuilder,
+                                                             flatBufferBaseLayer,
+                                                             flatBufferDescriptor);
+
+    // Add the AnyLayer to the FlatBufferLayers
+    CreateAnyLayer(flatBufferBatchMatMulLayer.o, serializer::Layer::Layer_BatchMatMulLayer);
+}
+
 // Build FlatBuffer for BatchToSpaceNd Layer
 void SerializerStrategy::SerializeBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer,
                                                       const armnn::BatchToSpaceNdDescriptor& descriptor,
@@ -1971,6 +1998,15 @@ void SerializerStrategy::ExecuteStrategy(const armnn::IConnectableLayer* layer,
             SerializeArgMinMaxLayer(layer, layerDescriptor, name);
             break;
         }
+        case armnn::LayerType::BatchMatMul:
+        {
+            const armnn::BatchMatMulDescriptor& layerDescriptor =
+                    static_cast<const armnn::BatchMatMulDescriptor&>(descriptor);
+            SerializeBatchMatMulLayer(layer,
+                                      layerDescriptor,
+                                      name);
+            break;
+        }
         case armnn::LayerType::BatchNormalization :
         {
             const armnn::BatchNormalizationDescriptor& layerDescriptor =
diff --git a/src/armnnSerializer/Serializer.hpp b/src/armnnSerializer/Serializer.hpp
index 216f4dc016..60fed4f6df 100644
--- a/src/armnnSerializer/Serializer.hpp
+++ b/src/armnnSerializer/Serializer.hpp
@@ -113,6 +113,10 @@ private:
                                  const armnn::ArgMinMaxDescriptor& argMinMaxDescriptor,
                                  const char* name = nullptr);
 
+    void SerializeBatchMatMulLayer(const armnn::IConnectableLayer* layer,
+                                   const armnn::BatchMatMulDescriptor& descriptor,
+                                   const char* name = nullptr);
+
     void SerializeBatchToSpaceNdLayer(const armnn::IConnectableLayer* layer,
                                       const armnn::BatchToSpaceNdDescriptor& descriptor,
                                       const char* name = nullptr);
diff --git a/src/armnnSerializer/test/SerializerTests.cpp b/src/armnnSerializer/test/SerializerTests.cpp
index 3c00fc43ae..a568bf15c9 100644
--- a/src/armnnSerializer/test/SerializerTests.cpp
+++ b/src/armnnSerializer/test/SerializerTests.cpp
@@ -92,6 +92,47 @@ TEST_CASE("SerializeArgMinMaxSigned64")
     SerializeArgMinMaxTest(armnn::DataType::Signed64);
 }
 
+TEST_CASE("SerializeBatchMatMul")
+{
+    const std::string layerName("batchMatMul");
+    const armnn::TensorInfo inputXInfo({2, 3, 4, 5}, armnn::DataType::Float32);
+    const armnn::TensorInfo inputYInfo({2, 4, 3, 5}, armnn::DataType::Float32);
+
+    const armnn::TensorInfo outputInfo({2, 3, 3, 5}, armnn::DataType::Float32);
+
+    armnn::BatchMatMulDescriptor descriptor(false,
+                                            false,
+                                            false,
+                                            false,
+                                            armnn::DataLayout::NHWC,
+                                            armnn::DataLayout::NHWC);
+
+    armnn::INetworkPtr network = armnn::INetwork::Create();
+    armnn::IConnectableLayer* const inputXLayer = network->AddInputLayer(0);
+    armnn::IConnectableLayer* const inputYLayer = network->AddInputLayer(1);
+
+    armnn::IConnectableLayer* const batchMatMulLayer =
+        network->AddBatchMatMulLayer(descriptor, layerName.c_str());
+    armnn::IConnectableLayer* const outputLayer = network->AddOutputLayer(0);
+
+    inputXLayer->GetOutputSlot(0).Connect(batchMatMulLayer->GetInputSlot(0));
+    inputYLayer->GetOutputSlot(0).Connect(batchMatMulLayer->GetInputSlot(1));
+    batchMatMulLayer->GetOutputSlot(0).Connect(outputLayer->GetInputSlot(0));
+
+    inputXLayer->GetOutputSlot(0).SetTensorInfo(inputXInfo);
+    inputYLayer->GetOutputSlot(0).SetTensorInfo(inputYInfo);
+    batchMatMulLayer->GetOutputSlot(0).SetTensorInfo(outputInfo);
+
+    armnn::INetworkPtr deserializedNetwork = DeserializeNetwork(SerializeNetwork(*network));
+    CHECK(deserializedNetwork);
+
+    LayerVerifierBaseWithDescriptor<armnn::BatchMatMulDescriptor> verifier(layerName,
+                                                                           {inputXInfo, inputYInfo},
+                                                                           {outputInfo},
+                                                                           descriptor);
+    deserializedNetwork->ExecuteStrategy(verifier);
+}
+
 TEST_CASE("SerializeBatchNormalization")
 {
     const std::string layerName("batchNormalization");