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 "exec/IPermuteFunction.h"
  21 #include "exec/IExecutor.h"
  22 #include "../ExternalContext.h"
  23 #include "ruy/thread_pool.h" // from @ruy
  24
  25 namespace onert
  26 {
  27 namespace backend
  28 {
  29 namespace controlflow
  30 {
  31 namespace kernel
  32 {
  33
  34 class PermuteLayer : public onert::exec::IPermuteFunction
  35 {
  36 public:
  37   PermuteLayer(const std::vector<ITensor *> &src_tensors, const std::vector<ITensor *> &dst_tensors,
  38                const std::shared_ptr<ExternalContext> &external_context);
  39
  40   void optimize() override;
  41
  42   void run() override;
  43
  44 private:
  45   std::shared_ptr<ExternalContext> _external_context;
  46
  47 private:
  48   void appendPermuteTasks(const ITensor *src_tensor, ITensor *dst_tensor,
  49                           const ir::Shape &loop_shape, size_t size);
  50
  51   void runPermuteTasks(backend::ITensor *src, uint8_t *dst_buffer);
  52
  53   struct PermuteWorkerTask : ruy::Task
  54   {
  55     using Strides = ir::Coordinates;
  56
  57     PermuteWorkerTask(const ITensor &src_tensor, ITensor &dst_tensor,
  58                       const ir::Coordinates &start_coords, const ir::Shape &loop_shape, size_t size)
  59         : _src_buffer{src_tensor.buffer()}, _dst_buffer{dst_tensor.buffer()},
  60           _src_start_offset{src_tensor.calcOffset(start_coords)},
  61           _dst_start_offset{dst_tensor.calcOffset(start_coords)}, _src_strides{}, _dst_strides{},
  62           _loop_shape{loop_shape}, _size{size}, _src_layout{src_tensor.layout()},
  63           _dst_layout{dst_tensor.layout()}, _is_permutation{true}
  64     {
  65       // Set strides
  66       setStrides(src_tensor, &_src_strides);
  67       setStrides(dst_tensor, &_dst_strides);
  68
  69       _is_permutation = (_src_layout != _dst_layout && loop_shape.rank() == 4);
  70     }
  71     // Constructor for a copy
  72     PermuteWorkerTask(const uint8_t *src_buffer, uint8_t *dst_buffer, uint32_t src_start_offset,
  73                       uint32_t dst_start_offset, size_t size)
  74         : _src_buffer{src_buffer}, _dst_buffer{dst_buffer}, _src_start_offset{src_start_offset},
  75           _dst_start_offset{dst_start_offset}, _src_strides{0}, _dst_strides{0}, _loop_shape{1},
  76           _size{size}, _src_layout{}, _dst_layout{}, _is_permutation{false}
  77     {
  78       // DO NOTHING
  79     }
  80     void setBuffers(const uint8_t *src_buffer, uint8_t *dst_buffer)
  81     {
  82       _src_buffer = src_buffer;
  83       _dst_buffer = dst_buffer;
  84     }
  85     void Run() override
  86     {
  87       ShapeLoop(_loop_shape, [&](const onert::ir::Coordinates &coords) {
  88         size_t src_offset = _src_start_offset;
  89         size_t dst_offset = _dst_start_offset;
  90         assert(static_cast<size_t>(_loop_shape.rank()) == coords.size());
  91         ir::Coordinates dst_coords = coords;
  92         if (_is_permutation)
  93         {
  94           dst_coords = ir::convertCoordinates(coords, _src_layout, _dst_layout);
  95         }
  96         for (auto i = 0; i < _loop_shape.rank(); ++i)
  97         {
  98           assert(coords[i] >= 0 && dst_coords[i] >= 0);
  99           src_offset += coords[i] * _src_strides[i];
 100           dst_offset += dst_coords[i] * _dst_strides[i];
 101         }
 102         memcpy(_dst_buffer + dst_offset, _src_buffer + src_offset, _size);
 103       });
 104     }
 105
 106   private:
 107     void setStrides(const ITensor &tensor, Strides *strides)
 108     {
 109       const size_t rank = tensor.num_dimensions();
 110       for (size_t i = 0; i < rank; ++i)
 111       {
 112         ir::Coordinates no_step(rank), one_step(rank);
 113         one_step.set(i, 1);
 114         if (tensor.dimension(i) > 1)
 115         {
 116           strides->set(i, tensor.calcOffset(one_step) - tensor.calcOffset(no_step));
 117         }
 118         else
 119         {
 120           // If dimension value is 0 or 1, the stride of the dimension will be not used
 121           // Do not call calcOffset() with coordinate value that is greater than dimension value
 122           strides->set(i, 0);
 123         }
 124         assert((*strides)[i] >= 0);
 125       }
 126     }
 127
 128   private:
 129     const uint8_t *_src_buffer;
 130     uint8_t *_dst_buffer;
 131     size_t _src_start_offset;
 132     size_t _dst_start_offset;
 133     Strides _src_strides;
 134     Strides _dst_strides;
 135     const ir::Shape _loop_shape;
 136     const size_t _size;
 137     const ir::Layout _src_layout;
 138     const ir::Layout _dst_layout;
 139     bool _is_permutation;
 140   };
 141   std::unordered_map<const ITensor *, std::vector<PermuteWorkerTask>> _tasks_map;
 142 };
 143
 144 } // namespace kernel
 145 } // namespace controlflow
 146 } // namespace backend
 147 } // namespace onert
 148
 149 #endif // __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__