IVGCVSW-2264 Remove input swizzling from ParseConv2D in the TF parser

 * Removed the input swizzling when the data layout is NHWC
 * Permuting weights depending on the data layout used
 * Added getter methods to ParsedConstTfOperation to get the tensor
   info and the storage memory area, needed for swizzling the weights
 * Added unit tests for both NHWC and NCHW data layouts

Change-Id: I6543900c594417df630b2663d8551158b93b7836

1 files changed, 103 insertions, 40 deletions

diff --git a/src/armnnTfParser/test/Convolution2d.cpp b/src/armnnTfParser/test/Convolution2d.cpp
index cc534df53d..aead1fe965 100644
--- a/src/armnnTfParser/test/Convolution2d.cpp
+++ b/src/armnnTfParser/test/Convolution2d.cpp
@@ -6,6 +6,8 @@
 #include <boost/test/unit_test.hpp>
 #include "armnnTfParser/ITfParser.hpp"
 #include "ParserPrototxtFixture.hpp"
+
+#include <array>
 #include <string>
 #include <iostream>
 
@@ -13,15 +15,28 @@ BOOST_AUTO_TEST_SUITE(TensorflowParser)
 
 struct Convolution2dFixture : public armnnUtils::ParserPrototxtFixture<armnnTfParser::ITfParser>
 {
-    explicit Convolution2dFixture(const char* paddingType)
-    : Convolution2dFixture(paddingType, 1)
+    explicit Convolution2dFixture(const std::string& dataLayout, const std::string& paddingType)
+    : Convolution2dFixture(dataLayout, paddingType, 1)
     {}
 
     // Dilation: 0 - dilations attribute is not included;
     // Dilation: >0 - dilations attribute set to [1,v,v,1], where v is the value of the dilation arg
-    explicit Convolution2dFixture(const char* paddingType, int stride, int dilation = 0)
+    explicit Convolution2dFixture(const std::string& dataLayout, const std::string& paddingType,
+                                  int stride, int dilation = 0)
     {
-        std::string strideString = std::to_string(stride);
+        std::string strideString ("        i: 1 \n"
+                                  "        i: 1 \n");
+        if (dataLayout == "NHWC")
+        {
+            strideString.append("        i: " + std::to_string(stride) + " \n"
+                                "        i: 1 \n");
+        }
+        else // dataLayout == "NCHW"
+        {
+            strideString.append("        i: 1 \n"
+                                "        i: " + std::to_string(stride) + " \n");
+        }
+
         std::string dilationString = std::to_string(dilation);
         m_Prototext = "node { \n"
             "    name: \"graphInput\" \n"
@@ -87,13 +102,15 @@ struct Convolution2dFixture : public armnnUtils::ParserPrototxtFixture<armnnTfPa
             "  attr { \n"
             "    key: \"data_format\" \n"
             "    value { \n"
-            "      s: \"NHWC\" \n"
-            "    } \n"
-            "  } \n"
-            "  attr { \n"
-            "    key: \"padding\" \n"
-            "    value { \n"
             "      s: \"";
+        m_Prototext.append(dataLayout);
+        m_Prototext.append("\"\n"
+                           "    } \n"
+                           "  } \n"
+                           "  attr { \n"
+                           "    key: \"padding\" \n"
+                           "    value { \n"
+                           "      s: \"");
         m_Prototext.append(paddingType);
         m_Prototext.append("\"\n"
                            "    } \n"
@@ -101,14 +118,10 @@ struct Convolution2dFixture : public armnnUtils::ParserPrototxtFixture<armnnTfPa
                            "  attr { \n"
                            "    key: \"strides\" \n"
                            "    value { \n"
-                           "      list { \n"
-                           "        i: 1 \n"
-                           "        i: 1 \n"
-                           "        i: ");
+                           "      list { \n");
         m_Prototext.append(strideString);
-        m_Prototext.append(" \n"
-                           "        i: 1 \n"
-                           "      } \n"
+
+        m_Prototext.append("      } \n"
                            "    } \n"
                            "  } \n");
 
@@ -139,67 +152,118 @@ struct Convolution2dFixture : public armnnUtils::ParserPrototxtFixture<armnnTfPa
                            "} \n");
 
         // Manual height computation based on stride parameter.
-        BOOST_ASSERT_MSG(stride == 1 || stride==2, "Add support for strides other than 1 or 2.");
-        unsigned int dims[] = {1,2,3,1};
-        if (stride == 2)
+        BOOST_ASSERT_MSG(stride == 1 || stride == 2, "Add support for strides other than 1 or 2.");
+        std::array<unsigned int, 4> dims;
+        if (dataLayout == "NHWC")
         {
-            dims[1]=3;
+            dims = { 1u, (stride == 2 ? 3u : 2u), 3u, 1u };
+        }
+        else // dataLayout == "NCHW"
+        {
+            dims = { 1u, 1u, (stride == 2 ? 3u : 2u), 3u };
         }
 
-        SetupSingleInputSingleOutput(armnn::TensorShape(4, dims), "graphInput", "potato");
+        SetupSingleInputSingleOutput(armnn::TensorShape(4, dims.data()), "graphInput", "potato");
     }
 };
 
 
