IVGCVSW-7343 Add Transpose support to TOSA Reference Backend

Signed-off-by: Cathal Corbett <cathal.corbett@arm.com>
Change-Id: I11505f672349e1f04143edfdc2df8775f685372d

10 files changed, 209 insertions, 1 deletions

diff --git a/src/backends/tosaCommon/TosaMappings.cpp b/src/backends/tosaCommon/TosaMappings.cpp
index b0f8fd9c2d..0b5fa1a158 100644
--- a/src/backends/tosaCommon/TosaMappings.cpp
+++ b/src/backends/tosaCommon/TosaMappings.cpp
@@ -79,6 +79,11 @@ TosaSerializationBasicBlock* GetTosaMapping(const Layer* layer,
             auto transposeConv2dDesc = PolymorphicDowncast<const TransposeConvolution2dDescriptor*>(&descriptor);
             return ConvertTransposeConv2dToTosaOperator(layer, inputs, outputs, transposeConv2dDesc);
         }
+        case LayerType::Transpose:
+        {
+            auto transposeDesc = PolymorphicDowncast<const TransposeDescriptor*>(&descriptor);
+            return ConvertTransposeToTosaOperator(layer, inputs, outputs, transposeDesc);
+        }
         default:
         {
             return CreateEmptyTosaSerializationBasicBlock();
diff --git a/src/backends/tosaCommon/operatorMappings/CMakeLists.txt b/src/backends/tosaCommon/operatorMappings/CMakeLists.txt
index 6e897aaa8c..2ec052cd43 100644
--- a/src/backends/tosaCommon/operatorMappings/CMakeLists.txt
+++ b/src/backends/tosaCommon/operatorMappings/CMakeLists.txt
@@ -23,6 +23,8 @@ list(APPEND armnnTosaBackendOperators_sources
         TosaOperatorUtils.hpp
         TransposeConv2dOperator.hpp
         TransposeConv2dOperator.cpp
+        TransposeOperator.hpp
+        TransposeOperator.cpp
     )
 
 add_library(armnnTosaBackendOperators OBJECT ${armnnTosaBackendOperators_sources})
diff --git a/src/backends/tosaCommon/operatorMappings/TosaCommonOperators.hpp b/src/backends/tosaCommon/operatorMappings/TosaCommonOperators.hpp
index 7b117d8443..3f27371295 100644
--- a/src/backends/tosaCommon/operatorMappings/TosaCommonOperators.hpp
+++ b/src/backends/tosaCommon/operatorMappings/TosaCommonOperators.hpp
@@ -14,3 +14,4 @@
 #include "ReshapeOperator.hpp"
 #include "SliceOperator.hpp"
 #include "TransposeConv2dOperator.hpp"
+#include "TransposeOperator.hpp"
diff --git a/src/backends/tosaCommon/operatorMappings/TransposeOperator.cpp b/src/backends/tosaCommon/operatorMappings/TransposeOperator.cpp
new file mode 100644
index 0000000000..56178e428b
--- /dev/null
+++ b/src/backends/tosaCommon/operatorMappings/TransposeOperator.cpp
@@ -0,0 +1,65 @@
+//
+// Copyright © 2022 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "TransposeOperator.hpp"
+
+TosaSerializationBasicBlock* ConvertTransposeToTosaOperator(const Layer* layer,
+                                                            const std::vector<const TensorInfo*>& inputs,
+                                                            const std::vector<const TensorInfo*>& outputs,
+                                                            const TransposeDescriptor* transposeDescriptor)
+{
+    std::string input0Name = std::string("input0_");
+    std::string outputName = std::string("output0_");
+    std::string blockName  = std::string("Op_TRANSPOSE_block_") + GetUniqueTosaMappingID();
+
+    // If a layer is present then the block will be used for execution, so input and output names need to be determined
+    // using the previous and following layers so the graph is connected correctly. For validation this doesn't matter.
+    if(layer != nullptr)
+    {
+        // Get the layers connected to the input slot and determine unique tensor name.
+        Layer& connectedLayer0 = layer->GetInputSlot(0).GetConnectedOutputSlot()->GetOwningLayer();
+        input0Name = GenerateUniqueName(connectedLayer0, 0);
+
+        // Determine unique output tensor name.
+        outputName = GenerateUniqueOutputName(*layer, 0);
+    }
+
+    std::vector<int32_t> mappings(transposeDescriptor->m_DimMappings.begin(),
+                                  transposeDescriptor->m_DimMappings.end());
+    TosaTransposeAttribute attribute(mappings);
+
+    auto* op = new TosaSerializationOperator(Op_TRANSPOSE,
+                                             Attribute_TransposeAttribute,
+                                             &attribute,
+                                             {input0Name},
+                                             {outputName});
+
+
+    std::vector<TosaSerializationTensor*> tensors;
+
+    // Only add input tensors if connected layer is an input layer.
+    // As intermediate or constant tensors will be created separately.
+    // There also can't be duplicate tensor.
+    if(input0Name.find("input0_") != std::string::npos)
+    {
+        std::vector<int32_t> inputShape0 = GetTosaTensorShape(inputs[0]->GetShape());
+        DType inputDType0 = ArmNNToDType(inputs[0]->GetDataType());
+
+        tensors.push_back(new TosaSerializationTensor(input0Name, inputShape0, inputDType0, {}));
+    }
+
+    std::vector<int32_t> outputShape0 = GetTosaTensorShape(outputs[0]->GetShape());
+    DType outputDType0 = ArmNNToDType(outputs[0]->GetDataType());
+
+    tensors.push_back(new TosaSerializationTensor(outputName, outputShape0, outputDType0, {}));
+
+    // operatorInputNames/operatorOutputNames ends up being the same as
+    // blockInputNames/blockOutputNames for one-to-one ArmNN to TOSA mappings
+    return new TosaSerializationBasicBlock(blockName, // name
+                                           {op}, // operators
+                                           tensors, // tensors
+                                           {input0Name}, // inputs
+                                           {outputName}); // outputs
+}
\ No newline at end of file
diff --git a/src/backends/tosaCommon/operatorMappings/TransposeOperator.hpp b/src/backends/tosaCommon/operatorMappings/TransposeOperator.hpp
new file mode 100644
index 0000000000..3d1e2acd14
--- /dev/null
+++ b/src/backends/tosaCommon/operatorMappings/TransposeOperator.hpp
@@ -0,0 +1,20 @@
+//
+// Copyright © 2022 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include "TosaOperatorUtils.hpp"
+
+#include <Layer.hpp>
+
+#include <tosa_serialization_handler.h>
+
+using namespace armnn;
+using namespace tosa;
+
+TosaSerializationBasicBlock* ConvertTransposeToTosaOperator(const Layer* layer,
+                                                            const std::vector<const TensorInfo*>& inputs,
+                                                            const std::vector<const TensorInfo*>& outputs,
+                                                            const TransposeDescriptor* transposeDescriptor);
diff --git a/src/backends/tosaCommon/test/OneToOneMappingTests.cpp b/src/backends/tosaCommon/test/OneToOneMappingTests.cpp
index 146a9cba37..4cc37918e5 100644
--- a/src/backends/tosaCommon/test/OneToOneMappingTests.cpp
+++ b/src/backends/tosaCommon/test/OneToOneMappingTests.cpp
@@ -664,6 +664,64 @@ TEST_CASE("GetTosaMappingFromLayer_TransposeConv2dLayer")
                         LayerType::TransposeConvolution2d);
 }
 
