1 files changed, 57 insertions, 0 deletions
diff --git a/src/armnn/optimizations/SquashEqualSiblings.hpp b/src/armnn/optimizations/SquashEqualSiblings.hpp
new file mode 100644
index 0000000000..2dfe91fdcc
--- /dev/null
+++ b/src/armnn/optimizations/SquashEqualSiblings.hpp
@@ -0,0 +1,57 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// See LICENSE file in the project root for full license information.
+//
+#pragma once
+
+#include "Optimization.hpp"
+
+namespace armnn
+{
+namespace optimizations
+{
+
+template <typename Comparable>
+class SquashEqualSiblingsImpl
+{
+public:
+    /// Run for every connection between a base Layer (any) and a child ComparableLayer.
+    /// For all siblings of the child layer that compare equal to it, bypasses and removes
+    /// them. I.e., moves the connections in the outputs of the siblings to the outputs of
+    /// the child layer, so the siblings are left unconnected (and later removed).
+    void Run(Graph& graph, InputSlot& connection) const
+    {
+        auto& child = connection.GetOwningLayer();
+
+        if (!child.IsOutputUnconnected())
+        {
+            OutputSlot& baseOutput = *connection.GetConnectedOutputSlot();
+            auto& comparableChild = *boost::polymorphic_downcast<Comparable*>(&child);
+
+            for (auto&& it : baseOutput.GetConnections())
+            {
+                Layer& sibling = it->GetOwningLayer();
+                if ((&sibling != &child) && comparableChild.IsEqual(sibling))
+                {
+                    // Bypass sibling. It will be removed as it's left unconnected.
+                    auto siblingOut = sibling.BeginOutputSlots();
+                    for (auto childOut = child.BeginOutputSlots(); childOut != child.EndOutputSlots(); ++childOut)
+                    {
+                        siblingOut->MoveAllConnections(*childOut);
+                        ++siblingOut;
+                    }
+                }
+            }
+        }
+    }
+
+protected:
+    SquashEqualSiblingsImpl() = default;
+    ~SquashEqualSiblingsImpl() = default;
+};
+
+using SquashEqualPermuteSiblings = OptimizeForConnection<Layer, PermuteLayer, SquashEqualSiblingsImpl<PermuteLayer>>;
+using SquashEqualReshapeSiblings = OptimizeForConnection<Layer, ReshapeLayer, SquashEqualSiblingsImpl<ReshapeLayer>>;
+
+} // namespace optimizations
+} // namespace armnn