runtime/onert/core/include/ir/NNPkg.h

   1 /*
   2  * Copyright (c) 2022 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_IR_NNPKG_H__
  18 #define __ONERT_IR_NNPKG_H__
  19
  20 #include <memory>
  21 #include <unordered_set>
  22 #include <vector>
  23
  24 #include "ir/Graph.h"
  25 #include "ir/Index.h"
  26 #include "ir/Model.h"
  27
  28 namespace onert
  29 {
  30 namespace ir
  31 {
  32
  33 using IODesc = std::tuple<ModelIndex, SubgraphIndex, IOIndex>;
  34
  35 struct ModelEdge
  36 {
  37   IODesc from;
  38   IODesc to;
  39 };
  40
  41 struct ModelEdgeEqual
  42 {
  43   bool operator()(const onert::ir::ModelEdge &lhs, const onert::ir::ModelEdge &rhs) const
  44   {
  45     return lhs.from == rhs.from && lhs.to == rhs.to;
  46   }
  47 };
  48
  49 struct ModelEdgeHash
  50 {
  51   size_t operator()(const ::onert::ir::ModelEdge &edge) const noexcept
  52   {
  53     unsigned long long h1 = (std::get<0>(edge.from).value() << 24) |
  54                             (std::get<1>(edge.from).value() << 16) | std::get<2>(edge.from).value();
  55     unsigned long long h2 = (std::get<0>(edge.to).value() << 24) |
  56                             (std::get<1>(edge.to).value() << 16) | std::get<2>(edge.to).value();
  57     return h1 + h2;
  58   }
  59 };
  60
  61 inline std::ostream &operator<<(std::ostream &o, const IODesc &od)
  62 {
  63   o << std::get<0>(od).value() << ":" << std::get<1>(od).value() << ":" << std::get<2>(od).value();
  64   return o;
  65 }
  66
  67 using ModelEdgeSet = std::unordered_set<ir::ModelEdge, ir::ModelEdgeHash, ir::ModelEdgeEqual>;
  68
  69 /**
  70  * @brief Struct to gather model I/O information in multimodel NN package
  71  *        Model I/O will have role one of below
  72  *        - Package input/output
  73  *        - Edge's start/finish point between model
  74  */
  75 struct ModelEdges
  76 {
  77   std::vector<ir::IODesc> pkg_inputs;
  78   std::vector<ir::IODesc> pkg_outputs;
  79   ModelEdgeSet edges;
  80 };
  81
  82 class NNPkg
  83 {
  84 public:
  85   NNPkg() = default;
  86   NNPkg(const NNPkg &obj) = default;
  87   NNPkg(NNPkg &&) = default;
  88   NNPkg &operator=(const NNPkg &) = default;
  89   NNPkg &operator=(NNPkg &&) = default;
  90   ~NNPkg() = default;
  91
  92   NNPkg(std::shared_ptr<Model> model) { _models[ModelIndex{0}] = model; }
  93   std::shared_ptr<Model> primary_model() const { return _models.at(onert::ir::ModelIndex{0}); }
  94
  95   /**
  96    * @brief Put model at index
  97    *
  98    * @param[in] model Model to be pushed
  99    * @param[in] index Index where Model is to be pushed
 100    */
 101   void push(ModelIndex index, const std::shared_ptr<Model> &model) { _models[index] = model; }
 102
 103   /**
 104    * @brief Get the count of model
 105    *
 106    * @return the count of models
 107    */
 108   size_t model_count() const { return _models.size(); }
 109
 110   /**
 111    * @brief Get model at index
 112    *
 113    * @param[in] index Index of the model to be returned
 114    * @return Model at index
 115    */
 116   const std::shared_ptr<Model> &model(const ModelIndex &index) const { return _models.at(index); }
 117   /**
 118    * @brief Get model at index
 119    *
 120    * @param[in] index Index of the model to be returned
 121    * @return Model at index
 122    */
 123   std::shared_ptr<Model> &model(const ModelIndex &index) { return _models.at(index); }
 124
 125   /**
 126    * @brief Get pkg_input at index
 127    *
 128    * @param[in] index Index of pkg_input to be returned
 129    * @return IODesc at index
 130    */
 131   const IODesc &input(uint32_t index) const { return _edges.pkg_inputs[index]; }
 132   /**
 133    * @brief Get pkg_input at index
 134    *
 135    * @param[in] index Index of pkg_input to be returned
 136    * @return IODesc at index
 137    */
 138   IODesc &input(uint32_t index) { return _edges.pkg_inputs[index]; }
 139   /**
 140    * @brief Add input at the end
 141    *
 142    * @param[in] input Input IODesc to be pushed
 143    */
 144   void addInput(const IODesc &input) { _edges.pkg_inputs.push_back(input); }
 145
 146   /**
 147    * @brief Get pkg_output at index
 148    *
 149    * @param[in] index Index of pkg_output to be returned
 150    * @return IODesc at index
 151    */
 152   const IODesc &output(uint32_t index) const { return _edges.pkg_outputs[index]; }
 153   /**
 154    * @brief Get pkg_output at index
 155    *
 156    * @param[in] index Index of pkg_output to be returned
 157    * @return IODesc at index
 158    */
 159   IODesc &output(uint32_t index) { return _edges.pkg_outputs[index]; }
 160   /**
 161    * @brief Add output at the end
 162    *
 163    * @param[in] output Output IODesc to be pushed
 164    */
 165   void addOutput(const IODesc &output) { _edges.pkg_outputs.push_back(output); }
 166
 167   /**
 168    * @brief Add edge between models at the end
 169    *
 170    * @param[in] from from IODesc
 171    * @param[in] to   to IODesc
 172    */
 173   void addEdge(const IODesc &from, const IODesc &to)
 174   {
 175     std::cout << from << " -> " << to << std::endl;
 176     _edges.edges.insert(ModelEdge{from, to});
 177   }
 178   /**
 179    * @brief   Get model edge set
 180    * @return  Edge set reference
 181    */
 182   const ModelEdges &model_edges() { return _edges; }
 183
 184   /**
 185    * @brief Verify NNPkg
 186    *
 187    */
 188   void verify(void)
 189   {
 190     // Verify edges information
 191     //
 192     // Only duplicates of nnpkg output and Edge `from` are possible.
 193     // | Whether duplicates are possible   | Edge `to` | Edge `from` |
 194     // | nnpkg input  (input of subgraph)  | X (*1)    | X (*2)      |
 195     // | nnpkg output (output of subgraph) | X (*2)    | O           |
 196     // *1. The subjects who determine values of each buffer are different.
 197     //    - nnpkg input : user input
 198     //    - Edge `to`   : output of another subgraph
 199     // *2. `IOIndex` of inputs and outputs of subgraph is distinct.
 200     //
 201     for (const auto &edge : _edges.edges)
 202     {
 203       if (std::find(_edges.pkg_inputs.begin(), _edges.pkg_inputs.end(), edge.to) !=
 204           _edges.pkg_inputs.end())
 205       {
 206         throw std::runtime_error{
 207           "Invalid edge information. NNPkg inputs and Edge `to` cannot be duplicated"};
 208       }
 209     }
 210   }
 211
 212   // TODO Find better way to handle single model NNPackage and multi model NNPackage on inputSize(),
 213   //      outputSize(), inputInfo(), outputInfo()
 214
 215   /**
 216    * @brief   Get model input size
 217    */
 218   uint32_t inputSize() const
 219   {
 220     return _models.size() == 1 ? primary_model()->primary_subgraph()->getInputs().size()
 221                                : _edges.pkg_inputs.size();
 222   }
 223
 224   /**
 225    * @brief   Get model output size
 226    */
 227   uint32_t outputSize() const
 228   {
 229     return _models.size() == 1 ? primary_model()->primary_subgraph()->getOutputs().size()
 230                                : _edges.pkg_outputs.size();
 231   }
 232
 233   /**
 234    * @brief   Get model input info
 235    */
 236   OperandInfo &inputInfo(uint32_t index) const
 237   {
 238     if (_models.size() == 1)
 239     {
 240       auto const graph = primary_model()->primary_subgraph();
 241       auto const operand_index = graph->getInputs().at(index);
 242       return graph->operands().at(operand_index).info();
 243     }
 244
 245     auto const &desc = input(index);
 246     auto const graph = model(std::get<ModelIndex>(desc))->primary_subgraph();
 247     auto const operand_index = graph->getInputs().at(std::get<IOIndex>(desc).value());
 248     return graph->operands().at(operand_index).info();
 249   }
 250
 251   /**
 252    * @brief   Get model output info
 253    */
 254   OperandInfo &outputInfo(uint32_t index) const
 255   {
 256     if (_models.size() == 1)
 257     {
 258       auto const graph = primary_model()->primary_subgraph();
 259       auto const operand_index = graph->getOutputs().at(index);
 260       return graph->operands().at(operand_index).info();
 261     }
 262
 263     auto const &desc = output(index);
 264     auto const graph = model(std::get<ModelIndex>(desc))->primary_subgraph();
 265     auto const operand_index = graph->getOutputs().at(std::get<IOIndex>(desc).value());
 266     return graph->operands().at(operand_index).info();
 267   }
 268
 269   // TODO: Add iterate() or getter for edges
 270
 271 private:
 272   std::unordered_map<ModelIndex, std::shared_ptr<Model>> _models;
 273   ModelEdges _edges;
 274 };
 275
 276 } // namespace ir
 277 } // namespace onert
 278
 279 namespace std
 280 {
 281
 282 template <> struct hash<onert::ir::IODesc>
 283 {
 284   size_t operator()(const ::onert::ir::IODesc &iodesc) const noexcept
 285   {
 286     return (std::get<0>(iodesc).value() << 24) | (std::get<1>(iodesc).value() << 16) |
 287            std::get<2>(iodesc).value();
 288   }
 289 };
 290
 291 } // namespace std
 292
 293 #endif // __ONERT_IR_NNPKG_H__