+TEST_CASE("GetTosaMapping_TransposeLayer")
+{
+    TensorInfo inputInfo  = TensorInfo({ 1, 1, 5, 3 }, DataType::Float32, 0.0f, 0, true);
+    TensorInfo outputInfo = TensorInfo({ 1, 5, 1, 3 }, DataType::Float32, 0.0f, 0, true);
+
+    std::vector<std::vector<int32_t>> inputShape  = {{ 1, 1, 5, 3 }};
+    std::vector<std::vector<int32_t>> outputShape = {{ 1, 5, 1, 3 }};
+
+    TransposeDescriptor transposeDescriptor = TransposeDescriptor({ 0, 2, 1 ,3 });
+
+    TosaSerializationBasicBlock* basicBlock =
+        GetTosaMapping(nullptr, LayerType::Transpose, {&inputInfo,}, {&outputInfo}, transposeDescriptor);
+    AssertTosaOneToOneMappingBasicBlock(basicBlock,
+                                        inputShape,
+                                        outputShape,
+                                        Op_TRANSPOSE,
+                                        Attribute_TransposeAttribute,
+                                        transposeDescriptor,
+                                        LayerType::Transpose);
+}
+
+TEST_CASE("GetTosaMappingFromLayer_TransposeLayer")
+{
+    IRuntime::CreationOptions options;
+    IRuntimePtr runtime(IRuntime::Create(options));
+
+    // Builds up the structure of the network.
+    INetworkPtr net(INetwork::Create());
+
+    TransposeDescriptor transposeDescriptor = TransposeDescriptor({ 0, 2, 1 ,3 });
+
+    IConnectableLayer* input     = net->AddInputLayer(0, "input0");
+    IConnectableLayer* transpose = net->AddTransposeLayer(transposeDescriptor, "transpose");
+    IConnectableLayer* output    = net->AddOutputLayer(0, "output");
+
+    input->GetOutputSlot(0).Connect(transpose->GetInputSlot(0));
+    transpose->GetOutputSlot(0).Connect(output->GetInputSlot(0));
+
+    TensorInfo inputInfo  = TensorInfo({ 1, 1, 5, 3 }, DataType::Float32, 0.0f, 0, true);
+    TensorInfo outputInfo = TensorInfo({ 1, 5, 1, 3 }, DataType::Float32, 0.0f, 0, true);
+
+    input->GetOutputSlot(0).SetTensorInfo(inputInfo);
+    transpose->GetOutputSlot(0).SetTensorInfo(outputInfo);
+
+    std::vector<std::vector<int32_t>> inputShape  = {{ 1, 1, 5, 3 }};
+    std::vector<std::vector<int32_t>> outputShape = {{ 1, 5, 1, 3 }};
+
+    TosaSerializationBasicBlock* basicBlock =
+        GetTosaMappingFromLayer(PolymorphicDowncast<Layer*>(transpose));
+    AssertTosaOneToOneMappingBasicBlock(basicBlock,
+                                        inputShape,
+                                        outputShape,
+                                        Op_TRANSPOSE,
+                                        Attribute_TransposeAttribute,
+                                        transposeDescriptor,
+                                        LayerType::Transpose);
+}
+
 TEST_CASE("GetTosaMapping_Unimplemented")
 {
     TosaSerializationBasicBlock* basicBlock =
diff --git a/src/backends/tosaCommon/test/TosaTestUtils.hpp b/src/backends/tosaCommon/test/TosaTestUtils.hpp
index 140cb83983..e24055371f 100644
--- a/src/backends/tosaCommon/test/TosaTestUtils.hpp
+++ b/src/backends/tosaCommon/test/TosaTestUtils.hpp
@@ -158,6 +158,14 @@ inline void VerifyTosaAttribute(const BaseDescriptor& descriptor,
             CHECK(stride == transposeConvAttribute.stride());
             break;
         }
+        case LayerType::Transpose:
+        {
+            auto transposeDesc = PolymorphicDowncast<const TransposeDescriptor*>(&descriptor);
+            std::vector<int> outPerm(transposeDesc->m_DimMappings.begin(), transposeDesc->m_DimMappings.end());
+            TosaTransposeAttribute transposeAttribute(attribute);
+            CHECK(outPerm == transposeAttribute.perms());
+            break;
+        }
         default:
             break;
     }
