src/armnn/optimizations/SquashEqualSiblings.hpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // See LICENSE file in the project root for full license information.
   4 //
   5 #pragma once
   6
   7 #include "Optimization.hpp"
   8
   9 namespace armnn
  10 {
  11 namespace optimizations
  12 {
  13
  14 template <typename Comparable>
  15 class SquashEqualSiblingsImpl
  16 {
  17 public:
  18     /// Run for every connection between a base Layer (any) and a child ComparableLayer.
  19     /// For all siblings of the child layer that compare equal to it, bypasses and removes
  20     /// them. I.e., moves the connections in the outputs of the siblings to the outputs of
  21     /// the child layer, so the siblings are left unconnected (and later removed).
  22     void Run(Graph& graph, InputSlot& connection) const
  23     {
  24         auto& child = connection.GetOwningLayer();
  25
  26         if (!child.IsOutputUnconnected())
  27         {
  28             OutputSlot& baseOutput = *connection.GetConnectedOutputSlot();
  29             auto& comparableChild = *boost::polymorphic_downcast<Comparable*>(&child);
  30
  31             for (auto&& it : baseOutput.GetConnections())
  32             {
  33                 Layer& sibling = it->GetOwningLayer();
  34                 if ((&sibling != &child) && comparableChild.IsEqual(sibling))
  35                 {
  36                     // Bypass sibling. It will be removed as it's left unconnected.
  37                     auto siblingOut = sibling.BeginOutputSlots();
  38                     for (auto childOut = child.BeginOutputSlots(); childOut != child.EndOutputSlots(); ++childOut)
  39                     {
  40                         siblingOut->MoveAllConnections(*childOut);
  41                         ++siblingOut;
  42                     }
  43                 }
  44             }
  45         }
  46     }
  47
  48 protected:
  49     SquashEqualSiblingsImpl() = default;
  50     ~SquashEqualSiblingsImpl() = default;
  51 };
  52
  53 using SquashEqualPermuteSiblings = OptimizeForConnection<Layer, PermuteLayer, SquashEqualSiblingsImpl<PermuteLayer>>;
  54 using SquashEqualReshapeSiblings = OptimizeForConnection<Layer, ReshapeLayer, SquashEqualSiblingsImpl<ReshapeLayer>>;
  55
  56 } // namespace optimizations
  57 } // namespace armnn