ONE - On-device Neural Engine
Loading...
Searching...
No Matches
Gather.cpp
Go to the documentation of this file.
1/*
2 * Copyright (c) 2022 Samsung Electronics Co., Ltd. All Rights Reserved
3 * Copyright 2021 The TensorFlow Authors. All Rights Reserved.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18#include "Builders.h"
19#include "kernels/Utils.h"
20#include "TISOKernel.h"
21
22#include <cassert>
23
24namespace luci_interpreter
25{
26namespace
27{
28
29template <typename InputT, typename CoordsT = int32_t>
30void gather(const circle::GatherOptions *options, kernels::TISOKernel *kernel)
31{
32 kernels::TISOData tiso_data = kernel->readData();
33
34 const InputT *input_data = kernels::getTensorData<InputT>(tiso_data.input1_data);
35 const CoordsT *coords_data = kernels::getTensorData<CoordsT>(tiso_data.input2_data);
36 InputT *output_data = kernels::getTensorData<InputT>(tiso_data.output_data);
37
38 const circle::Tensor *input = kernel->input1();
39 const circle::Tensor *coords = kernel->input2();
40
41 const int input_dims_size = Tensor::num_dims(input);
42 int axis = options->axis();
43 if (axis < 0)
44 {
45 axis += input_dims_size;
46 }
47
48 int batch_dims = options->batch_dims();
49 // batch_dims should be in range: [-rank(coords), rank(coords)].
50 // Negative batch_dims is added with rank of coords.
51 const int coords_dims_size = Tensor::num_dims(coords);
52 if (batch_dims < 0)
53 {
54 batch_dims += coords_dims_size;
55 }
56
57 const int axis_size = Tensor::dim(input, axis);
58
59 int batch_size = 1;
60 for (int i = 0; i < batch_dims; ++i)
61 {
62 batch_size *= Tensor::dim(input, i);
63 }
64 int outer_size = 1;
65 for (int i = batch_dims; i < axis; ++i)
66 {
67 outer_size *= Tensor::dim(input, i);
68 }
69 int inner_size = 1;
70 for (int i = axis + 1; i < input_dims_size; ++i)
71 {
72 inner_size *= Tensor::dim(input, i);
73 }
74 int coord_size = 1;
75 for (int i = batch_dims; i < coords_dims_size; ++i)
76 {
77 coord_size *= Tensor::dim(coords, i);
78 }
79
80 for (int batch = 0; batch < batch_size; ++batch)
81 {
82 for (int outer = 0; outer < outer_size; ++outer)
83 {
84 for (int coord = 0; coord < coord_size; ++coord)
85 {
86 auto x = coords_data[coord];
87 std::memcpy(
88 output_data + (((batch * outer_size) + outer) * coord_size + coord) * inner_size,
89 input_data +
90 (((batch * outer_size) + outer) * axis_size + coords_data[batch * coord_size + coord]) *
91 inner_size,
92 sizeof(InputT) * inner_size);
93 }
94 }
95 }
96}
97
98} // namespace
99
100void configure_kernel_CircleGather(const circle::Operator *cur_op, BaseRuntimeGraph *runtime_graph)
101{
102 kernels::TISOKernel kernel(cur_op, runtime_graph);
103
104 const auto *options = cur_op->builtin_options_as_GatherOptions();
105
106 LUCI_INTERPRETER_CHECK(Tensor::element_type(kernel.input2()) == DataType::S32);
107 LUCI_INTERPRETER_CHECK(Tensor::element_type(kernel.input1()) == DataType::FLOAT32 or
108 Tensor::element_type(kernel.input1()) == DataType::S8 or
109 Tensor::element_type(kernel.input1()) == DataType::S32);
110
111 int32_t axis = options->axis();
112 int32_t num_dims = Tensor::num_dims(kernel.input1());
113 if (axis < 0)
114 {
115 axis += num_dims;
116 }
117
118 LUCI_INTERPRETER_CHECK(axis >= 0 and axis < num_dims);
119
120 int32_t batch_dims = options->batch_dims();
121 int32_t coords_num_dims = Tensor::num_dims(kernel.input2());
122 // batch_dims should be in range: [-rank(coords), rank(coords)].
123 // Negative batch_dims is added with rank of coords.
124 if (batch_dims < 0)
125 {
126 batch_dims += coords_num_dims;
127 }
128 LUCI_INTERPRETER_CHECK(batch_dims <= axis);
129 LUCI_INTERPRETER_CHECK(batch_dims >= 0 and batch_dims < num_dims);
130 LUCI_INTERPRETER_CHECK(batch_dims <= coords_num_dims);
131 for (int i = 0; i < batch_dims; ++i)
132 {
133 LUCI_INTERPRETER_CHECK(Tensor::dim(kernel.input1(), i) == Tensor::dim(kernel.input2(), i));
134 }
135}
136
137void execute_kernel_CircleGather(const circle::Operator *cur_op, BaseRuntimeGraph *runtime_graph)
138{
139 kernels::TISOKernel kernel(cur_op, runtime_graph);
140
141 const auto *options = cur_op->builtin_options_as_GatherOptions();
142
143 switch (Tensor::element_type(kernel.input1()))
144 {
145#ifndef DIS_FLOAT
146 case DataType::FLOAT32:
147 return gather<float, int32_t>(options, &kernel);
148#endif // DIS_FLOAT
149#ifndef DIS_QUANT
150 case DataType::S8:
151 return gather<int8_t, int32_t>(options, &kernel);
152#endif // DIS_QUANT
153 case DataType::S32:
154 return gather<int32_t, int32_t>(options, &kernel);
155 default:
156 assert(false && "Unsupported type");
157 }
158}
159
160} // namespace luci_interpreter
Array< CornerBox > coords
const circle::Tensor * input2() const
Definition TISOKernel.h:61
const circle::Tensor * input1() const
Definition TISOKernel.h:60
#define LUCI_INTERPRETER_CHECK(cond)
Definition Utils.h:36
list input_data
Definition infer.py:29
void configure_kernel_CircleGather(const circle::Operator *cur_op, BaseRuntimeGraph *runtime_graph)
Definition Gather.cpp:100
void execute_kernel_CircleGather(const circle::Operator *cur_op, BaseRuntimeGraph *runtime_graph)
Definition Gather.cpp:137
const loco::Dimension & dim(uint32_t axis) const
Definition Tensor.h:44