runtime/onert/core/src/backend/controlflow/kernel/PermuteLayer.h

   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 #ifndef __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__
  18 #define __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__
  19
  20 #include "backend/ITensorBuilder.h"
  21 #include "exec/IPermuteFunction.h"
  22 #include "exec/IExecutor.h"
  23
  24 namespace onert
  25 {
  26 namespace backend
  27 {
  28 namespace controlflow
  29 {
  30 namespace kernel
  31 {
  32
  33 class PermuteLayer : public onert::exec::IPermuteFunction
  34 {
  35 public:
  36   PermuteLayer(const std::vector<std::shared_ptr<ITensor>> &src_tensors,
  37                const std::vector<std::shared_ptr<ITensor>> &dst_tensors,
  38                const exec::DynAllocInfoMap &dst_dyn_alloc_info_map)
  39       : _dst_dyn_alloc_info_map{dst_dyn_alloc_info_map}
  40   {
  41     assert(src_tensors.size() == dst_tensors.size());
  42     _src_tensors = src_tensors;
  43     _dst_tensors = dst_tensors;
  44   }
  45
  46   void optimize() override
  47   {
  48     // Remove copying of tensor as nullptr
  49     auto src_it = _src_tensors.begin();
  50     auto dst_it = _dst_tensors.begin();
  51     while (src_it != _src_tensors.end())
  52     {
  53       if ((*src_it == *dst_it) || (*src_it == nullptr || *dst_it == nullptr))
  54       {
  55         src_it = _src_tensors.erase(src_it);
  56         dst_it = _dst_tensors.erase(dst_it);
  57       }
  58       else
  59       {
  60         ++src_it;
  61         ++dst_it;
  62       }
  63     }
  64   }
  65
  66   void run() override;
  67
  68 private:
  69   const exec::DynAllocInfoMap _dst_dyn_alloc_info_map;
  70 };
  71
  72 } // namespace kernel
  73 } // namespace controlflow
  74 } // namespace backend
  75 } // namespace onert
  76
  77 #endif // __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__