runtime/onert/core/src/backend/controlflow/KernelGenerator.cc

   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 #include "KernelGenerator.h"
  18
  19 #include <backend/BackendContext.h>
  20 #include <util/Utils.h>
  21 #include "kernel/IfLayer.h"
  22 #include "kernel/WhileLayer.h"
  23 #include "kernel/PermuteLayer.h"
  24 #include "exec/ExecutorBase.h"
  25 #include "exec/FunctionSequence.h"
  26
  27 namespace onert
  28 {
  29 namespace backend
  30 {
  31 namespace controlflow
  32 {
  33
  34 KernelGenerator::KernelGenerator(const ir::Graph &graph,
  35                                  const std::shared_ptr<TensorBuilder> &tensor_builder)
  36     : _graph{graph}, _tensor_builder{tensor_builder}, _tensor_builder_set{}, _executor_map{nullptr}
  37 {
  38   UNUSED_RELEASE(_graph);
  39   UNUSED_RELEASE(_tensor_builder_set);
  40   UNUSED_RELEASE(_executor_map);
  41 }
  42
  43 void KernelGenerator::visit(const ir::OpSequence &op_seq)
  44 {
  45   assert(!_return_fn_seq);
  46   assert(_tensor_builder->dynamicTensorManager());
  47   assert(_tensor_builder->tensorRegistry());
  48
  49   auto dyn_tensor_manager = _tensor_builder->dynamicTensorManager();
  50   auto dyn_shape_inferer = std::make_unique<exec::DynamicShapeInferer>(
  51       _graph.operands(), dyn_tensor_manager, _tensor_builder->tensorRegistry());
  52
  53   _return_fn_seq = std::make_unique<exec::FunctionSequence>();
  54
  55   // Prepare to handle dynamic tensors later
  56   auto dyn_ctx = std::make_shared<exec::FunctionSequence::DynamicTensorCtx>();
  57   {
  58     dyn_ctx->op_seq = &op_seq;
  59     dyn_ctx->operations = &_graph.operations();
  60     dyn_ctx->dynamic_shape_inferer = std::move(dyn_shape_inferer);
  61     dyn_ctx->tensor_registry = _tensor_builder->tensorRegistry();
  62     dyn_ctx->dynamic_tensor_manager = _tensor_builder->dynamicTensorManager();
  63
  64     _return_fn_seq->dynamic_tensor_ctx(dyn_ctx);
  65   }
  66   _return_fn_seq->enableDynamicShapeInferer(true);
  67
  68   for (const auto &op_idx : op_seq.operations())
  69   {
  70     const auto &node = _graph.operations().at(op_idx);
  71     node.accept(*this);
  72     _return_fn_seq->append(releaseFunction());
  73   }
  74 }
  75
  76 void KernelGenerator::visit(const ir::operation::If &node)
  77 {
  78   const auto then_subg_index = node.param().then_subg_index;
  79   const auto else_subg_index = node.param().else_subg_index;
  80
  81   std::vector<std::shared_ptr<backend::ITensor>> input_tensors;
  82   for (const auto input_index : node.getInputs())
  83   {
  84     auto input_tensor = getTensor(input_index);
  85
  86     input_tensors.emplace_back(input_tensor);
  87   }
  88
  89   std::vector<std::shared_ptr<backend::ITensor>> output_tensors;
  90   exec::DynAllocInfoMap outputs_dyn_alloc_info;
  91   for (const auto output_index : node.getOutputs())
  92   {
  93     auto output_tensor = getTensor(output_index);
  94
  95     output_tensors.emplace_back(output_tensor);
  96     const auto output_tensor_builder = getTensorBuilder(output_index);
  97     if (output_tensor_builder->supportDynamicTensor())
  98     {
  99       auto output_dyn_manager = output_tensor_builder->dynamicTensorManager();
 100       outputs_dyn_alloc_info[output_tensor] = exec::DynAllocInfo{output_index, output_dyn_manager};
 101     }
 102   }
 103
 104   // IfLayer just set ExecutorMap instead of then and else executor to avoid complexity of
 105   // creating executor recusively
 106   const auto cond_tensor = input_tensors.front();
 107   input_tensors.erase(input_tensors.begin());
 108   auto fn = std::make_unique<::onert::backend::controlflow::kernel::IfLayer>(
 109       cond_tensor, input_tensors, output_tensors, node.getOutputs(), _graph, outputs_dyn_alloc_info,
 110       then_subg_index, else_subg_index, _executor_map);
 111
 112   _return_fn = std::move(fn);
 113 }
 114
 115 void KernelGenerator::visit(const ir::operation::Permute &node)
 116 {
 117   const auto output_index{node.getOutputs().at(0)};
 118   const auto input_index{node.getInputs().at(0)};
 119
 120   // Add PermuteLayer
 121   std::vector<std::shared_ptr<ITensor>> output_tensors{getTensor(output_index)};
 122   std::vector<std::shared_ptr<ITensor>> input_tensors{getTensor(input_index)};
 123   std::unordered_map<std::shared_ptr<ITensor>, exec::DynAllocInfo> outputs_dyn_alloc_info;
 124   const auto output_tensor_builder = getTensorBuilder(output_index);
 125   VERBOSE(PERMUTE_FIND_TB) << output_index << " -> " << output_tensor_builder.get() << std::endl;
 126   assert(output_tensor_builder != nullptr);
 127   if (output_tensor_builder->supportDynamicTensor())
 128   {
 129     outputs_dyn_alloc_info[output_tensors.at(0)] =
 130         exec::DynAllocInfo{output_index, output_tensor_builder->dynamicTensorManager()};
 131   }
 132
 133   auto fn =
 134       std::make_unique<kernel::PermuteLayer>(input_tensors, output_tensors, outputs_dyn_alloc_info);
 135
 136   _return_fn = std::move(fn);
 137 }
 138
 139 void KernelGenerator::visit(const ir::operation::While &node)
 140 {
 141   const auto cond_subg_index = node.param().cond_subg_index;
 142   const auto body_subg_index = node.param().body_subg_index;
 143
 144   // This op does not support input as a constant, because controlflow backend does not have
 145   // TensorBuilder
 146   std::vector<std::shared_ptr<backend::ITensor>> input_tensors;
 147   for (const auto input_index : node.getInputs())
 148   {
 149     auto input_tensor = getTensor(input_index);
 150
 151     input_tensors.emplace_back(input_tensor);
 152   }
 153
 154   std::vector<std::shared_ptr<backend::ITensor>> output_tensors;
 155   std::unordered_map<std::shared_ptr<ITensor>, exec::DynAllocInfo> outputs_dyn_alloc_info;
 156   for (const auto output_index : node.getOutputs())
 157   {
 158     auto output_tensor = getTensor(output_index);
 159
 160     output_tensors.emplace_back(output_tensor);
 161
 162     const auto output_tensor_builder = getTensorBuilder(output_index);
 163     if (output_tensor_builder->supportDynamicTensor())
 164     {
 165       auto output_dyn_manager = output_tensor_builder->dynamicTensorManager();
 166       outputs_dyn_alloc_info[output_tensor] = exec::DynAllocInfo{output_index, output_dyn_manager};
 167     }
 168   }
 169
 170   // WhileLayer just set ExecutorMap instead of cond and body executor to avoid complexity of
 171   // creating executor recusively
 172   auto fn = std::make_unique<::onert::backend::controlflow::kernel::WhileLayer>(
 173       input_tensors, output_tensors, node.getOutputs(), _graph, outputs_dyn_alloc_info,
 174       cond_subg_index, body_subg_index, _executor_map);
 175
 176   _return_fn = std::move(fn);
 177 }
 178
 179 std::shared_ptr<backend::ITensor> KernelGenerator::getTensor(const ir::OperandIndex &index)
 180 {
 181   std::shared_ptr<backend::ITensor> ret;
 182   for (auto tensor_builder : _tensor_builder_set)
 183   {
 184     auto tensor = tensor_builder->tensorAt(index);
 185     if (tensor)
 186     {
 187       ret = tensor;
 188       break;
 189     }
 190   }
 191   assert(ret != nullptr);
 192   return ret;
 193 }
 194
 195 std::shared_ptr<backend::ITensorBuilder>
 196 KernelGenerator::getTensorBuilder(const ir::OperandIndex &index)
 197 {
 198   std::shared_ptr<backend::ITensorBuilder> ret;
 199   for (auto tensor_builder : _tensor_builder_set)
 200   {
 201     auto reg = tensor_builder->tensorRegistry();
 202     auto tensor = reg ? reg->getNativeITensor(index) : tensor_builder->tensorAt(index);
 203     if (tensor)
 204     {
 205       ret = tensor_builder;
 206       break;
 207     }
 208   }
 209   assert(ret != nullptr);
 210   return ret;
 211 }
 212
 213 } // namespace controlflow
 214 } // namespace backend
 215 } // namespace onert