diff --git a/src/backends/tosaReference/TosaRefLayerSupport.cpp b/src/backends/tosaReference/TosaRefLayerSupport.cpp
index b37ecc4ab7..6113b5861a 100644
--- a/src/backends/tosaReference/TosaRefLayerSupport.cpp
+++ b/src/backends/tosaReference/TosaRefLayerSupport.cpp
@@ -71,7 +71,7 @@ bool TosaRefLayerSupport::IsLayerSupported(const LayerType& type,
         case LayerType::Pooling2d:
         case LayerType::Reshape:
         case LayerType::Slice:
-            // Setup inputs and outputs
+        case LayerType::Transpose:
             inputInfos.push_back(&infos[0]);
             outputInfos.push_back(&infos[1]);
             break;
diff --git a/src/backends/tosaReference/test/TosaRefEndToEndTests.cpp b/src/backends/tosaReference/test/TosaRefEndToEndTests.cpp
index 67b87ae8b9..e19462e986 100644
--- a/src/backends/tosaReference/test/TosaRefEndToEndTests.cpp
+++ b/src/backends/tosaReference/test/TosaRefEndToEndTests.cpp
@@ -14,6 +14,7 @@
 #include "backendsCommon/test/SliceEndToEndTestImpl.hpp"
 #include "backendsCommon/test/SubtractionEndToEndTestImpl.hpp"
 #include "backendsCommon/test/TransposeConvolution2dEndToEndTestImpl.hpp"
+#include "backendsCommon/test/TransposeEndToEndTestImpl.hpp"
 
 #include <doctest/doctest.h>
 
@@ -195,4 +196,10 @@ TEST_CASE("TosaRefSimpleTransposeConvolution2dEndToEndFloatNhwcTest")
         tosaDefaultBackends, armnn::DataLayout::NHWC);
 }
 
