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