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/gl/kernels/registry.h"
17
18 #include <functional>
19 #include <memory>
20 #include <string>
21 #include <utility>
22 #include <vector>
23
24 #include "absl/container/flat_hash_map.h"
25 #include "absl/memory/memory.h"
26 #include "absl/strings/str_cat.h"
27 #include "absl/strings/str_join.h"
28 #include "tensorflow/lite/delegates/gpu/common/operations.h"
29 #include "tensorflow/lite/delegates/gpu/common/status.h"
30 #include "tensorflow/lite/delegates/gpu/gl/kernels/add.h"
31 #include "tensorflow/lite/delegates/gpu/gl/kernels/concat.h"
32 #include "tensorflow/lite/delegates/gpu/gl/kernels/conv.h"
33 #include "tensorflow/lite/delegates/gpu/gl/kernels/custom_registry.h"
34 #include "tensorflow/lite/delegates/gpu/gl/kernels/depthwise_conv.h"
35 #include "tensorflow/lite/delegates/gpu/gl/kernels/elementwise.h"
36 #include "tensorflow/lite/delegates/gpu/gl/kernels/fully_connected.h"
37 #include "tensorflow/lite/delegates/gpu/gl/kernels/lstm.h"
38 #include "tensorflow/lite/delegates/gpu/gl/kernels/mean.h"
39 #include "tensorflow/lite/delegates/gpu/gl/kernels/mul.h"
40 #include "tensorflow/lite/delegates/gpu/gl/kernels/pad.h"
41 #include "tensorflow/lite/delegates/gpu/gl/kernels/pooling.h"
42 #include "tensorflow/lite/delegates/gpu/gl/kernels/prelu.h"
43 #include "tensorflow/lite/delegates/gpu/gl/kernels/quantize_and_dequantize.h"
44 #include "tensorflow/lite/delegates/gpu/gl/kernels/relu.h"
45 #include "tensorflow/lite/delegates/gpu/gl/kernels/reshape.h"
46 #include "tensorflow/lite/delegates/gpu/gl/kernels/resize.h"
47 #include "tensorflow/lite/delegates/gpu/gl/kernels/slice.h"
48 #include "tensorflow/lite/delegates/gpu/gl/kernels/softmax.h"
49 #include "tensorflow/lite/delegates/gpu/gl/kernels/space_to_depth.h"
50 #include "tensorflow/lite/delegates/gpu/gl/kernels/tile.h"
51 #include "tensorflow/lite/delegates/gpu/gl/kernels/transpose_conv.h"
52
53 #ifndef TFLITE_GPU_BINARY_RELEASE
54 #include "tensorflow/lite/delegates/gpu/gl/kernels/max_unpooling.h"
55 #endif // TFLITE_GPU_BINARY_RELEASE
56
57 namespace tflite {
58 namespace gpu {
59 namespace gl {
60 namespace {
61
62 class Registry : public NodeShader {
63 public:
Registry()64 Registry() {
65 using Type = OperationType;
66 using NewShaderFunc = std::function<std::unique_ptr<NodeShader>()>;
67
68 const auto insert_op = [&](Type type, NewShaderFunc func) {
69 shaders_[ToString(type)].push_back(func());
70 };
71 const auto insert_elementwise_op = [&](Type operation_type) {
72 shaders_[ToString(operation_type)].push_back(
73 NewElementwiseNodeShader(operation_type));
74 };
75
76 insert_op(Type::ADD, NewAddNodeShader);
77 insert_op(Type::CONCAT, NewAlignedConcatNodeShader);
78 insert_op(Type::CONCAT, NewFlatConcatNodeShader);
79 insert_op(Type::CONCAT, NewConcatNodeShader);
80 insert_op(Type::CONVOLUTION_2D, NewConvolution1x1NodeShader);
81 insert_op(Type::CONVOLUTION_2D, NewConvolutionNodeShader);
82 insert_op(Type::CONVOLUTION_TRANSPOSED, NewConvolutionTransposedNodeShader);
83 insert_op(Type::DEPTHWISE_CONVOLUTION, NewDepthwiseConvolutionNodeShader);
84 insert_op(Type::DEPTH_TO_SPACE, NewDepthToSpaceNodeShader);
85 insert_op(Type::FULLY_CONNECTED, NewFullyConnectedNodeShader);
86 insert_op(Type::LSTM, NewLstmNodeShader);
87 insert_op(Type::MEAN, NewMeanNodeShader);
88 // TODO(b/162763635): implement MeanStddevNormalization for OpenGL.
89 insert_op(Type::MUL, NewMultiplyNodeShader);
90 insert_op(Type::PAD, NewPadNodeShader);
91 insert_op(Type::POOLING_2D, NewPoolingNodeShader);
92 insert_op(Type::PRELU, NewPReLUNodeShader);
93 insert_op(Type::QUANTIZE_AND_DEQUANTIZE,
94 NewQuantizeAndDequantizeNodeShader);
95 insert_op(Type::RELU, NewReLUNodeShader);
96 insert_op(Type::RESIZE, NewResizeNodeShader);
97 insert_op(Type::RESHAPE, NewReshapeNodeShader);
98 insert_op(Type::SLICE, NewSliceNodeShader);
99 insert_op(Type::SOFTMAX, NewSoftmaxNodeShader);
100 insert_op(Type::SPACE_TO_DEPTH, NewSpaceToDepthNodeShader);
101 insert_op(Type::TILE, NewTileNodeShader);
102
103 insert_elementwise_op(Type::ABS);
104 insert_elementwise_op(Type::COPY);
105 insert_elementwise_op(Type::COS);
106 insert_elementwise_op(Type::DIV);
107 insert_elementwise_op(Type::ELU);
108 insert_elementwise_op(Type::EXP);
109 insert_elementwise_op(Type::FLOOR);
110 insert_elementwise_op(Type::FLOOR_DIV);
111 insert_elementwise_op(Type::FLOOR_MOD);
112 insert_elementwise_op(Type::HARD_SWISH);
113 insert_elementwise_op(Type::LOG);
114 insert_elementwise_op(Type::NEG);
115 insert_elementwise_op(Type::MAXIMUM);
116 insert_elementwise_op(Type::MINIMUM);
117 insert_elementwise_op(Type::POW);
118 insert_elementwise_op(Type::RSQRT);
119 insert_elementwise_op(Type::SIGMOID);
120 insert_elementwise_op(Type::SIN);
121 insert_elementwise_op(Type::SQRT);
122 insert_elementwise_op(Type::SQUARE);
123 insert_elementwise_op(Type::SQUARED_DIFF);
124 insert_elementwise_op(Type::SUB);
125 insert_elementwise_op(Type::TANH);
126
127 #ifndef TFLITE_GPU_BINARY_RELEASE
128 insert_op(Type::MAX_UNPOOLING_2D, NewMaxUnpoolingNodeShader);
129 RegisterCustomOps(&shaders_);
130 #endif // TFLITE_GPU_BINARY_RELEASE
131 }
132
133 ~Registry() final = default;
134
GenerateCode(const GenerationContext & ctx,GeneratedCode * generated_code) const135 absl::Status GenerateCode(const GenerationContext& ctx,
136 GeneratedCode* generated_code) const final {
137 auto it = shaders_.find(ctx.op_type);
138 if (it == shaders_.end()) {
139 return absl::NotFoundError(
140 absl::StrCat("No shader implementation for ", ctx.op_type));
141 }
142 std::vector<std::string> errors;
143 for (const auto& shader : it->second) {
144 const auto status = shader->GenerateCode(ctx, generated_code);
145 // Return the first suitable shader.
146 if (status.ok()) return absl::OkStatus();
147 errors.push_back(std::string(status.message()));
148 }
149 return errors.empty() ? absl::OkStatus()
150 : absl::UnknownError(absl::StrJoin(errors, ", "));
151 }
152
153 private:
154 absl::flat_hash_map<std::string, std::vector<std::unique_ptr<NodeShader>>>
155 shaders_;
156 };
157
158 } // namespace
159
NewNodeShaderRegistry()160 std::unique_ptr<NodeShader> NewNodeShaderRegistry() {
161 return absl::make_unique<Registry>();
162 }
163
164 } // namespace gl
165 } // namespace gpu
166 } // namespace tflite
167