src/armnn/optimizations/Optimization.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 "Graph.hpp"
   8 #include "LayersFwd.hpp"
   9
  10 namespace armnn
  11 {
  12
  13 class Optimization
  14 {
  15 public:
  16     virtual void Run(Graph& graph, Layer& base) const = 0;
  17 protected:
  18     ~Optimization() = default;
  19 };
  20
  21 // Wrappers
  22 // The implementation of the following wrappers make use of the CRTP C++ idiom
  23 // (curiously recurring template pattern).
  24 // For details, see https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern
  25
  26 /// Wrapper Optimization base class that calls Wrapped::Run() for every layer of type BaseType.
  27 /// - Wrapped class mustn't remove the base layer. The optimizer will remove it if left unconnected
  28 ///   after applying each optimization.
  29 template <typename BaseType, typename Wrapped>
  30 class OptimizeForTypeImpl : public armnn::Optimization, public Wrapped
  31 {
  32 public:
  33     using Wrapped::Wrapped;
  34
  35     void Run(Graph& graph, Layer& base) const override
  36     {
  37         if (base.GetType() == LayerEnumOf<BaseType>())
  38         {
  39             Wrapped::Run(graph, *boost::polymorphic_downcast<BaseType*>(&base));
  40         }
  41     }
  42
  43 protected:
  44     ~OptimizeForTypeImpl() = default;
  45 };
  46
  47 /// Specialization that calls Wrapped::Run() for any layer type
  48 template <typename Wrapped>
  49 class OptimizeForTypeImpl<Layer, Wrapped> : public armnn::Optimization, public Wrapped
  50 {
  51 public:
  52     using Wrapped::Wrapped;
  53
  54     void Run(Graph& graph, Layer& base) const override
  55     {
  56         Wrapped::Run(graph, base);
  57     }
  58
  59 protected:
  60     ~OptimizeForTypeImpl() = default;
  61 };
  62
  63 template <typename BaseType, typename Wrapped>
  64 class OptimizeForType final : public OptimizeForTypeImpl<BaseType, Wrapped>
  65 {
  66 public:
  67     using OptimizeForTypeImpl<BaseType, Wrapped>::OptimizeForTypeImpl;
  68 };
  69
  70 /// Wrapper Optimization class that calls Wrapped::Run for every connection BaseType -> ChildType.
  71 /// - Wrapped class mustn't remove the base layer. The optimizer will remove it if left unconnected
  72 ///   after applying each optimization.
  73 /// - Wrapped class mustn't affect existing connections in the same output. It might add new ones.
  74 /// - Children layers are removed if left unconnected after applying the wrapped optimization.
  75 template <typename BaseType, typename ChildType, typename Wrapped>
  76 class OptimizeForConnectionImpl : public Wrapped
  77 {
  78 public:
  79     using Wrapped::Wrapped;
  80
  81     void Run(Graph& graph, BaseType& base) const
  82     {
  83         for (auto output = base.BeginOutputSlots(); output != base.EndOutputSlots(); ++output)
  84         {
  85             for (auto&& childInput : output->GetConnections())
  86             {
  87                 if (childInput->GetOwningLayer().GetType() == LayerEnumOf<ChildType>())
  88                 {
  89                     Wrapped::Run(graph, *childInput);
  90                 }
  91             }
  92
  93             // Remove unconnected children
  94             for (unsigned int i = 0; i < output->GetNumConnections();)
  95             {
  96                 Layer* child = &output->GetConnection(i)->GetOwningLayer();
  97
  98                 if (child->IsOutputUnconnected())
  99                 {
 100                     graph.EraseLayer(child);
 101                 }
 102                 else
 103                 {
 104                     ++i;
 105                 }
 106             }
 107         }
 108     }
 109
 110 protected:
 111     ~OptimizeForConnectionImpl() = default;
 112 };
 113
 114 template <typename BaseType, typename ChildType, typename Wrapped>
 115 class OptimizeForConnection final
 116     : public OptimizeForTypeImpl<BaseType, OptimizeForConnectionImpl<BaseType, ChildType, Wrapped>>
 117 {
 118 public:
 119     using OptimizeForTypeImpl<BaseType, OptimizeForConnectionImpl<BaseType, ChildType, Wrapped>>::OptimizeForTypeImpl;
 120 };
 121
 122 } // namespace armnn