-struct Convolution2dSameFixture : Convolution2dFixture
+struct Convolution2dNhwcSameFixture : Convolution2dFixture
+{
+    Convolution2dNhwcSameFixture() : Convolution2dFixture("NHWC", "SAME", 1){}
+};
+BOOST_FIXTURE_TEST_CASE(ParseConv2dNhwcSame, Convolution2dNhwcSameFixture)
+{
+    RunTest<4>({1, 2, 3, 4, 5, 6}, {2, 4, 4, 6.5f, 10 , 8.5f});
+}
+
+struct Convolution2dNchwSameFixture : Convolution2dFixture
 {
-    Convolution2dSameFixture() : Convolution2dFixture("SAME", 1){}
+    Convolution2dNchwSameFixture() : Convolution2dFixture("NCHW", "SAME", 1){}
 };
-BOOST_FIXTURE_TEST_CASE(ParseConv2DSame, Convolution2dSameFixture)
+BOOST_FIXTURE_TEST_CASE(ParseConv2dNchwSame, Convolution2dNchwSameFixture)
 {
     RunTest<4>({1, 2, 3, 4, 5, 6}, {2, 4, 4, 6.5f, 10 , 8.5f});
 }
 
-struct Convolution2dValidFixture : Convolution2dFixture
+
+struct Convolution2dNhwcValidFixture : Convolution2dFixture
+{
+    Convolution2dNhwcValidFixture() : Convolution2dFixture("NHWC", "VALID", 1){}
+};
+BOOST_FIXTURE_TEST_CASE(ParseConv2dNhwcValid, Convolution2dNhwcValidFixture)
+{
+    RunTest<4>({1, 2, 3, 4, 5, 6}, {4, 10});
+}
+
+struct Convolution2dNchwValidFixture : Convolution2dFixture
 {
-    Convolution2dValidFixture() : Convolution2dFixture("VALID", 1){}
+    Convolution2dNchwValidFixture() : Convolution2dFixture("NCHW", "VALID", 1){}
 };
-BOOST_FIXTURE_TEST_CASE(ParseConv2DValid, Convolution2dValidFixture)
+BOOST_FIXTURE_TEST_CASE(ParseConv2dNchwValid, Convolution2dNchwValidFixture)
 {
     RunTest<4>({1, 2, 3, 4, 5, 6}, {4, 10});
 }
 
 
-struct Convolution2dStride2SameFixture : Convolution2dFixture
+struct Convolution2dStride2NhwcSameFixture : Convolution2dFixture
 {
-    Convolution2dStride2SameFixture() : Convolution2dFixture("SAME", 2){}
+    Convolution2dStride2NhwcSameFixture() : Convolution2dFixture("NHWC", "SAME", 2){}
 };
