runtime/onert/core/include/compiler/StaticShapeInference.h

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *      http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #ifndef __ONERT_COMPILER_STATIC_SHAPE_INFERENCE_H__
  18 #define __ONERT_COMPILER_STATIC_SHAPE_INFERENCE_H__
  19
  20 #include "ir/OperationVisitor.h"
  21 #include "ir/OpSequence.h"
  22 #include "ir/LoweredGraph.h"
  23 #include "ir/Index.h"
  24
  25 #include <memory>
  26 #include <unordered_map>
  27
  28 namespace onert
  29 {
  30 namespace compiler
  31 {
  32
  33 /**
  34  * @brief Class to infer shape before running kernels. It does the following:
  35  *        - re-calculate and set output shape at compile time (before running kernels)
  36  *        - if calculation cannot be done at compile time, mark the outputs to be dynamic, meaning
  37  *          shapes of outputs will be calculated during running kernels
  38  */
  39 class StaticShapeInferer : public ir::OperationVisitor
  40 {
  41 public:
  42   StaticShapeInferer(
  43       const ir::SubgraphIndex &subg_idx,
  44       const std::unordered_map<ir::SubgraphIndex, std::unique_ptr<ir::LoweredGraph>> &lowered_subgs)
  45       : _lowered_subgs(lowered_subgs), _operands(lowered_subgs.at(subg_idx)->graph().operands()),
  46         _operations(lowered_subgs.at(subg_idx)->graph().operations()),
  47         _return_has_dynamic_tensor(false)
  48   { /* empty */
  49   }
  50   virtual ~StaticShapeInferer() = default;
  51
  52 public:
  53   /**
  54    * @brief Infer shape of operands beloning to ops and set the output shape.
  55    *        If output shape cannot be known without running op, mark it so that it can be allocated
  56    *        when running kernel.
  57    * @param op_seq sequence of operations
  58    * @return @c true if op_seq's input or output has any dynamic tensor; @c false otherwise.
  59    */
  60   bool infer(const ir::OpSequence &op_seq)
  61   {
  62     bool has_dynamic_tensor = false;
  63
  64     _return_has_dynamic_tensor = false; // this is used as a return value inside operation's visit()
  65
  66     for (const auto &operation_idx : op_seq.operations())
  67     {
  68       _operations.at(operation_idx).accept(*this);
  69
  70       has_dynamic_tensor = has_dynamic_tensor || _return_has_dynamic_tensor;
  71     }
  72
  73     return has_dynamic_tensor;
  74   }
  75
  76   void dump();
  77
  78 private:
  79   // TODO Define visitors for operations. List them in alphabetic order.
  80   void visit(const ir::operation::Abs &op) override;
  81   void visit(const ir::operation::Add &op) override;
  82   void visit(const ir::operation::ArgMax &op) override;
  83   void visit(const ir::operation::BatchMatMul &op) override;
  84   void visit(const ir::operation::BroadcastTo &op) override;
  85   void visit(const ir::operation::Cast &op) override;
  86   void visit(const ir::operation::Comparison &op) override;
  87   void visit(const ir::operation::Concat &op) override;
  88   void visit(const ir::operation::Conv2D &op) override;
  89   void visit(const ir::operation::Cos &op) override;
  90   void visit(const ir::operation::Div &op) override;
  91   void visit(const ir::operation::Exp &op) override;
  92   void visit(const ir::operation::ExpandDims &op) override;
  93   void visit(const ir::operation::Fill &op) override;
  94   void visit(const ir::operation::FullyConnected &op) override;
  95   void visit(const ir::operation::FusedBatchNorm &op) override;
  96   void visit(const ir::operation::Gather &op) override;
  97   void visit(const ir::operation::If &op) override;
  98   void visit(const ir::operation::Log &op) override;
  99   void visit(const ir::operation::LogicalNot &op) override;
 100   void visit(const ir::operation::LogicalOr &op) override;
 101   void visit(const ir::operation::Logistic &op) override;
 102   void visit(const ir::operation::L2Normalization &op) override;
 103   void visit(const ir::operation::MatrixBandPart &op) override;
 104   void visit(const ir::operation::Max &op) override;
 105   void visit(const ir::operation::Min &op) override;
 106   void visit(const ir::operation::Mul &op) override;
 107   void visit(const ir::operation::Neg &op) override;
 108   void visit(const ir::operation::OneHot &op) override;
 109   void visit(const ir::operation::Pack &op) override;
 110   void visit(const ir::operation::Pad &op) override;
 111   void visit(const ir::operation::Permute &op) override;
 112   void visit(const ir::operation::Pow &op) override;
 113   void visit(const ir::operation::Range &op) override;
 114   void visit(const ir::operation::Reduce &op) override;
 115   void visit(const ir::operation::Reshape &op) override;
 116   void visit(const ir::operation::Round &op) override;
 117   void visit(const ir::operation::RSQRT &op) override;
 118   void visit(const ir::operation::ResizeBilinear &op) override;
 119   void visit(const ir::operation::Reverse &op) override;
 120   void visit(const ir::operation::Select &op) override;
 121   void visit(const ir::operation::Shape &op) override;
 122   void visit(const ir::operation::Sin &op) override;
 123   void visit(const ir::operation::Slice &op) override;
 124   void visit(const ir::operation::Softmax &op) override;
 125   void visit(const ir::operation::SpaceToBatchND &op) override;
 126   void visit(const ir::operation::Split &op) override;
 127   void visit(const ir::operation::Squeeze &op) override;
 128   void visit(const ir::operation::StridedSlice &op) override;
 129   void visit(const ir::operation::Sub &op) override;
 130   void visit(const ir::operation::SquaredDifference &op) override;
 131   void visit(const ir::operation::Tanh &op) override;
 132   void visit(const ir::operation::Tile &op) override;
 133   void visit(const ir::operation::Transpose &op) override;
 134   void visit(const ir::operation::Unpack &op) override;
 135   void visit(const ir::operation::While &op) override;
 136   void visit(const ir::operation::ZerosLike &op) override;
 137
 138 private:
 139   /**
 140    * @brief Performs shape inference for arithmetic operation
 141    */
 142   void handleBinaryArithmeticOp(const ir::Operation &op, const ir::OperandIndex lhs_idx,
 143                                 const ir::OperandIndex rhs_idx);
 144
 145   /**
 146    * @brief Performs shape inference for unary op whose output shape is
 147    *        always same with input shape
 148    */
 149   void handleSimpleUnaryOp(const ir::Operation &op, const ir::OperandIndex input_idx);
 150
 151 private:
 152   const std::unordered_map<ir::SubgraphIndex, std::unique_ptr<ir::LoweredGraph>> &_lowered_subgs;
 153   // _operands and _operations can be changed by controlflow operation
 154   ir::Operands &_operands;     // operands of current subgraph
 155   ir::Operations &_operations; // operations of current subgraph
 156   bool _return_has_dynamic_tensor;
 157 };
 158
 159 } // namespace compiler
 160 } // namespace onert
 161
 162 #endif // __ONERT_COMPILER_STATIC_SHAPE_INFERENCE_H__