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.
17 #ifndef __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__
18 #define __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__
20 #include "exec/IPermuteFunction.h"
21 #include "exec/IExecutor.h"
22 #include "../ExternalContext.h"
23 #include "ruy/thread_pool.h" // from @ruy
34 class PermuteLayer : public onert::exec::IPermuteFunction
37 PermuteLayer(const std::vector<ITensor *> &src_tensors, const std::vector<ITensor *> &dst_tensors,
38 const std::shared_ptr<ExternalContext> &external_context);
40 void optimize() override;
45 std::shared_ptr<ExternalContext> _external_context;
48 void appendPermuteTasks(const ITensor *src_tensor, ITensor *dst_tensor,
49 const ir::Shape &loop_shape, size_t size);
51 void runPermuteTasks(backend::ITensor *src, uint8_t *dst_buffer);
53 struct PermuteWorkerTask : ruy::Task
55 using Strides = ir::Coordinates;
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}
66 setStrides(src_tensor, &_src_strides);
67 setStrides(dst_tensor, &_dst_strides);
69 _is_permutation = (_src_layout != _dst_layout && loop_shape.rank() == 4);
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}
80 void setBuffers(const uint8_t *src_buffer, uint8_t *dst_buffer)
82 _src_buffer = src_buffer;
83 _dst_buffer = dst_buffer;
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;
94 dst_coords = ir::convertCoordinates(coords, _src_layout, _dst_layout);
96 for (auto i = 0; i < _loop_shape.rank(); ++i)
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];
102 memcpy(_dst_buffer + dst_offset, _src_buffer + src_offset, _size);
107 void setStrides(const ITensor &tensor, Strides *strides)
109 const size_t rank = tensor.num_dimensions();
110 for (size_t i = 0; i < rank; ++i)
112 ir::Coordinates no_step(rank), one_step(rank);
114 if (tensor.dimension(i) > 1)
116 strides->set(i, tensor.calcOffset(one_step) - tensor.calcOffset(no_step));
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
124 assert((*strides)[i] >= 0);
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;
137 const ir::Layout _src_layout;
138 const ir::Layout _dst_layout;
139 bool _is_permutation;
141 std::unordered_map<const ITensor *, std::vector<PermuteWorkerTask>> _tasks_map;
144 } // namespace kernel
145 } // namespace controlflow
146 } // namespace backend
149 #endif // __ONERT_BACKEND_CONTROLFLOW_KERNEL_PERMUTELAYER_H__