2 * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
19 #include <backend/ITensor.h>
20 #include "exec/ExecutorBase.h"
21 #include <misc/polymorphic_downcast.h>
22 #include "PermuteLayer.h"
33 IfLayer::IfLayer(const std::shared_ptr<backend::ITensor> &cond_tensor,
34 const std::vector<std::shared_ptr<backend::ITensor>> input_tensors,
35 const std::vector<std::shared_ptr<backend::ITensor>> output_tensors,
36 const ir::OperandIndexSequence &output_indices, const ir::Graph &graph,
37 const exec::DynAllocInfoMap &outputs_dyn_alloc_info,
38 const ir::SubgraphIndex &then_subg_index, const ir::SubgraphIndex &else_subg_index,
39 exec::ExecutorMap *executor_map)
40 : _cond_tensor{cond_tensor}, _input_tensors{input_tensors}, _output_tensors{output_tensors},
41 _output_indices{output_indices}, _graph{graph},
42 _outputs_dyn_alloc_info{outputs_dyn_alloc_info}, _then_subg_index{then_subg_index},
43 _else_subg_index{else_subg_index}, _executor_map{executor_map}
45 // At this point, executor_map may not have executors of then subg and else subg
52 // // // Copy _input_tensors -> then subg's inputs
53 // // // Run then subg
54 // // // Copy outputs of then subg -> _output_tensors
56 // // // Copy _input_tensors -> else subg's inputs if false
57 // // // Run else subg
58 // // // Copy outputs of else subg -> _output_tensors
59 auto getResultCond = [](backend::ITensor *tensor) -> bool {
61 tensor->access([&](ITensor &tensor) { ret = *reinterpret_cast<bool *>(tensor.buffer()); });
65 exec::ExecutorBase *subg_exec = nullptr;
66 if (getResultCond(_cond_tensor.get()))
68 subg_exec = nnfw::misc::polymorphic_downcast<exec::ExecutorBase *>(
69 _executor_map->at(_then_subg_index).get());
73 subg_exec = nnfw::misc::polymorphic_downcast<exec::ExecutorBase *>(
74 _executor_map->at(_else_subg_index).get());
77 const auto &subg_graph = subg_exec->graph();
79 std::vector<std::shared_ptr<backend::ITensor>> src_tensors;
80 std::vector<std::shared_ptr<backend::ITensor>> dst_tensors;
81 // Add tensors used in subgraph or contained in outputs of subgraph
82 assert(subg_graph.getInputs().size() == _input_tensors.size());
83 assert(subg_graph.getInputs().size() == subg_exec->getInputTensors().size());
84 for (uint32_t i = 0; i < subg_graph.getInputs().size(); ++i)
86 const auto &subg_input_index = subg_graph.getInputs().at(i);
87 const auto &subg_input = subg_graph.operands().at(subg_input_index);
88 if (subg_input.getUses().size() > 0 || subg_graph.getOutputs().contains(subg_input_index))
90 src_tensors.emplace_back(_input_tensors.at(i));
91 dst_tensors.emplace_back(subg_exec->getInputTensors().at(i));
94 const auto &subg_inputs_dyn_alloc_info = subg_exec->getInputsDynamicAllocInfo();
95 const auto permute_op_input_to_subg_input =
96 std::make_shared<PermuteLayer>(src_tensors, dst_tensors, subg_inputs_dyn_alloc_info);
98 // Add tensors used as output of operation or contained in outputs of operation
101 assert(_output_indices.size() == subg_exec->getOutputTensors().size());
102 assert(_output_indices.size() == _output_tensors.size());
103 for (uint32_t i = 0; i < _output_indices.size(); ++i)
105 const auto &output_index = _output_indices.at(i);
106 const auto &output = _graph.operands().at(output_index);
107 if (output.getUses().size() > 0 || _graph.getOutputs().contains(output_index))
109 src_tensors.emplace_back(subg_exec->getOutputTensors().at(i));
110 dst_tensors.emplace_back(_output_tensors.at(i));
113 const auto permute_subg_output_to_op_output =
114 std::make_shared<PermuteLayer>(src_tensors, dst_tensors, _outputs_dyn_alloc_info);
116 // Remove copying of unused tensor
117 permute_op_input_to_subg_input->prepare();
118 permute_subg_output_to_op_output->prepare();
121 subg_exec->execute(_input_tensors, permute_op_input_to_subg_input);
122 permute_subg_output_to_op_output->run();
125 } // namespace kernel
126 } // namespace controlflow
127 } // namespace backend