IVGCVSW-2616: static quantization of reshape

Change-Id: Iaefa145b8a218b3708c8a9208fdabd4984376ae9
Signed-off-by: Nina Drozd <nina.drozd@arm.com>

5 files changed, 57 insertions, 0 deletions

diff --git a/src/armnn/QuantizerVisitor.cpp b/src/armnn/QuantizerVisitor.cpp
index 88062068e8..95b593971d 100644
--- a/src/armnn/QuantizerVisitor.cpp
+++ b/src/armnn/QuantizerVisitor.cpp
@@ -262,4 +262,13 @@ void QuantizerVisitor::VisitConstantLayer(const IConnectableLayer* layer,
     RecordLayer(layer, newLayer);
 }
 
+void QuantizerVisitor::VisitReshapeLayer(const IConnectableLayer* layer,
+                                         const ReshapeDescriptor& reshapeDescriptor,
+                                         const char* name)
+{
+    IConnectableLayer* newLayer = m_QuantizedNetwork->AddReshapeLayer(reshapeDescriptor, name);
+    RecordLayer(layer, newLayer);
+    SetQuantizedInputConnections(layer, newLayer);
+}
+
 } //namespace armnn
diff --git a/src/armnn/QuantizerVisitor.hpp b/src/armnn/QuantizerVisitor.hpp
index ac90bfb1e5..aaa888e5a0 100644
--- a/src/armnn/QuantizerVisitor.hpp
+++ b/src/armnn/QuantizerVisitor.hpp
@@ -88,6 +88,10 @@ public:
                           const OriginsDescriptor& mergerDescriptor,
                           const char* name = nullptr) override;
 
+    void VisitReshapeLayer(const IConnectableLayer* layer,
+                           const ReshapeDescriptor& reshapeDescriptor,
+                           const char* name = nullptr) override;
+
     /// Extract the quantized network
     INetworkPtr RetrieveFinalNetwork() { return std::move(m_QuantizedNetwork); }
 
diff --git a/src/armnn/StaticRangeVisitor.cpp b/src/armnn/StaticRangeVisitor.cpp
index 44b05ca6cb..4b715bdfb6 100644
--- a/src/armnn/StaticRangeVisitor.cpp
+++ b/src/armnn/StaticRangeVisitor.cpp
@@ -198,4 +198,13 @@ void StaticRangeVisitor::VisitMergerLayer(const IConnectableLayer* layer,
     SetRange(layer, 0, min, max);
 }
 
+void StaticRangeVisitor::VisitReshapeLayer(const IConnectableLayer* layer,
+                                           const ReshapeDescriptor& reshapeDescriptor,
+                                           const char* name)
+{
+    boost::ignore_unused(reshapeDescriptor);
+    boost::ignore_unused(name);
+    ForwardParentParameters(layer);
+}
+
 } //namespace armnn
diff --git a/src/armnn/StaticRangeVisitor.hpp b/src/armnn/StaticRangeVisitor.hpp
index 7576e96032..145db20fc3 100644
--- a/src/armnn/StaticRangeVisitor.hpp
+++ b/src/armnn/StaticRangeVisitor.hpp
@@ -79,6 +79,10 @@ public:
                           const OriginsDescriptor& mergerDescriptor,
                           const char* name = nullptr) override;
 
+    void VisitReshapeLayer(const IConnectableLayer* layer,
+                           const ReshapeDescriptor& reshapeDescriptor,
+                           const char* name = nullptr) override;
+
 private:
     /// Set the range for an output slot on a layer
     void SetRange(const IConnectableLayer* layer, unsigned int outputIdx, float min, float max);
diff --git a/src/armnn/test/QuantizerTest.cpp b/src/armnn/test/QuantizerTest.cpp
index a46b443fa7..b73ac20457 100644
--- a/src/armnn/test/QuantizerTest.cpp
+++ b/src/armnn/test/QuantizerTest.cpp
@@ -1052,5 +1052,36 @@ BOOST_AUTO_TEST_CASE(QuantizeMerger)
     VisitLayersTopologically(quantizedNetwork.get(), validator);
 }
 
+BOOST_AUTO_TEST_CASE(QuantizeReshape)
+{
+    class TestReshapeQuantization : public TestLeakyReLuActivationQuantization
+    {
+    public:
+        virtual void VisitReshapeLayer(const IConnectableLayer* layer,
+                                       const ReshapeDescriptor& reshapeDescriptor,
+                                       const char* name = nullptr) override
+        {
+            CheckForwardedQuantizationSettings(layer);
+        }
+    };
+
+    INetworkPtr network = INetwork::Create();
+
+    TensorShape shape{1U};
+    TensorInfo info(shape, DataType::Float32);
+
+    IConnectableLayer* activation = CreateStartOfLeakyReluNetwork(network.get(), info);
+
+    // Add the layer under test
+    ReshapeDescriptor descriptor({1, 2, 3, 4});
+    IConnectableLayer* reshape = network->AddReshapeLayer(descriptor);
+
+    CompleteLeakyReluNetwork(network.get(), activation, reshape, info);
+
+    auto quantizedNetwork = INetworkQuantizer::Create(network.get())->ExportNetwork();
+    TestReshapeQuantization validator;
+    VisitLayersTopologically(quantizedNetwork.get(), validator);
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 } // namespace armnn