-BOOST_FIXTURE_TEST_CASE(ParseConv2DStride2Same, Convolution2dStride2SameFixture)
+BOOST_FIXTURE_TEST_CASE(ParseConv2dStride2NhwcSame, Convolution2dStride2NhwcSameFixture)
+{
+    RunTest<4>({1, 2, 3, 4, 5, 6, 7, 8, 9}, {2, 4, 6.5, 8.5, 11, 13});
+}
+
+struct Convolution2dStride2NchwSameFixture : Convolution2dFixture
+{
+    Convolution2dStride2NchwSameFixture() : Convolution2dFixture("NCHW", "SAME", 2){}
+};
+BOOST_FIXTURE_TEST_CASE(ParseConv2dStride2NchwSame, Convolution2dStride2NchwSameFixture)
 {
     RunTest<4>({1, 2, 3, 4, 5, 6, 7, 8, 9}, {2, 4, 6.5, 8.5, 11, 13});
 }
 
 
-struct Convolution2dStride2ValidFixture : Convolution2dFixture
+struct Convolution2dStride2NhwcValidFixture : Convolution2dFixture
 {
-    Convolution2dStride2ValidFixture() : Convolution2dFixture("VALID", 2){}
+    Convolution2dStride2NhwcValidFixture() : Convolution2dFixture("NHWC", "VALID", 2){}
 };
-BOOST_FIXTURE_TEST_CASE(ParseConv2DStride2Valid, Convolution2dStride2ValidFixture)
+BOOST_FIXTURE_TEST_CASE(ParseConv2dStride2NhwcValid, Convolution2dStride2NhwcValidFixture)
 {
     RunTest<4>({1, 2, 3, 4, 5, 6, 7, 8, 9}, {4, 10, 16});
 }
 
+struct Convolution2dStride2NchwValidFixture : Convolution2dFixture
+{
+    Convolution2dStride2NchwValidFixture() : Convolution2dFixture("NCHW", "VALID", 2){}
+};
+BOOST_FIXTURE_TEST_CASE(ParseConv2dStride2NchwValid, Convolution2dStride2NchwValidFixture)
+{
+    RunTest<4>({1, 2, 3, 4, 5, 6, 7, 8, 9}, {4, 10, 16});
+}
+
+
+struct Convolution2dDilation1NhwcFixture : Convolution2dFixture
+{
+    Convolution2dDilation1NhwcFixture() : Convolution2dFixture("NHWC", "SAME", 1, 1){}
+};
+BOOST_FIXTURE_TEST_CASE(ParseConv2dDilation1Nhwc, Convolution2dDilation1NhwcFixture)
+{
+    RunTest<4>({1, 2, 3, 4, 5, 6}, {2, 4, 4, 6.5f, 10 , 8.5f});
+}
 
-struct Convolution2dDilation1Fixture : Convolution2dFixture
+struct Convolution2dDilation1NchwFixture : Convolution2dFixture
 {
-    Convolution2dDilation1Fixture() : Convolution2dFixture("SAME", 1, 1){}
+    Convolution2dDilation1NchwFixture() : Convolution2dFixture("NCHW", "SAME", 1, 1){}
 };
-BOOST_FIXTURE_TEST_CASE(ParseConv2DDilation1, Convolution2dDilation1Fixture)
+BOOST_FIXTURE_TEST_CASE(ParseConv2dDilation1Nchw, Convolution2dDilation1NchwFixture)
 {
     RunTest<4>({1, 2, 3, 4, 5, 6}, {2, 4, 4, 6.5f, 10 , 8.5f});
 }
 
-BOOST_AUTO_TEST_CASE(ParseConv2DDilation2)
+
+BOOST_AUTO_TEST_CASE(ParseConv2dDilation2)
 {
     const char* prototext = ""
         "node {\n"
@@ -309,8 +373,7 @@ BOOST_AUTO_TEST_CASE(ParseConv2DDilation2)
     armnn::TensorShape tensorShape = { 1, 3, 3, 1 };
     inputShapes["graphInput"] = tensorShape;
     armnnTfParser::ITfParserPtr parser = armnnTfParser::ITfParser::Create();
-    BOOST_CHECK_THROW(parser->CreateNetworkFromString(prototext, inputShapes, { "potato" }),
-                          armnn::ParseException);
+    BOOST_CHECK_THROW(parser->CreateNetworkFromString(prototext, inputShapes, { "potato" }), armnn::ParseException);
 }