1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_OPERATIONS_H_
17 #define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_OPERATIONS_H_
18 
19 #include <cstdint>
20 #include <set>
21 #include <string>
22 #include <vector>
23 
24 #include "absl/types/variant.h"
25 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
26 #include "tensorflow/lite/delegates/gpu/common/shape.h"
27 #include "tensorflow/lite/delegates/gpu/common/status.h"
28 #include "tensorflow/lite/delegates/gpu/common/tensor.h"
29 
30 namespace tflite {
31 namespace gpu {
32 
33 // Non exhaustive list of operations.
34 enum class OperationType {
35   UNKNOWN = 0,
36   ABS,
37   ADD,
38   BATCH_TO_SPACE,
39   BATCH_NORMALIZATION,
40   BATCHED_MATMUL,
41   CONCAT,
42   CONSTANT,
43   CONVOLUTION_2D,
44   CONVOLUTION_TRANSPOSED,
45   COPY,
46   COS,
47   DEPTHWISE_CONVOLUTION,
48   DIV,
49   ELU,
50   EQUAL,
51   EXP,
52   FULLY_CONNECTED,
53   GREATER,
54   GREATER_EQUAL,
55   HARD_SWISH,
56   LESS,
57   LESS_EQUAL,
58   LOG,
59   LSTM,
60   MAXIMUM,
61   MAX_UNPOOLING_2D,
62   MEAN,
63   MEAN_STDDEV_NORMALIZATION,
64   MINIMUM,
65   MUL,
66   NEG,
67   NOT_EQUAL,
68   PAD,
69   POOLING_2D,
70   POW,
71   PRELU,
72   // Used to accurately run inference on quantized models.
73   QUANTIZE_AND_DEQUANTIZE,
74   REDUCE_MAXIMUM,
75   REDUCE_MINIMUM,
76   REDUCE_PRODUCT,
77   REDUCE_SUM,
78   RELU,
79   RESHAPE,
80   RESIZE,
81   RSQRT,
82   SIGMOID,
83   SIN,
84   SLICE,
85   SOFTMAX,
86   SPACE_TO_BATCH,
87   SPACE_TO_DEPTH,
88   SPLIT,
89   SQRT,
90   SQUARE,
91   SQUARED_DIFF,
92   SUB,
93   TANH,
94   TRANSPOSE,
95 };
96 
97 std::string ToString(enum OperationType op);
98 
99 OperationType OperationTypeFromString(const std::string& name);
100 
101 typedef absl::variant<absl::monostate, Tensor<HWC, DataType::FLOAT32>,
102                       Tensor<Linear, DataType::FLOAT32>, float>
103     TensorOrScalar;
104 
105 struct Padding2D {
106   Padding2D() = default;
107   Padding2D& operator=(const Padding2D& value);
108   bool operator==(const Padding2D& value);
109   bool operator!=(const Padding2D& value);
110   Padding2D& operator-(const Padding2D& value);
111 
112   // Padding values for every axis (if needed), where 'prepended' defines
113   // padding for the beginning of each axis and 'appended' represents end part
114   // of the corresponding axis.
115   HW prepended = HW(-1, -1);
116   HW appended = HW(-1, -1);
117 };
118 
119 struct Padding3D {
120   Padding3D() = default;
121   Padding3D& operator=(const Padding3D& value);
122   bool operator==(const Padding3D& value);
123   bool operator!=(const Padding3D& value);
124   Padding3D& operator-(const Padding3D& value);
125 
126   // Padding values for every axis (if needed), where 'prepended' defines
127   // padding for the beginning of each axis and 'appended' represents end part
128   // of the corresponding axis.
129   HWD prepended = HWD(0, 0, 0);
130   HWD appended = HWD(0, 0, 0);
131 };
132 
133 struct Crop2D : public Padding2D {};
134 
135 struct SpaceToBatchAttributes {
136   HW block;
137   Padding2D padding;
138 };
139 
140 struct BatchToSpaceAttributes {
141   HW block;
142   Crop2D crop;
143 };
144 
145 enum class PoolingType {
146   UNDEFINED = 0,
147 
148   // average pooling
149   AVERAGE = 1,
150 
151   // max pooling
152   MAX = 2,
153 };
154 
155 struct Pooling2DAttributes {
156   PoolingType type = PoolingType::UNDEFINED;
157   // Strides for every axis.
158   HW strides = HW(-1, -1);
159   HW kernel = HW(-1, -1);
160   Padding2D padding;
161   // NOTE(akulik): technically the number of outputs from Pooling node indicates
162   // whether indices are needed or not, but I decided to keep it inside
163   // attributes to simplify processing.
164   bool output_indices = false;
165 };
166 
167 struct Pooling3DAttributes {
168   PoolingType type = PoolingType::UNDEFINED;
169   // Strides for every axis.
170   HWD strides = HWD(0, 0, 0);
171   HWD kernel = HWD(0, 0, 0);
172   Padding3D padding;
173   // NOTE(akulik): technically the number of outputs from Pooling node indicates
174   // whether indices are needed or not, but I decided to keep it inside
175   // attributes to simplify processing.
176   bool output_indices = false;
177 };
178 
179 struct MaxUnpooling2DAttributes {
180   // Strides for every axis.
181   HW strides = HW(-1, -1);
182   HW kernel = HW(-1, -1);
183   Padding2D padding;
184 };
185 
186 struct MaxUnpooling3DAttributes {
187   // Strides for every axis.
188   HWD strides = HWD(0, 0, 0);
189   HWD kernel = HWD(0, 0, 0);
190   Padding3D padding;
191 };
192 
193 struct MeanAttributes {
194   // The vector of dimensions to calculate mean along.
195   std::set<Axis> dims;
196 };
197 
198 struct ConcatAttributes {
199   // Defines axis by which to concat on.
200   Axis axis = Axis::UNKNOWN;
201 };
202 
203 // @return shape of a tensor after MaxUnpooling2D operation is applied to
204 //         the given input.
205 BHWC CalculateOutputShape(const BHWC& input,
206                           const MaxUnpooling2DAttributes& attr);
207 
208 // @return shape of a tensor after MaxUnpooling3D operation is applied to
209 //         the given input.
210 BHWDC CalculateOutputShape(const BHWDC& input,
211                            const MaxUnpooling3DAttributes& attr);
212 
213 // @return shape of a tensor after Pooling2D operation is applied to the given
214 //         input.
215 BHWC CalculateOutputShape(const BHWC& input, const Pooling2DAttributes& attr);
216 
217 // @return shape of a tensor after Pooling3D operation is applied to the given
218 //         input.
219 BHWDC CalculateOutputShape(const BHWDC& input, const Pooling3DAttributes& attr);
220 
221 // @return shape of a tensor after Concat operation is applied to the given
222 //         input.
223 absl::Status CalculateOutputShape(const std::vector<BHWC>& input,
224                                   const ConcatAttributes& attr,
225                                   BHWC* output_shape);
226 
227 // @return shape of a tensor after Concat operation is applied to the given
228 //         input.
229 absl::Status CalculateOutputShape(const std::vector<BHWDC>& input,
230                                   const ConcatAttributes& attr,
231                                   BHWDC* output_shape);
232 
233 // @return padding for pooling operation to make sure output keep the same shape
234 // as the given input.
235 Padding2D CalculateSamePadding(const BHWC& input,
236                                const Pooling2DAttributes& attr);
237 
238 // @return padding for pooling operation to make sure output keep the same shape
239 // as the given input.
240 Padding3D CalculateSamePadding(const BHWDC& input,
241                                const Pooling3DAttributes& attr);
242 
243 // @return padding for max unpooling operation to make sure output keep the same
244 // shape as the given input.
245 Padding2D CalculateSamePadding(const BHWC& input,
246                                const MaxUnpooling2DAttributes& attr);
247 
248 // @return padding for max unpooling operation to make sure output keep the same
249 // shape as the given input.
250 Padding3D CalculateSamePadding(const BHWDC& input,
251                                const MaxUnpooling3DAttributes& attr);
252 
253 struct Convolution2DAttributes {
254   HW strides = HW(1, 1);    // Along each axis.
255   HW dilations = HW(1, 1);  // Along each axis.
256   Padding2D padding;
257 
258   Tensor<OHWI, DataType::FLOAT32> weights;
259   Tensor<Linear, DataType::FLOAT32> bias;  // optional
260 };
261 
262 struct Convolution3DAttributes {
263   HWD strides = HWD(0, 0, 0);    // Along each axis.
264   HWD dilations = HWD(0, 0, 0);  // Along each axis.
265   Padding3D padding;
266 
267   Tensor<OHWDI, DataType::FLOAT32> weights;
268   Tensor<Linear, DataType::FLOAT32> bias;  // optional
269 };
270 
271 // @return shape of a tensor after Convolution2D operation is applied to
272 //         the given input.
273 BHWC CalculateOutputShape(const BHWC& input,
274                           const Convolution2DAttributes& attr);
275 
276 // @return shape of a tensor after Convolution3D operation is applied to
277 //         the given input.
278 BHWDC CalculateOutputShape(const BHWDC& input,
279                            const Convolution3DAttributes& attr);
280 
281 // @return padding for convolution operation to make sure output keep the same
282 // shape as the given input.
283 Padding2D CalculateSamePadding(const BHWC& input,
284                                const Convolution2DAttributes& attr);
285 
286 // @return padding for convolution operation to make sure output keep the same
287 // shape as the given input.
288 Padding3D CalculateSamePadding(const BHWDC& input,
289                                const Convolution3DAttributes& attr);
290 
291 struct ConvolutionTransposedAttributes {
292   HW stride = HW(1, 1);  // Along each axis.
293   HW adjacent;           // TODO(sorokin): No op on Flow.
294   Padding2D padding;
295 
296   Tensor<OHWI, DataType::FLOAT32> weights;
297   Tensor<Linear, DataType::FLOAT32> bias;  // optional
298 };
299 
300 struct ConvolutionTransposed3DAttributes {
301   HWD stride = HWD(0, 0, 0);  // Along each axis.
302   Padding3D padding;
303 
304   Tensor<OHWDI, DataType::FLOAT32> weights;
305   Tensor<Linear, DataType::FLOAT32> bias;  // optional
306 };
307 
308 Padding2D CalculateSamePadding(const BHWC& input,
309                                const ConvolutionTransposedAttributes& attr);
310 
311 Padding3D CalculateSamePadding(const BHWDC& input,
312                                const ConvolutionTransposed3DAttributes& attr);
313 
314 // @return shape of a tensor after ConvolutionTransposed operation is applied to
315 //         the given input.
316 BHWC CalculateOutputShape(const BHWC& input,
317                           const ConvolutionTransposedAttributes& attr);
318 
319 // @return shape of a tensor after ConvolutionTransposed3D operation is applied
320 // to
321 //         the given input.
322 BHWDC CalculateOutputShape(const BHWDC& input,
323                            const ConvolutionTransposed3DAttributes& attr);
324 
325 struct DepthwiseConvolution2DAttributes : public Convolution2DAttributes {};
326 struct DepthwiseConvolution3DAttributes : public Convolution3DAttributes {};
327 
328 // @return shape of a tensor after DepthwiseConvolution2D operation is applied
329 //         to the given input.
330 BHWC CalculateOutputShape(const BHWC& input,
331                           const DepthwiseConvolution2DAttributes& attr);
332 
333 // @return shape of a tensor after DepthwiseConvolution3D operation is applied
334 //         to the given input.
335 BHWDC CalculateOutputShape(const BHWDC& input,
336                            const DepthwiseConvolution3DAttributes& attr);
337 
338 // @return padding for depthwise convolution operation to make sure output keep
339 // the same shape as the given input.
340 Padding2D CalculateSamePadding(const BHWC& input,
341                                const DepthwiseConvolution2DAttributes& attr);
342 
343 // @return padding for depthwise convolution operation to make sure output keep
344 // the same shape as the given input.
345 Padding3D CalculateSamePadding(const BHWDC& input,
346                                const DepthwiseConvolution3DAttributes& attr);
347 
348 // f(x):= {
349 //   if x < 0  : x -> alpha * x
350 //   if x >= 0 : x -> min(clip, x)
351 // }
352 //
353 // Examples:
354 //   - ReLU: clip = 0, alpha = 0
355 //   - ReLU6: clip = 6, alpha = 0
356 //   - Leaky ReLU: clip = 0, alpha = a
357 struct ReLUAttributes {
358   // clip <= 0 mean it is not set.
359   float clip = 0;
360 
361   float alpha = 0;
362 };
363 
364 struct PReLUAttributes {
365   // clip <= 0 mean it is not set.
366   float clip = 0;
367 
368   // If alpha is linear, then it is sharded across CHANNELS axis, otherwise
369   // full shape alpha is required.
370   absl::variant<Tensor<Linear, DataType::FLOAT32>,
371                 Tensor<HWC, DataType::FLOAT32>>
372       alpha;
373 };
374 
375 struct ReduceAttributes {
376   std::set<Axis> dims;
377 };
378 
379 struct SoftmaxAttributes {
380   Axis axis = Axis::UNKNOWN;
381 };
382 
383 enum LstmKernelType {
384   FULL = 0,
385   BASIC = 1,  // Currently, only basic is supported.
386 };
387 
388 struct LstmAttributes {
389   LstmKernelType kernel_type = LstmKernelType::BASIC;
390 };
391 
392 enum class SamplingType {
393   UNKNOWN = 0,
394   NEAREST = 1,
395   BILINEAR = 2,
396 };
397 
398 struct Resize2DAttributes {
399   HW new_shape;
400 
401   SamplingType type = SamplingType::UNKNOWN;
402 
403   // If true, the centers of the 4 corner pixels of the input and output tensors
404   // are aligned, preserving the values at the corner pixels. Defaults to false.
405   bool align_corners = false;
406 
407   bool half_pixel_centers = false;
408 };
409 
410 // TODO(b/147771327): rename to Resize3D
411 struct Resize3DAttributes {
412   HWD new_shape;
413 
414   SamplingType type = SamplingType::NEAREST;
415 
416   // If true, the centers of the 8 corner pixels of the input and output tensors
417   // are aligned, preserving the values at the corner pixels. Defaults to false.
418   bool align_corners = false;
419 
420   bool half_pixel_centers = false;
421 };
422 
423 float CalculateResizeScale(int32_t input_size, int32_t output_size,
424                            const Resize2DAttributes& attr);
425 
426 float CalculateResizeScale(int32_t input_size, int32_t output_size,
427                            const Resize3DAttributes& attr);
428 
429 // @return shape of a tensor after scale operation is applied to the given
430 // input.
431 BHWC CalculateOutputShape(const BHWC& input, const Resize2DAttributes& attr);
432 
433 // @return shape of a tensor after scale operation is applied to the given
434 // input.
435 BHWDC CalculateOutputShape(const BHWDC& input, const Resize3DAttributes& attr);
436 
437 enum class PaddingContentType {
438   ZEROS = 0,
439   REFLECT = 1,
440   EDGE = 2,
441 };
442 
443 struct PadAttributes {
444   PaddingContentType type = PaddingContentType::ZEROS;
445 
446   BHWC prepended;
447   BHWC appended;
448 };
449 
450 // @return shape of a tensor after Pad operation is applied to the given input.
451 BHWC CalculateOutputShape(const BHWC& input, const PadAttributes& attr);
452 
453 struct Pad3DAttributes {
454   PaddingContentType type = PaddingContentType::ZEROS;
455 
456   BHWDC prepended;
457   BHWDC appended;
458 };
459 
460 // @return shape of a tensor after Pad3D operation is applied to the given
461 // input.
462 BHWDC CalculateOutputShape(const BHWDC& input, const Pad3DAttributes& attr);
463 
464 struct ConstTensorAttributes {
465   Tensor<BHWC, DataType::FLOAT32> tensor;
466 };
467 
468 // Simple slicing without advanced support for shrinking, reverse slicing etc.
469 struct SliceAttributes {
470   // Specifies start and end dimensions for slicing.
471   BHWC starts;
472   BHWC ends;
473 
474   // Stride should be >= 1.
475   BHWC strides;
476 };
477 
478 // @return shape of a tensor after Slice2D operation is applied to the given
479 //         input.
480 BHWC CalculateOutputShape(const BHWC& input, const SliceAttributes& attr);
481 
482 // Simple slicing without advanced support for shrinking, reverse slicing etc.
483 struct Slice3DAttributes {
484   // Specifies start and end dimensions for slicing.
485   BHWDC starts;
486   BHWDC ends;
487 
488   // Stride should be >= 1.
489   BHWDC strides;
490 };
491 
492 // @return shape of a tensor after Slice3D operation is applied to the given
493 //         input.
494 BHWDC CalculateOutputShape(const BHWDC& input, const Slice3DAttributes& attr);
495 
496 struct FullyConnectedAttributes {
497   Tensor<OHWI, DataType::FLOAT32> weights;
498   Tensor<Linear, DataType::FLOAT32> bias;
499 };
500 
501 // @return shape of a tensor after FullyConnected operation is applied to
502 // the given input.
503 BHWC CalculateOutputShape(const BHWC& input,
504                           const FullyConnectedAttributes& attr);
505 
506 // @return shape of a tensor after Mean operation is applied to the given input.
507 BHWC CalculateOutputShape(const BHWC& input, const MeanAttributes& attr);
508 
509 // @return shape of a tensor after Mean operation is applied to the given input.
510 BHWDC CalculateOutputShape(const BHWDC& input, const MeanAttributes& attr);
511 
512 struct ElementwiseAttributes {
513   TensorOrScalar param;
514   // For elementwise operation with 2 inputs op(A, B), runtime_tensor_is_second
515   // true when runtime tensor is B(on second position). this is important for
516   // ops that non commutative, for example substract.
517   bool runtime_tensor_is_second = false;
518 };
519 
520 struct ReshapeAttributes {
521   BHWC new_shape;
522 };
523 
524 struct Reshape3DAttributes {
525   BHWDC new_shape;
526 };
527 
528 struct TransposeAttributes {
529   // A permutation of the dimensions of input tensor
530   BHWC perm;
531 };
532 
533 // @return shape of a tensor after Transpose operation is applied to
534 // the given input.
535 BHWC CalculateOutputShape(const BHWC& input, const TransposeAttributes& attr);
536 
537 struct Transpose3DAttributes {
538   // A permutation of the dimensions of input tensor
539   BHWDC perm;
540 };
541 
542 // @return shape of a tensor after Transpose3D operation is applied to
543 // the given input.
544 BHWDC CalculateOutputShape(const BHWDC& input,
545                            const Transpose3DAttributes& attr);
546 
547 struct SpaceToDepthAttributes {
548   int block_size;
549 };
550 
551 struct SplitAttributes {
552   // Defines axis by which to split.
553   Axis axis = Axis::UNKNOWN;
554 };
555 
556 // These help perform a combination of Quantize & Dequantize to adjust float
557 // values like quantized inference would.
558 struct QuantizeAndDequantizeAttributes {
559   float min = 0;
560   float max = 0;
561   float scale = 0;
562 };
563 
564 }  // namespace gpu
565 }  // namespace tflite
566 
567 #endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_OPERATIONS_H_
568