IVGCVSW-4496 Add Flatten support to ONNX parser

 * Added ParseFlatten method
 * Added Read int64 attribute method
 - Modified ComputeReshapeInfo method
 - Modified ParseReshape
 * Added unit tests
 - Reorganised OnnxParser.cpp/.hpp

Signed-off-by: Ryan OShea <Ryan.OShea2@arm.com>
Change-Id: I8a9553438dd1e8c702d821b093587e0074c027d5

2 files changed, 548 insertions, 0 deletions

diff --git a/src/armnnOnnxParser/test/Flatten.cpp b/src/armnnOnnxParser/test/Flatten.cpp
new file mode 100644
index 0000000000..1ba509e5e6
--- /dev/null
+++ b/src/armnnOnnxParser/test/Flatten.cpp
@@ -0,0 +1,443 @@
+//
+// Copyright © 2020 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include <boost/test/unit_test.hpp>
+#include "armnnOnnxParser/IOnnxParser.hpp"
+#include  "ParserPrototxtFixture.hpp"
+
+BOOST_AUTO_TEST_SUITE(OnnxParser)
+
+struct FlattenMainFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    FlattenMainFixture(const std::string& dataType)
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: )" + dataType + R"(
+                            shape {
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                     node {
+                         input: "Input"
+                         output: "Output"
+                         name: "flatten"
+                         op_type: "Flatten"
+                         attribute {
+                           name: "axis"
+                           i: 2
+                           type: INT
+                         }
+                      }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 4
+                                   }
+                                   dim {
+                                       dim_value: 9
+                                   }
+                               }
+                            }
+                          }
+                       }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
+struct FlattenDefaultAxisFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    FlattenDefaultAxisFixture(const std::string& dataType)
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: )" + dataType + R"(
+                            shape {
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                     node {
+                         input: "Input"
+                         output: "Output"
+                         name: "flatten"
+                         op_type: "Flatten"
+                      }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 2
+                                   }
+                                   dim {
+                                       dim_value: 18
+                                   }
+                               }
+                            }
+                          }
+                       }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
+struct FlattenAxisZeroFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    FlattenAxisZeroFixture(const std::string& dataType)
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: )" + dataType + R"(
+                            shape {
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                     node {
+                         input: "Input"
+                         output: "Output"
+                         name: "flatten"
+                         op_type: "Flatten"
+                         attribute {
+                           name: "axis"
+                           i: 0
+                           type: INT
+                         }
+                      }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 1
+                                   }
+                                   dim {
+                                       dim_value: 36
+                                   }
+                               }
+                            }
+                          }
+                       }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
+struct FlattenNegativeAxisFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    FlattenNegativeAxisFixture(const std::string& dataType)
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: )" + dataType + R"(
+                            shape {
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                     node {
+                         input: "Input"
+                         output: "Output"
+                         name: "flatten"
+                         op_type: "Flatten"
+                         attribute {
+                           name: "axis"
+                           i: -1
+                           type: INT
+                         }
+                      }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 12
+                                   }
+                                   dim {
+                                       dim_value: 3
+                                   }
+                               }
+                            }
+                          }
+                       }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
+struct FlattenInvalidNegativeAxisFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    FlattenInvalidNegativeAxisFixture(const std::string& dataType)
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: )" + dataType + R"(
+                            shape {
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                     node {
+                         input: "Input"
+                         output: "Output"
+                         name: "flatten"
+                         op_type: "Flatten"
+                         attribute {
+                           name: "axis"
+                           i: -5
+                           type: INT
+                         }
+                      }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 12
+                                   }
+                                   dim {
+                                       dim_value: 3
+                                   }
+                               }
+                            }
+                          }
+                       }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
+struct FlattenValidFixture : FlattenMainFixture
+{
+    FlattenValidFixture() : FlattenMainFixture("1") {
+        Setup();
+    }
+};
+
+struct FlattenDefaultValidFixture : FlattenDefaultAxisFixture
+{
+    FlattenDefaultValidFixture() : FlattenDefaultAxisFixture("1") {
+        Setup();
+    }
+};
+
+struct FlattenAxisZeroValidFixture : FlattenAxisZeroFixture
+{
+    FlattenAxisZeroValidFixture() : FlattenAxisZeroFixture("1") {
+        Setup();
+    }
+};
+
+struct FlattenNegativeAxisValidFixture : FlattenNegativeAxisFixture
+{
+    FlattenNegativeAxisValidFixture() : FlattenNegativeAxisFixture("1") {
+        Setup();
+    }
+};
+
+struct FlattenInvalidFixture : FlattenMainFixture
+{
+    FlattenInvalidFixture() : FlattenMainFixture("10") { }
+};
+
+struct FlattenInvalidAxisFixture : FlattenInvalidNegativeAxisFixture
+{
+    FlattenInvalidAxisFixture() : FlattenInvalidNegativeAxisFixture("1") { }
+};
+
+BOOST_FIXTURE_TEST_CASE(ValidFlattenTest, FlattenValidFixture)
+{
+    RunTest<2>({{"Input",
+                          { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                            1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                            1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}},
+                {{"Output",
+                          { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                            1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                            1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}});
+}
+
+BOOST_FIXTURE_TEST_CASE(ValidFlattenDefaultTest, FlattenDefaultValidFixture)
+{
+    RunTest<2>({{"Input",
+                    { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}},
+               {{"Output",
+                    { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}});
+}
+
+BOOST_FIXTURE_TEST_CASE(ValidFlattenAxisZeroTest, FlattenAxisZeroValidFixture)
+{
+    RunTest<2>({{"Input",
+                    { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}},
+               {{"Output",
+                    { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}});
+}
+
+BOOST_FIXTURE_TEST_CASE(ValidFlattenNegativeAxisTest, FlattenNegativeAxisValidFixture)
+{
+    RunTest<2>({{"Input",
+                    { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}},
+               {{"Output",
+                    { 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                        1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f }}});
+}
+
+BOOST_FIXTURE_TEST_CASE(IncorrectDataTypeFlatten, FlattenInvalidFixture)
+{
+    BOOST_CHECK_THROW(Setup(), armnn::ParseException);
+}
+
+BOOST_FIXTURE_TEST_CASE(IncorrectAxisFlatten, FlattenInvalidAxisFixture)
+{
+    BOOST_CHECK_THROW(Setup(), armnn::ParseException);
+}
+
+BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/armnnOnnxParser/test/Reshape.cpp b/src/armnnOnnxParser/test/Reshape.cpp
index 60937f0666..119a406d7e 100644
--- a/src/armnnOnnxParser/test/Reshape.cpp
+++ b/src/armnnOnnxParser/test/Reshape.cpp
@@ -85,6 +85,91 @@ struct ReshapeMainFixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxPa
     }
 };
 
+struct ReshapeRank4Fixture : public armnnUtils::ParserPrototxtFixture<armnnOnnxParser::IOnnxParser>
+{
+    ReshapeRank4Fixture(const std::string& dataType)
+    {
+        m_Prototext = R"(
+                   ir_version: 3
+                   producer_name:  "CNTK"
+                   producer_version:  "2.5.1"
+                   domain:  "ai.cntk"
+                   model_version: 1
+                   graph {
+                     name:  "CNTKGraph"
+                     input {
+                        name: "Input"
+                        type {
+                          tensor_type {
+                            elem_type: )" + dataType + R"(
+                            shape {
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 2
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                              dim {
+                                dim_value: 3
+                              }
+                            }
+                          }
+                        }
+                      }
+                      input {
+                         name: "Shape"
+                         type {
+                           tensor_type {
+                             elem_type: 7
+                             shape {
+                               dim {
+                                 dim_value: 2
+                               }
+                             }
+                           }
+                         }
+                       }
+                     node {
+                         input: "Input"
+                         input: "Shape"
+                         output: "Output"
+                         name: "reshape"
+                         op_type: "Reshape"
+
+                      }
+                      initializer {
+                        dims: 2
+                        data_type: 7
+                        int64_data: 2
+                        int64_data: 2
+                        name: "Shape"
+                     }
+                      output {
+                          name: "Output"
+                          type {
+                             tensor_type {
+                               elem_type: 1
+                               shape {
+                                   dim {
+                                       dim_value: 6
+                                   }
+                                   dim {
+                                       dim_value: 6
+                                   }
+                               }
+                            }
+                          }
+                       }
+                    }
+                   opset_import {
+                      version: 7
+                    })";
+    }
+};
+
 struct ReshapeValidFixture : ReshapeMainFixture
 {
     ReshapeValidFixture() : ReshapeMainFixture("1") {
@@ -92,6 +177,13 @@ struct ReshapeValidFixture : ReshapeMainFixture
     }
 };
 
+struct ReshapeValidRank4Fixture : ReshapeRank4Fixture
+{
+    ReshapeValidRank4Fixture() : ReshapeRank4Fixture("1") {
+        Setup();
+    }
+};
+
 struct ReshapeInvalidFixture : ReshapeMainFixture
 {
     ReshapeInvalidFixture() : ReshapeMainFixture("10") { }
@@ -102,6 +194,19 @@ BOOST_FIXTURE_TEST_CASE(ValidReshapeTest, ReshapeValidFixture)
     RunTest<2>({{"Input", { 0.0f, 1.0f, 2.0f, 3.0f }}}, {{"Output", { 0.0f, 1.0f, 2.0f, 3.0f }}});
 }
 
+BOOST_FIXTURE_TEST_CASE(ValidRank4ReshapeTest, ReshapeValidRank4Fixture)
+{
+    RunTest<2>(
+        {{"Input",
+                   {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                    1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                    1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}}},
+        {{"Output",
+                    {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                     1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f,
+                     1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}}});
+}
+
 BOOST_FIXTURE_TEST_CASE(IncorrectDataTypeReshape, ReshapeInvalidFixture)
 {
    BOOST_CHECK_THROW(Setup(), armnn::ParseException);