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