2 * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
3 * Copyright 2020 The TensorFlow Authors. All Rights Reserved.
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
9 * http://www.apache.org/licenses/LICENSE-2.0
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.
18 #ifndef LUCI_INTERPRETER_PAL_STRIDED_SLICE_H
19 #define LUCI_INTERPRETER_PAL_STRIDED_SLICE_H
23 namespace luci_interpreter_pal
28 // Use until std::clamp() is available from C++17.
29 inline int clamp(const int v, const int lo, const int hi)
38 inline bool loopCondition(int index, int stop, int stride)
40 // True when we have reached the end of an axis and should loop.
41 return stride > 0 ? index >= stop : index <= stop;
44 // Return the "real" index for the end of iteration along that axis. This is an
45 // "end" in the traditional C sense, in that it points to one past the last
46 // element. ie. So if you were iterating through all elements of a 1D array of
47 // size 4, this function would return 4 as the stop, because it is one past the
48 // "real" indices of 0, 1, 2 & 3.
49 inline int stopForAxis(const StridedSliceParams ¶ms,
50 const luci_interpreter::RuntimeShape &input_shape, int axis,
53 const auto end_mask = params.end_mask;
54 const auto shrink_axis_mask = params.shrink_axis_mask;
55 const auto *stop_indices = params.stop_indices;
56 const auto *strides = params.strides;
57 const int axis_size = input_shape.dims(axis);
63 // Begin with the specified index
64 const bool shrink_axis = shrink_axis_mask & (1 << axis);
65 int stop = stop_indices[axis];
67 // When shrinking an axis, the end position does not matter (and can be
68 // incorrect when negative indexing is used, see Issue #19260). Always use
69 // start_for_axis + 1 to generate a length 1 slice, since start_for_axis has
70 // already been adjusted for negative indices.
73 return start_for_axis + 1;
77 if (end_mask & (1 << axis))
79 if (strides[axis] > 0)
81 // Forward iteration - use the last element. These values will get
83 stop = std::numeric_limits<int>::max();
87 // Backward iteration - use the first element.
88 stop = std::numeric_limits<int>::lowest();
92 // Handle negative indices
99 // Because the end index points one past the last element, we need slightly
100 // different clamping ranges depending on the direction.
101 if (strides[axis] > 0)
104 stop = clamp(stop, 0, axis_size);
108 // Backward iteration
109 stop = clamp(stop, -1, axis_size - 1);
115 // Return the index for the first element along that axis. This index will be a
116 // positive integer between [0, axis_size] (or [-1, axis_size -1] if stride < 0)
117 // that can be used to index directly into the data.
118 inline int startForAxis(const StridedSliceParams ¶ms,
119 const luci_interpreter::RuntimeShape &input_shape, int axis)
121 const auto begin_mask = params.begin_mask;
122 const auto *start_indices = params.start_indices;
123 const auto *strides = params.strides;
124 const int axis_size = input_shape.dims(axis);
129 // Begin with the specified index.
130 int start = start_indices[axis];
132 // begin_mask override
133 if (begin_mask & 1 << axis)
135 if (strides[axis] > 0)
137 // Forward iteration - use the first element. These values will get
138 // clamped below (Note: We could have set them to 0 and axis_size-1, but
139 // use lowest() and max() to maintain symmetry with StopForAxis())
140 start = std::numeric_limits<int>::lowest();
144 // Backward iteration - use the last element.
145 start = std::numeric_limits<int>::max();
149 // Handle negative indices
156 if (strides[axis] > 0)
159 start = clamp(start, 0, axis_size);
163 // Backward iteration
164 start = clamp(start, -1, axis_size - 1);
170 inline void stridedSlicePadIndices(StridedSliceParams *p, int dim_count)
172 const int pad_count = dim_count - p->start_indices_count;
174 // Pad indices at start, so move arrays by pad_count.
175 for (int i = p->start_indices_count - 1; i >= 0; --i)
177 p->strides[i + pad_count] = p->strides[i];
178 p->start_indices[i + pad_count] = p->start_indices[i];
179 p->stop_indices[i + pad_count] = p->stop_indices[i];
181 for (int i = 0; i < pad_count; ++i)
183 p->start_indices[i] = 0;
184 p->stop_indices[i] = 1;
188 // Pad masks with 0s or 1s as required.
189 p->shrink_axis_mask <<= pad_count;
190 p->ellipsis_mask <<= pad_count;
191 p->new_axis_mask <<= pad_count;
192 p->begin_mask <<= pad_count;
193 p->end_mask <<= pad_count;
194 p->begin_mask |= (1 << pad_count) - 1;
195 p->end_mask |= (1 << pad_count) - 1;
197 p->start_indices_count = dim_count;
198 p->stop_indices_count = dim_count;
199 p->strides_count = dim_count;
204 template <typename T>
205 inline void StridedSlice(StridedSliceParams &op_params,
206 const luci_interpreter::RuntimeShape &unextended_input_shape,
207 const T *input_data, T *output_data)
209 const luci_interpreter::RuntimeShape input_shape =
210 luci_interpreter::RuntimeShape::extendedShape(5, unextended_input_shape);
212 // Reverse and pad to 5 dimensions because that is what the runtime code
213 // requires (ie. all shapes must be 5D and are given backwards).
214 stridedSlicePadIndices(&op_params, 5);
216 const int start_0 = startForAxis(op_params, input_shape, 0);
217 const int stop_0 = stopForAxis(op_params, input_shape, 0, start_0);
218 const int start_1 = startForAxis(op_params, input_shape, 1);
219 const int stop_1 = stopForAxis(op_params, input_shape, 1, start_1);
220 const int start_2 = startForAxis(op_params, input_shape, 2);
221 const int stop_2 = stopForAxis(op_params, input_shape, 2, start_2);
222 const int start_3 = startForAxis(op_params, input_shape, 3);
223 const int stop_3 = stopForAxis(op_params, input_shape, 3, start_3);
224 const int start_4 = startForAxis(op_params, input_shape, 4);
225 const int stop_4 = stopForAxis(op_params, input_shape, 4, start_4);
227 for (int offset_0 = start_0 * input_shape.dims(1), end_0 = stop_0 * input_shape.dims(1),
228 step_0 = op_params.strides[0] * input_shape.dims(1);
229 !loopCondition(offset_0, end_0, op_params.strides[0]); offset_0 += step_0)
231 for (int offset_1 = (offset_0 + start_1) * input_shape.dims(2),
232 end_1 = (offset_0 + stop_1) * input_shape.dims(2),
233 step_1 = op_params.strides[1] * input_shape.dims(2);
234 !loopCondition(offset_1, end_1, op_params.strides[1]); offset_1 += step_1)
236 for (int offset_2 = (offset_1 + start_2) * input_shape.dims(3),
237 end_2 = (offset_1 + stop_2) * input_shape.dims(3),
238 step_2 = op_params.strides[2] * input_shape.dims(3);
239 !loopCondition(offset_2, end_2, op_params.strides[2]); offset_2 += step_2)
241 for (int offset_3 = (offset_2 + start_3) * input_shape.dims(4),
242 end_3 = (offset_2 + stop_3) * input_shape.dims(4),
243 step_3 = op_params.strides[3] * input_shape.dims(4);
244 !loopCondition(offset_3, end_3, op_params.strides[3]); offset_3 += step_3)
246 for (int offset_4 = offset_3 + start_4, end_4 = offset_3 + stop_4;
247 !loopCondition(offset_4, end_4, op_params.strides[4]);
248 offset_4 += op_params.strides[4])
250 *output_data++ = input_data[offset_4];
258 } // namespace luci_interpreter_pal
260 #endif // LUCI_INTERPRETER_PAL_STRIDED_SLICE_H