• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 #include "tensorflow/lite/delegates/gpu/common/selectors/simple_selectors.h"
17 
18 #include <memory>
19 #include <set>
20 
21 #include "absl/memory/memory.h"
22 #include "tensorflow/lite/delegates/gpu/common/status.h"
23 #include "tensorflow/lite/delegates/gpu/common/tasks/add.h"
24 #include "tensorflow/lite/delegates/gpu/common/tasks/concat_xy.h"
25 #include "tensorflow/lite/delegates/gpu/common/tasks/concat_z.h"
26 #include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv.h"
27 #include "tensorflow/lite/delegates/gpu/common/tasks/gather.h"
28 #include "tensorflow/lite/delegates/gpu/common/tasks/lstm.h"
29 #include "tensorflow/lite/delegates/gpu/common/tasks/max_unpooling.h"
30 #include "tensorflow/lite/delegates/gpu/common/tasks/padding.h"
31 #include "tensorflow/lite/delegates/gpu/common/tasks/pooling.h"
32 #include "tensorflow/lite/delegates/gpu/common/tasks/prelu.h"
33 #include "tensorflow/lite/delegates/gpu/common/tasks/quantize_and_dequantize.h"
34 #include "tensorflow/lite/delegates/gpu/common/tasks/reduce.h"
35 #include "tensorflow/lite/delegates/gpu/common/tasks/relu.h"
36 #include "tensorflow/lite/delegates/gpu/common/tasks/resampler.h"
37 #include "tensorflow/lite/delegates/gpu/common/tasks/reshape.h"
38 #include "tensorflow/lite/delegates/gpu/common/tasks/reshapex4.h"
39 #include "tensorflow/lite/delegates/gpu/common/tasks/resize.h"
40 #include "tensorflow/lite/delegates/gpu/common/tasks/softmax.h"
41 #include "tensorflow/lite/delegates/gpu/common/tasks/softmax1x1.h"
42 #include "tensorflow/lite/delegates/gpu/common/tasks/space_to_depth.h"
43 #include "tensorflow/lite/delegates/gpu/common/tasks/split.h"
44 #include "tensorflow/lite/delegates/gpu/common/tasks/strided_slice.h"
45 #include "tensorflow/lite/delegates/gpu/common/tasks/tile.h"
46 #include "tensorflow/lite/delegates/gpu/common/tasks/transpose.h"
47 #include "tensorflow/lite/delegates/gpu/common/tasks/winograd.h"
48 
49 namespace tflite {
50 namespace gpu {
51 
SelectLSTM(const OperationDef & op_def,const GpuInfo & gpu_info)52 std::unique_ptr<GPUOperation> SelectLSTM(const OperationDef& op_def,
53                                          const GpuInfo& gpu_info) {
54   return absl::make_unique<GPUOperation>(CreateLSTM(op_def, gpu_info));
55 }
56 
SelectReLU(const ReLUAttributes & attr,const OperationDef & op_def)57 std::unique_ptr<GPUOperation> SelectReLU(const ReLUAttributes& attr,
58                                          const OperationDef& op_def) {
59   return absl::make_unique<GPUOperation>(CreateReLU(op_def, attr));
60 }
61 
SelectPReLU(const PReLUAttributes & attr,const GpuInfo & gpu_info,const OperationDef & op_def)62 std::unique_ptr<GPUOperation> SelectPReLU(const PReLUAttributes& attr,
63                                           const GpuInfo& gpu_info,
64                                           const OperationDef& op_def) {
65   return absl::make_unique<GPUOperation>(CreatePReLU(gpu_info, op_def, attr));
66 }
67 
SelectPooling(const Pooling2DAttributes & attr,const OperationDef & op_def)68 std::unique_ptr<GPUOperation> SelectPooling(const Pooling2DAttributes& attr,
69                                             const OperationDef& op_def) {
70   return absl::make_unique<GPUOperation>(CreatePooling(op_def, attr));
71 }
72 
SelectMaxUnpooling(const MaxUnpooling2DAttributes & attr,const OperationDef & op_def)73 std::unique_ptr<GPUOperation> SelectMaxUnpooling(
74     const MaxUnpooling2DAttributes& attr, const OperationDef& op_def) {
75   return absl::make_unique<GPUOperation>(CreateMaxUnpooling(op_def, attr));
76 }
77 
SelectAdd(const OperationDef & op_def,const std::vector<int> & channels,int dst_channels,std::unique_ptr<GPUOperation> * ptr)78 void SelectAdd(const OperationDef& op_def, const std::vector<int>& channels,
79                int dst_channels, std::unique_ptr<GPUOperation>* ptr) {
80   GPUOperation operation = CreateAdd(op_def, channels, dst_channels);
81   *ptr = absl::make_unique<GPUOperation>(std::move(operation));
82 }
83 
SelectGather(const GatherAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)84 absl::Status SelectGather(const GatherAttributes& attr,
85                           const OperationDef& op_def,
86                           std::unique_ptr<GPUOperation>* ptr) {
87   if (attr.axis != Axis::WIDTH) {
88     return absl::UnimplementedError(
89         "No gather for this axis. Only Width axis supported.");
90   }
91   GPUOperation operation = CreateGather(op_def, attr);
92   *ptr = absl::make_unique<GPUOperation>(std::move(operation));
93   return absl::OkStatus();
94 }
95 
SelectResampler(const OperationDef & op_def)96 std::unique_ptr<GPUOperation> SelectResampler(const OperationDef& op_def) {
97   GPUOperation operation = CreateResampler(op_def);
98   return absl::make_unique<GPUOperation>(std::move(operation));
99 }
100 
SelectResize(const Resize2DAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)101 absl::Status SelectResize(const Resize2DAttributes& attr,
102                           const OperationDef& op_def,
103                           std::unique_ptr<GPUOperation>* ptr) {
104   Resize operation = CreateResize(op_def, attr);
105   *ptr = absl::make_unique<Resize>(std::move(operation));
106   return absl::OkStatus();
107 }
108 
SelectConcat(const ConcatAttributes & attr,const std::vector<int> & channels,const OperationDef & op_def,const GpuInfo & gpu_info,std::unique_ptr<GPUOperation> * ptr)109 absl::Status SelectConcat(const ConcatAttributes& attr,
110                           const std::vector<int>& channels,
111                           const OperationDef& op_def, const GpuInfo& gpu_info,
112                           std::unique_ptr<GPUOperation>* ptr) {
113   switch (attr.axis) {
114     case Axis::CHANNELS: {
115       GPUOperation operation = CreateConcatZ(op_def, channels, gpu_info);
116       *ptr = absl::make_unique<GPUOperation>(std::move(operation));
117       return absl::OkStatus();
118     }
119     case Axis::BATCH:
120     case Axis::DEPTH:
121     case Axis::HEIGHT:
122     case Axis::WIDTH: {
123       GPUOperation operation = CreateConcatXY(op_def, attr);
124       *ptr = absl::make_unique<GPUOperation>(std::move(operation));
125       return absl::OkStatus();
126     }
127     default:
128       return absl::UnimplementedError("No concat for this axis.");
129   }
130 }
131 
SelectDWConvolutionDynamicWeights(const DepthwiseConvolution2DAttributes & attr,const GpuInfo & gpu_info,const OperationDef & op_def)132 std::unique_ptr<GPUOperation> SelectDWConvolutionDynamicWeights(
133     const DepthwiseConvolution2DAttributes& attr, const GpuInfo& gpu_info,
134     const OperationDef& op_def) {
135   return absl::make_unique<GPUOperation>(
136       CreateDepthwiseConvolution2DDynamicWeights(gpu_info, op_def, attr));
137 }
138 
SelectReshape(int src_channels,int dst_channels,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)139 void SelectReshape(int src_channels, int dst_channels,
140                    const OperationDef& op_def,
141                    std::unique_ptr<GPUOperation>* ptr) {
142   if (src_channels % 4 == 0 && dst_channels % 4 == 0) {
143     GPUOperation operation = CreateReshapex4(op_def);
144     *ptr = absl::make_unique<GPUOperation>(std::move(operation));
145   } else {
146     GPUOperation operation = CreateReshape(op_def);
147     *ptr = absl::make_unique<GPUOperation>(std::move(operation));
148   }
149 }
150 
SelectSpaceToDepth(const SpaceToDepthAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)151 void SelectSpaceToDepth(const SpaceToDepthAttributes& attr,
152                         const OperationDef& op_def,
153                         std::unique_ptr<GPUOperation>* ptr) {
154   GPUOperation operation = CreateSpaceToDepth(op_def, attr);
155   *ptr = absl::make_unique<GPUOperation>(std::move(operation));
156 }
157 
SelectDepthToSpace(const SpaceToDepthAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)158 void SelectDepthToSpace(const SpaceToDepthAttributes& attr,
159                         const OperationDef& op_def,
160                         std::unique_ptr<GPUOperation>* ptr) {
161   GPUOperation operation = CreateDepthToSpace(op_def, attr);
162   *ptr = absl::make_unique<GPUOperation>(std::move(operation));
163 }
164 
SelectSplit(const SplitAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)165 void SelectSplit(const SplitAttributes& attr, const OperationDef& op_def,
166                  std::unique_ptr<GPUOperation>* ptr) {
167   Split operation = CreateSplit(op_def, attr);
168   *ptr = absl::make_unique<Split>(std::move(operation));
169 }
170 
SelectPadding(const PadAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)171 void SelectPadding(const PadAttributes& attr, const OperationDef& op_def,
172                    std::unique_ptr<GPUOperation>* ptr) {
173   GPUOperation operation = CreatePadding(op_def, attr);
174   *ptr = absl::make_unique<GPUOperation>(std::move(operation));
175 }
176 
SelectStridedSlice(const SliceAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)177 void SelectStridedSlice(const SliceAttributes& attr, const OperationDef& op_def,
178                         std::unique_ptr<GPUOperation>* ptr) {
179   StridedSlice operation = CreateStridedSlice(op_def, attr);
180   *ptr = absl::make_unique<StridedSlice>(std::move(operation));
181 }
182 
SelectReduce(const std::set<Axis> & axis_to_reduce,const BHWC & src_shape,OperationType op_type,const OperationDef & op_def,const GpuInfo & gpu_info)183 std::unique_ptr<GPUOperation> SelectReduce(const std::set<Axis>& axis_to_reduce,
184                                            const BHWC& src_shape,
185                                            OperationType op_type,
186                                            const OperationDef& op_def,
187                                            const GpuInfo& gpu_info) {
188   return absl::make_unique<Reduce>(
189       CreateReduce(axis_to_reduce, src_shape, op_type, op_def, gpu_info));
190 }
191 
SelectSoftmax(const BHWC & shape,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)192 void SelectSoftmax(const BHWC& shape, const OperationDef& op_def,
193                    std::unique_ptr<GPUOperation>* ptr) {
194   if (shape.w == 1 && shape.h == 1) {
195     Softmax1x1 operation = CreateSoftmax1x1(op_def);
196     *ptr = absl::make_unique<Softmax1x1>(std::move(operation));
197   } else {
198     GPUOperation operation = CreateSoftmax(op_def);
199     *ptr = absl::make_unique<GPUOperation>(std::move(operation));
200   }
201 }
202 
SelectTile(const OperationDef & op_def,const BHWC & src_shape)203 std::unique_ptr<GPUOperation> SelectTile(const OperationDef& op_def,
204                                          const BHWC& src_shape) {
205   return absl::make_unique<GPUOperation>(CreateTile(op_def, src_shape.c));
206 }
207 
SelectTranspose(const TransposeAttributes & attr,const OperationDef & op_def,std::unique_ptr<GPUOperation> * ptr)208 void SelectTranspose(const TransposeAttributes& attr,
209                      const OperationDef& op_def,
210                      std::unique_ptr<GPUOperation>* ptr) {
211   GPUOperation operation = CreateTranspose(op_def, attr);
212   *ptr = absl::make_unique<GPUOperation>(std::move(operation));
213 }
214 
SelectWinograd4x4To36(const GpuInfo & gpu_info,const Padding2D & padding,const OperationDef & op_def)215 std::unique_ptr<GPUOperation> SelectWinograd4x4To36(
216     const GpuInfo& gpu_info, const Padding2D& padding,
217     const OperationDef& op_def) {
218   if (gpu_info.IsApple()) {
219     const auto src_storage = op_def.src_tensors[0].storage_type;
220     const auto dst_storage = op_def.src_tensors[0].storage_type;
221     if ((src_storage == TensorStorageType::BUFFER ||
222          src_storage == TensorStorageType::IMAGE_BUFFER) &&
223         (dst_storage == TensorStorageType::BUFFER ||
224          dst_storage == TensorStorageType::IMAGE_BUFFER)) {
225       Winograd4x4To36 operation = CreateWinograd4x4To36(op_def, padding);
226       return absl::make_unique<Winograd4x4To36>(std::move(operation));
227     }
228   }
229   return absl::make_unique<Winograd4x4To36TileX6>(
230       CreateWinograd4x4To36TileX6(gpu_info, op_def, padding));
231 }
232 
SelectWinograd36To4x4(const GpuInfo & gpu_info,const OperationDef & op_def,const tflite::gpu::Tensor<Linear,DataType::FLOAT32> & biases)233 std::unique_ptr<GPUOperation> SelectWinograd36To4x4(
234     const GpuInfo& gpu_info, const OperationDef& op_def,
235     const tflite::gpu::Tensor<Linear, DataType::FLOAT32>& biases) {
236   if (gpu_info.IsApple()) {
237     const auto src_storage = op_def.src_tensors[0].storage_type;
238     const auto dst_storage = op_def.src_tensors[0].storage_type;
239     if ((src_storage == TensorStorageType::BUFFER ||
240          src_storage == TensorStorageType::IMAGE_BUFFER) &&
241         (dst_storage == TensorStorageType::BUFFER ||
242          dst_storage == TensorStorageType::IMAGE_BUFFER)) {
243       Winograd36To4x4 operation = CreateWinograd36To4x4(op_def, biases);
244       return absl::make_unique<Winograd36To4x4>(std::move(operation));
245     }
246   }
247   return absl::make_unique<Winograd36To4x4Tile4x1>(
248       CreateWinograd36To4x4Tile4x1(gpu_info, op_def, biases));
249 }
250 
SelectQuantizeAndDequantize(const QuantizeAndDequantizeAttributes & attr,const OperationDef & op_def)251 std::unique_ptr<GPUOperation> SelectQuantizeAndDequantize(
252     const QuantizeAndDequantizeAttributes& attr, const OperationDef& op_def) {
253   return absl::make_unique<GPUOperation>(
254       CreateQuantizeAndDequantize(op_def, attr));
255 }
256 
257 }  // namespace gpu
258 }  // namespace tflite
259