runtime/onert/core/src/backend/controlflow/kernel/IfLayer.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 "IfLayer.h"
  18
  19 #include <backend/ITensor.h>
  20 #include "exec/ExecutorBase.h"
  21 #include <misc/polymorphic_downcast.h>
  22 #include "PermuteLayer.h"
  23
  24 namespace onert
  25 {
  26 namespace backend
  27 {
  28 namespace controlflow
  29 {
  30 namespace kernel
  31 {
  32
  33 IfLayer::IfLayer(backend::ITensor *cond_tensor, const std::vector<backend::ITensor *> input_tensors,
  34                  const std::vector<backend::ITensor *> output_tensors,
  35                  const ir::OperandIndexSequence &output_indices, const ir::Graph &graph,
  36                  const ir::SubgraphIndex &then_subg_index, const ir::SubgraphIndex &else_subg_index,
  37                  exec::ExecutorMap *executor_map,
  38                  const std::shared_ptr<ExternalContext> &external_context)
  39     : _cond_tensor{cond_tensor}, _input_tensors{input_tensors}, _output_tensors{output_tensors},
  40       _output_indices{output_indices}, _graph{graph}, _then_subg_index{then_subg_index},
  41       _else_subg_index{else_subg_index}, _executor_map{executor_map},
  42       _external_context{external_context}
  43 {
  44   // At this point, executor_map may not have executors of then subg and else subg
  45 }
  46
  47 void IfLayer::run()
  48 {
  49   // Check condition
  50   // // If true
  51   // // // Copy _input_tensors -> then subg's inputs
  52   // // // Run then subg
  53   // // // Copy outputs of then subg -> _output_tensors
  54   // // Else
  55   // // // Copy _input_tensors -> else subg's inputs if false
  56   // // // Run else subg
  57   // // // Copy outputs of else subg -> _output_tensors
  58   auto getResultCond = [](backend::ITensor *tensor) -> bool {
  59     bool ret = false;
  60     tensor->access([&](ITensor &tensor) { ret = *reinterpret_cast<bool *>(tensor.buffer()); });
  61     return ret;
  62   };
  63
  64   exec::ExecutorBase *subg_exec = nullptr;
  65   bool cond_result = getResultCond(_cond_tensor);
  66   if (cond_result)
  67   {
  68     VERBOSE(If) << "Call to $" << _then_subg_index << " (then)" << std::endl;
  69     subg_exec = nnfw::misc::polymorphic_downcast<exec::ExecutorBase *>(
  70         _executor_map->at(_then_subg_index).get());
  71   }
  72   else
  73   {
  74     VERBOSE(If) << "Call to $" << _else_subg_index << " (else)" << std::endl;
  75     subg_exec = nnfw::misc::polymorphic_downcast<exec::ExecutorBase *>(
  76         _executor_map->at(_else_subg_index).get());
  77   }
  78
  79   const auto &subg_graph = subg_exec->graph();
  80
  81   std::vector<backend::ITensor *> src_tensors;
  82   std::vector<backend::ITensor *> dst_tensors;
  83   // Add tensors used in subgraph or contained in outputs of subgraph
  84   assert(subg_graph.getInputs().size() == _input_tensors.size());
  85   assert(subg_graph.getInputs().size() == subg_exec->getInputTensors().size());
  86   for (uint32_t i = 0; i < subg_graph.getInputs().size(); ++i)
  87   {
  88     const auto &subg_input_index = subg_graph.getInputs().at(i);
  89     const auto &subg_input = subg_graph.operands().at(subg_input_index);
  90     if (subg_input.getUses().size() > 0 || subg_graph.getOutputs().contains(subg_input_index))
  91     {
  92       src_tensors.emplace_back(_input_tensors.at(i));
  93       dst_tensors.emplace_back(subg_exec->getInputTensors().at(i));
  94     }
  95   }
  96   const auto permute_op_input_to_subg_input =
  97       std::make_shared<PermuteLayer>(src_tensors, dst_tensors, _external_context);
  98
  99   // Add tensors used as output of operation or contained in outputs of operation
 100   src_tensors.clear();
 101   dst_tensors.clear();
 102   assert(_output_indices.size() == subg_exec->getOutputTensors().size());
 103   assert(_output_indices.size() == _output_tensors.size());
 104   for (uint32_t i = 0; i < _output_indices.size(); ++i)
 105   {
 106     const auto &output_index = _output_indices.at(i);
 107     const auto &output = _graph.operands().at(output_index);
 108     if (output.getUses().size() > 0 || _graph.getOutputs().contains(output_index))
 109     {
 110       src_tensors.emplace_back(subg_exec->getOutputTensors().at(i));
 111       dst_tensors.emplace_back(_output_tensors.at(i));
 112     }
 113   }
 114   const auto permute_subg_output_to_op_output =
 115       std::make_shared<PermuteLayer>(src_tensors, dst_tensors, _external_context);
 116
 117   // Remove copying of unused tensor
 118   permute_op_input_to_subg_input->prepare();
 119   permute_subg_output_to_op_output->prepare();
 120
 121   // Copy & run
 122   subg_exec->execute(_input_tensors, permute_op_input_to_subg_input);
 123   permute_subg_output_to_op_output->run();
 124   VERBOSE(If) << "Return from $" << (cond_result ? _then_subg_index : _else_subg_index)
 125               << std::endl;
 126 }
 127
 128 } // namespace kernel
 129 } // namespace controlflow
 130 } // namespace backend
 131 } // namespace onert