Add Split support to TOSA Reference Backend

* Resolves IVGCVSW-7918

Signed-off-by: Kevin May <kevin.may@arm.com>
Change-Id: Ic2afaa55f7ee88ce4c9b8ea696eef5f28663f8c6

1 files changed, 160 insertions, 0 deletions

diff --git a/src/backends/backendsCommon/test/SplitterEndToEndTestImpl.hpp b/src/backends/backendsCommon/test/SplitterEndToEndTestImpl.hpp
index da4c6a62be..5a4990f748 100644
--- a/src/backends/backendsCommon/test/SplitterEndToEndTestImpl.hpp
+++ b/src/backends/backendsCommon/test/SplitterEndToEndTestImpl.hpp
@@ -48,6 +48,7 @@ INetworkPtr CreateSplitterNetwork(const TensorShape& inputShape,
     splitterDimSizes[splitAxis] /= numSplit;
 
     SplitterDescriptor splitDesc(numSplit, inputShape.GetNumDimensions());
+    splitDesc.SetAxis(static_cast<int32_t>(splitAxis));
     for (unsigned int g = 0; g < numSplit; ++g)
     {
         // Set the size of the views.
@@ -101,6 +102,35 @@ void Splitter1dEndToEnd(const std::vector<BackendId>& backends)
 }
 
 template<armnn::DataType ArmnnType>
+void Splitter1dEndToEndFloat16(const std::vector<BackendId>& backends)
+{
+    using namespace armnn;
+    using namespace half_float::literal;
+    using Half = half_float::half;
+
+    unsigned int splitAxis = 0;
+    unsigned int numSplit = 2;
+    const TensorShape& inputShape = { 4 };
+    const std::vector<TensorShape> outputShapes{{ 2 }, { 2 }};
+
+    // Builds up the structure of the network
+    INetworkPtr net = CreateSplitterNetwork<ArmnnType>(inputShape, outputShapes, splitAxis, numSplit);
+
+    CHECK(net);
+
+    // Creates structures for input & output.
+    std::vector<Half> inputData{ 1._h, 2._h, 3._h, 4._h };
+
+    std::vector<Half> expectedOutput0{ 1._h, 2._h };
+    std::vector<Half> expectedOutput1{ 3._h, 4._h };
+
+    std::map<int, std::vector<Half>> inputTensorData = { { 0, inputData } };
+    std::map<int, std::vector<Half>> expectedOutputData = { { 0, expectedOutput0 }, {1, expectedOutput1} };
+
+    EndToEndLayerTestImpl<ArmnnType, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
+template<armnn::DataType ArmnnType>
 void Splitter2dDim0EndToEnd(const std::vector<BackendId>& backends)
 {
     using namespace armnn;
@@ -270,6 +300,55 @@ void Splitter3dDim1EndToEnd(const std::vector<BackendId>& backends)
 }
 
 template<armnn::DataType ArmnnType>
+void Splitter3dDim1EndToEndFloat16(const std::vector<BackendId>& backends)
+{
+    using namespace armnn;
+    using namespace half_float::literal;
+    using Half = half_float::half;
+
+    unsigned int splitAxis = 1;
+    unsigned int numSplit = 2;
+    const TensorShape& inputShape = { 2, 4, 3 };
+    const std::vector<TensorShape> outputShapes{{ 2, 2, 3 }, { 2, 2, 3 }};
+
+    // Builds up the structure of the network
+    INetworkPtr net = CreateSplitterNetwork<ArmnnType>(inputShape, outputShapes, splitAxis, numSplit);
+
+    CHECK(net);
+
+    // Creates structures for input & output.
+    std::vector<Half> inputData{
+            1._h, 2._h, 3._h,
+            4._h, 5._h, 6._h,
+            7._h, 8._h, 9._h,
+            10._h, 11._h, 12._h,
+            13._h, 14._h, 15._h,
+            16._h, 17._h, 18._h,
+            19._h, 20._h, 21._h,
+            22._h, 23._h, 24._h
+    };
+
+    std::vector<Half> expectedOutput0{
+            1._h, 2._h, 3._h,
+            4._h, 5._h, 6._h,
+            13._h, 14._h, 15._h,
+            16._h, 17._h, 18._h
+    };
+    std::vector<Half> expectedOutput1{
+            7._h, 8._h, 9._h,
+            10._h, 11._h, 12._h,
+            19._h, 20._h, 21._h,
+            22._h, 23._h, 24._h
+    };
+
+    std::map<int, std::vector<Half>> inputTensorData = { { 0, inputData } };
+    std::map<int, std::vector<Half>> expectedOutputData = { { 0, expectedOutput0 },
+                                                         { 1, expectedOutput1 } };
+
+    EndToEndLayerTestImpl<ArmnnType, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
+template<armnn::DataType ArmnnType>
 void Splitter3dDim2EndToEnd(const std::vector<BackendId>& backends)
 {
     using namespace armnn;
@@ -549,6 +628,87 @@ void Splitter4dDim2EndToEnd(const std::vector<BackendId>& backends)
     EndToEndLayerTestImpl<ArmnnType, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
 }
 
+template<armnn::DataType ArmnnType>
+void Splitter4dDim2EndToEndFloat16(const std::vector<BackendId>& backends)
+{
+    using namespace armnn;
+    using namespace half_float::literal;
+    using Half = half_float::half;
+
+    unsigned int splitAxis = 2;
+    unsigned int numSplit = 2;
+    const TensorShape& inputShape = { 2, 3, 4, 2 };
+    const std::vector<TensorShape> outputShapes{{ 2, 3, 2, 2 }, { 2, 3, 2, 2 }};
+
+    // Builds up the structure of the network
+    INetworkPtr net = CreateSplitterNetwork<ArmnnType>(inputShape, outputShapes, splitAxis, numSplit);
+
+    CHECK(net);
+
+    // Creates structures for input & output.
+    std::vector<Half> inputData{
+            1._h, 2._h,
+            3._h, 4._h,
+            5._h, 6._h,
+            7._h, 8._h,
+            9._h, 10._h,
+            11._h, 12._h,
+            13._h, 14._h,
+            15._h, 16._h,
+            17._h, 18._h,
+            19._h, 20._h,
+            21._h, 22._h,
+            23._h, 24._h,
+            25._h, 26._h,
+            27._h, 28._h,
+            29._h, 30._h,
+            31._h, 32._h,
+            33._h, 34._h,
+            35._h, 36._h,
+            37._h, 38._h,
+            39._h, 40._h,
+            41._h, 42._h,
+            43._h, 44._h,
+            45._h, 46._h,
+            47._h, 48._h
+    };
+
+    std::vector<Half> expectedOutput0{
+            1._h, 2._h,
+            3._h, 4._h,
+            9._h, 10._h,
+            11._h, 12._h,
+            17._h, 18._h,
+            19._h, 20._h,
+            25._h, 26._h,
+            27._h, 28._h,
+            33._h, 34._h,
+            35._h, 36._h,
+            41._h, 42._h,
+            43._h, 44._h
+    };
+
+    std::vector<Half> expectedOutput1{
+            5._h, 6._h,
+            7._h, 8._h,
+            13._h, 14._h,
+            15._h, 16._h,
+            21._h, 22._h,
+            23._h, 24._h,
+            29._h, 30._h,
+            31._h, 32._h,
+            37._h, 38._h,
+            39._h, 40._h,
+            45._h, 46._h,
+            47._h, 48._h
+    };
+
+    std::map<int, std::vector<Half>> inputTensorData = {{ 0,inputData }};
+    std::map<int, std::vector<Half>> expectedOutputData = {{ 0, expectedOutput0 }, { 1, expectedOutput1 }};
+
+    EndToEndLayerTestImpl<ArmnnType, ArmnnType>(std::move(net), inputTensorData, expectedOutputData, backends);
+}
+
 template<armnn::DataType ArmnnType, typename T = armnn::ResolveType<ArmnnType>>
 void Splitter4dDim3EndToEnd(const std::vector<BackendId>& backends)
 {