+// Transpose
+TEST_CASE("TosaRefTransposeEndtoEndTestFloat32")
+{
+    TransposeEndToEnd<armnn::DataType::Float32>(tosaDefaultBackends);
+}
+
 }
\ No newline at end of file
diff --git a/src/backends/tosaReference/test/TosaRefLayerSupportTests.cpp b/src/backends/tosaReference/test/TosaRefLayerSupportTests.cpp
index 9119b13557..66dfbe8dff 100644
--- a/src/backends/tosaReference/test/TosaRefLayerSupportTests.cpp
+++ b/src/backends/tosaReference/test/TosaRefLayerSupportTests.cpp
@@ -509,4 +509,46 @@ TEST_CASE("IsLayerSupportedTosaReferenceTransposeConv2dUnsupported")
     CHECK(!supported);
 }
 
+TEST_CASE("IsLayerSupportedTosaReferenceTranspose")
+{
+    TensorShape inShape = { 1, 1, 5, 3 };
+    TensorShape outShape = { 1, 5, 1, 3 };
+    TensorInfo in(inShape, DataType::Float32);
+    TensorInfo out(outShape, DataType::Float32);
+
+    TransposeDescriptor transposeDescriptor = TransposeDescriptor({ 0, 2, 1 ,3 });
+
+    TosaRefLayerSupport supportChecker;
+    std::string reasonIfNotSupported;
+    auto supported = supportChecker.IsLayerSupported(LayerType::Transpose,
+                                                     {in, out},
+                                                     transposeDescriptor,
+                                                     EmptyOptional(),
+                                                     EmptyOptional(),
+                                                     reasonIfNotSupported);
+
+    CHECK(supported);
+}
+
+TEST_CASE("IsLayerSupportedTosaReferenceTransposeUnsupported")
+{
+    TensorShape inShape = { 1, 1, 5, 3 };
+    TensorShape outShape = { 1, 5, 1, 3 };
+    TensorInfo in(inShape, DataType::Signed64);
+    TensorInfo out(outShape, DataType::Signed64);
+
+    TransposeDescriptor transposeDescriptor = TransposeDescriptor({ 0, 2, 1 ,3 });
+
+    TosaRefLayerSupport supportChecker;
+    std::string reasonIfNotSupported;
+    auto supported = supportChecker.IsLayerSupported(LayerType::Transpose,
+                                                     {in, out},
+                                                     transposeDescriptor,
+                                                     EmptyOptional(),
+                                                     EmptyOptional(),
+                                                     reasonIfNotSupported);
+
+    CHECK(!supported);
+}
+
 }