runtime/onert/core/include/compiler/StaticShapeInferer.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_INFERER_H__
  18 #define __ONERT_COMPILER_STATIC_SHAPE_INFERER_H__
  19
  20 #include "ir/OperationVisitor.h"
  21 #include "ir/OpSequence.h"
  22 #include "compiler/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<compiler::LoweredGraph>>
  45           &lowered_subgs)
  46       : _lowered_subgs(lowered_subgs), _operands(lowered_subgs.at(subg_idx)->graph().operands()),
  47         _operations(lowered_subgs.at(subg_idx)->graph().operations()),
  48         _return_has_dynamic_tensor(false)
  49   { /* empty */
  50   }
  51   virtual ~StaticShapeInferer() = default;
  52
  53 public:
  54   /**
  55    * @brief Infer shape of operands beloning to ops and set the output shape.
  56    *        If output shape cannot be known without running op, mark it so that it can be allocated
  57    *        when running kernel.
  58    * @param op_seq sequence of operations
  59    * @return @c true if op_seq's input or output has any dynamic tensor; @c false otherwise.
  60    */
  61   bool infer(const ir::OpSequence &op_seq);
  62
  63   void dump();
  64
  65 private:
  66   bool checkDynamicInput(const ir::Operation &op);
  67   void setDynamicOutput(const ir::Operation &op);
  68
  69 private:
  70   // TODO Define visitors for operations. List them in alphabetic order.
  71   void visit(const ir::operation::ArgMinMax &op) override;
  72   void visit(const ir::operation::BatchMatMul &op) override;
  73   void visit(const ir::operation::BCQFullyConnected &op) override;
  74   void visit(const ir::operation::BCQGather &op) override;
  75   void visit(const ir::operation::BinaryArithmetic &op) override;
  76   void visit(const ir::operation::BroadcastTo &op) override;
  77   void visit(const ir::operation::Comparison &op) override;
  78   void visit(const ir::operation::Concat &op) override;
  79   void visit(const ir::operation::Conv2D &op) override;
  80   void visit(const ir::operation::ElementwiseActivation &op) override;
  81   void visit(const ir::operation::ElementwiseBinary &op) override;
  82   void visit(const ir::operation::ElementwiseUnary &op) override;
  83   void visit(const ir::operation::ExpandDims &op) override;
  84   void visit(const ir::operation::Fill &op) override;
  85   void visit(const ir::operation::FullyConnected &op) override;
  86   void visit(const ir::operation::FusedBatchNorm &op) override;
  87   void visit(const ir::operation::Gather &op) override;
  88   void visit(const ir::operation::If &op) override;
  89   void visit(const ir::operation::L2Normalization &op) override;
  90   void visit(const ir::operation::LSTM &op) override;
  91   void visit(const ir::operation::MatrixBandPart &op) override;
  92   void visit(const ir::operation::OneHot &op) override;
  93   void visit(const ir::operation::Pack &op) override;
  94   void visit(const ir::operation::Pad &op) override;
  95   void visit(const ir::operation::Permute &op) override;
  96   void visit(const ir::operation::Pow &op) override;
  97   void visit(const ir::operation::Range &op) override;
  98   void visit(const ir::operation::Reduce &op) override;
  99   void visit(const ir::operation::Reshape &op) override;
 100   void visit(const ir::operation::ResizeBilinear &op) override;
 101   void visit(const ir::operation::Reverse &op) override;
 102   void visit(const ir::operation::Select &op) override;
 103   void visit(const ir::operation::Shape &op) override;
 104   void visit(const ir::operation::Slice &op) override;
 105   void visit(const ir::operation::Softmax &op) override;
 106   void visit(const ir::operation::SpaceToBatchND &op) override;
 107   void visit(const ir::operation::Split &op) override;
 108   void visit(const ir::operation::Squeeze &op) override;
 109   void visit(const ir::operation::StridedSlice &op) override;
 110   void visit(const ir::operation::SquaredDifference &op) override;
 111   void visit(const ir::operation::Tile &op) override;
 112   void visit(const ir::operation::Transpose &op) override;
 113   void visit(const ir::operation::Unpack &op) override;
 114   void visit(const ir::operation::While &op) override;
 115
 116 private:
 117   /**
 118    * @brief Performs shape inference for arithmetic operation
 119    */
 120   void handleBinaryArithmeticOp(const ir::Operation &op, const ir::OperandIndex lhs_idx,
 121                                 const ir::OperandIndex rhs_idx);
 122
 123   /**
 124    * @brief Performs shape inference for unary op whose output shape is
 125    *        always same with input shape
 126    */
 127   void handleSimpleUnaryOp(const ir::Operation &op, const ir::OperandIndex input_idx);
 128
 129 private:
 130   const std::unordered_map<ir::SubgraphIndex, std::unique_ptr<compiler::LoweredGraph>>
 131       &_lowered_subgs;
 132   // _operands and _operations can be changed by controlflow operation
 133   ir::Operands &_operands;     // operands of current subgraph
 134   ir::Operations &_operations; // operations of current subgraph
 135   bool _return_has_dynamic_tensor;
 136 };
 137
 138 } // namespace compiler
 139 } // namespace onert
 140
 141 #endif // __ONERT_COMPILER_STATIC_SHAPE_INFERER_H__