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/Index.h"
  25 #include "ir/Model.h"
  26
  27 namespace onert
  28 {
  29 namespace ir
  30 {
  31
  32 using IODesc = std::tuple<ModelIndex, SubgraphIndex, IOIndex>;
  33
  34 struct ModelEdge
  35 {
  36   IODesc from;
  37   IODesc to;
  38 };
  39
  40 struct ModelEdgeEqual
  41 {
  42   bool operator()(const onert::ir::ModelEdge &lhs, const onert::ir::ModelEdge &rhs) const
  43   {
  44     return lhs.from == rhs.from && lhs.to == rhs.to;
  45   }
  46 };
  47
  48 struct ModelEdgeHash
  49 {
  50   size_t operator()(const ::onert::ir::ModelEdge &edge) const noexcept
  51   {
  52     unsigned long long h1 = (std::get<0>(edge.from).value() << 24) |
  53                             (std::get<1>(edge.from).value() << 16) | std::get<2>(edge.from).value();
  54     unsigned long long h2 = (std::get<0>(edge.to).value() << 24) |
  55                             (std::get<1>(edge.to).value() << 16) | std::get<2>(edge.to).value();
  56     return h1 + h2;
  57   }
  58 };
  59
  60 inline std::ostream &operator<<(std::ostream &o, const IODesc &od)
  61 {
  62   o << std::get<0>(od).value() << ":" << std::get<1>(od).value() << ":" << std::get<2>(od).value();
  63   return o;
  64 }
  65
  66 using ModelEdgeSet = std::unordered_set<ir::ModelEdge, ir::ModelEdgeHash, ir::ModelEdgeEqual>;
  67
  68 /**
  69  * @brief Struct to gather model I/O information in multimodel NN package
  70  *        Model I/O will have role one of below
  71  *        - Package input/output
  72  *        - Edge's start/finish point between model
  73  */
  74 struct ModelEdges
  75 {
  76   std::vector<ir::IODesc> pkg_inputs;
  77   std::vector<ir::IODesc> pkg_outputs;
  78   ModelEdgeSet edges;
  79 };
  80
  81 class NNPkg
  82 {
  83 public:
  84   NNPkg() = default;
  85   NNPkg(const NNPkg &obj) = default;
  86   NNPkg(NNPkg &&) = default;
  87   NNPkg &operator=(const NNPkg &) = default;
  88   NNPkg &operator=(NNPkg &&) = default;
  89   ~NNPkg() = default;
  90
  91   NNPkg(std::shared_ptr<Model> model) { _models[ModelIndex{0}] = model; }
  92   std::shared_ptr<Model> primary_model() { return _models.at(onert::ir::ModelIndex{0}); }
  93
  94   /**
  95    * @brief Put model at index
  96    *
  97    * @param[in] model Model to be pushed
  98    * @param[in] index Index where Model is to be pushed
  99    */
 100   void push(ModelIndex index, const std::shared_ptr<Model> &model) { _models[index] = model; }
 101
 102   /**
 103    * @brief Get the count of model
 104    *
 105    * @return the count of models
 106    */
 107   size_t model_count() const { return _models.size(); }
 108
 109   /**
 110    * @brief Get model at index
 111    *
 112    * @param[in] index Index of the model to be returned
 113    * @return Model at index
 114    */
 115   const std::shared_ptr<Model> &model(const ModelIndex &index) const { return _models.at(index); }
 116   /**
 117    * @brief Get model at index
 118    *
 119    * @param[in] index Index of the model to be returned
 120    * @return Model at index
 121    */
 122   std::shared_ptr<Model> &model(const ModelIndex &index) { return _models.at(index); }
 123
 124   /**
 125    * @brief Get pkg_input at index
 126    *
 127    * @param[in] index Index of pkg_input to be returned
 128    * @return IODesc at index
 129    */
 130   const IODesc &input(uint32_t index) const { return _edges.pkg_inputs[index]; }
 131   /**
 132    * @brief Get pkg_input at index
 133    *
 134    * @param[in] index Index of pkg_input to be returned
 135    * @return IODesc at index
 136    */
 137   IODesc &input(uint32_t index) { return _edges.pkg_inputs[index]; }
 138   /**
 139    * @brief Add input at the end
 140    *
 141    * @param[in] input Input IODesc to be pushed
 142    */
 143   void addInput(const IODesc &input) { _edges.pkg_inputs.push_back(input); }
 144
 145   /**
 146    * @brief Get pkg_output at index
 147    *
 148    * @param[in] index Index of pkg_output to be returned
 149    * @return IODesc at index
 150    */
 151   const IODesc &output(uint32_t index) const { return _edges.pkg_outputs[index]; }
 152   /**
 153    * @brief Get pkg_output at index
 154    *
 155    * @param[in] index Index of pkg_output to be returned
 156    * @return IODesc at index
 157    */
 158   IODesc &output(uint32_t index) { return _edges.pkg_outputs[index]; }
 159   /**
 160    * @brief Add output at the end
 161    *
 162    * @param[in] output Output IODesc to be pushed
 163    */
 164   void addOutput(const IODesc &output) { _edges.pkg_outputs.push_back(output); }
 165
 166   /**
 167    * @brief Add edge between models at the end
 168    *
 169    * @param[in] from from IODesc
 170    * @param[in] to   to IODesc
 171    */
 172   void addEdge(const IODesc &from, const IODesc &to)
 173   {
 174     std::cout << from << " -> " << to << std::endl;
 175     _edges.edges.insert(ModelEdge{from, to});
 176   }
 177   /**
 178    * @brief   Get model edge set
 179    * @return  Edge set reference
 180    */
 181   const ModelEdges &model_edges() { return _edges; }
 182
 183   // TODO: Add iterate() or getter for edges
 184
 185 private:
 186   std::unordered_map<ModelIndex, std::shared_ptr<Model>> _models;
 187   ModelEdges _edges;
 188 };
 189
 190 } // namespace ir
 191 } // namespace onert
 192
 193 #endif // __ONERT_IR_NNPKG_H__