1 /* Copyright 2019 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/tasks/softmax1x1.h"
17
18 #include <string>
19
20 #include "tensorflow/lite/delegates/gpu/common/status.h"
21 #include "tensorflow/lite/delegates/gpu/common/task/util.h"
22
23 namespace tflite {
24 namespace gpu {
25
Softmax1x1(const OperationDef & definition)26 Softmax1x1::Softmax1x1(const OperationDef& definition)
27 : GPUOperation(definition) {
28 work_group_size_ = int3(32, 1, 1);
29 code_ = GetSoftmaxKernelCode(definition_);
30 }
31
Softmax1x1(Softmax1x1 && kernel)32 Softmax1x1::Softmax1x1(Softmax1x1&& kernel) : GPUOperation(std::move(kernel)) {}
33
operator =(Softmax1x1 && kernel)34 Softmax1x1& Softmax1x1::operator=(Softmax1x1&& kernel) {
35 if (this != &kernel) {
36 GPUOperation::operator=(std::move(kernel));
37 }
38 return *this;
39 }
40
GetSoftmaxKernelCode(const OperationDef & op_def)41 std::string Softmax1x1::GetSoftmaxKernelCode(const OperationDef& op_def) {
42 AddSrcTensor("src_tensor", op_def.src_tensors[0]);
43 AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
44 args_.AddFloat("mask_x");
45 args_.AddFloat("mask_y");
46 args_.AddFloat("mask_z");
47 args_.AddFloat("mask_w");
48
49 std::string c;
50 c += "MAIN_FUNCTION($0) {\n";
51 if (op_def.IsBatchSupported()) {
52 c += " int batch_id = GLOBAL_ID_1;\n";
53 c += " if (batch_id >= args.dst_tensor.Batch()) return;\n";
54 c += " args.dst_tensor.SetBatchRef(batch_id);\n";
55 c += " args.src_tensor.SetBatchRef(batch_id);\n";
56 }
57 c += " float4 mask = INIT_FLOAT4v4(args.mask_x, args.mask_y, args.mask_z, "
58 "args.mask_w);\n";
59 c +=
60 " float4 maxx4 = INIT_FLOAT4(args.src_tensor.Read<float>(0, 0, 0).x);\n";
61 c += " int tid = LOCAL_ID_0;\n";
62 c += " for (int s = tid; s < args.src_tensor.Slices(); s += 32) {\n";
63 c += " float4 mask_a = s == args.src_tensor.Slices() - 1 ? mask : "
64 "INIT_FLOAT4(1.0f);\n";
65 c += " float4 mask_b = INIT_FLOAT4(1.0f) - mask_a;\n";
66 c += " float4 src = args.src_tensor.Read<float>(0, 0, s);\n";
67 c += " src = src * mask_a + mask_b * src.x;\n";
68 c += " maxx4 = max(maxx4, src);\n";
69 c += " }\n";
70 c += " float maximum = max(maxx4.x, maxx4.y);\n";
71 c += " maximum = max(maximum, maxx4.z);\n";
72 c += " maximum = max(maximum, maxx4.w);\n";
73 c += " __local float4 tmp[8];\n";
74 c += " __local float* tmpx1 = (__local float*)tmp;\n";
75 c += " tmpx1[tid] = maximum;\n";
76 c += " LOCAL_MEM_BARRIER;\n";
77 c += " if (tid == 0) {\n";
78 c += " maxx4 = max(tmp[0], tmp[1]);\n";
79 c += " maxx4 = max(maxx4, tmp[2]);\n";
80 c += " maxx4 = max(maxx4, tmp[3]);\n";
81 c += " maxx4 = max(maxx4, tmp[4]);\n";
82 c += " maxx4 = max(maxx4, tmp[5]);\n";
83 c += " maxx4 = max(maxx4, tmp[6]);\n";
84 c += " maxx4 = max(maxx4, tmp[7]);\n";
85 c += " maximum = max(maxx4.x, maxx4.y);\n";
86 c += " maximum = max(maximum, maxx4.z);\n";
87 c += " maximum = max(maximum, maxx4.w);\n";
88 c += " tmpx1[0] = maximum;\n";
89 c += " }\n";
90 c += " LOCAL_MEM_BARRIER;\n";
91 c += " maximum = tmpx1[0];\n";
92 c += " float sum = 0.0f;\n";
93 c += " for (int s = tid; s < args.src_tensor.Slices(); s += 32) {\n";
94 c += " float4 mask_temp = s == args.src_tensor.Slices() - 1 ? mask : "
95 "INIT_FLOAT4(1.0f);\n";
96 c += " float4 src = args.src_tensor.Read<float>(0, 0, s) - "
97 "INIT_FLOAT4(maximum);\n";
98 c += " sum += dot(mask_temp, exp(src));\n";
99 c += " }\n";
100 c += " LOCAL_MEM_BARRIER;\n";
101 c += " tmpx1[tid] = sum;\n";
102 c += " LOCAL_MEM_BARRIER;\n";
103 c += " if (tid == 0) {\n";
104 c += " sum = dot((float4)(1.0f), tmp[0]);\n";
105 c += " sum += dot((float4)(1.0f), tmp[1]);\n";
106 c += " sum += dot((float4)(1.0f), tmp[2]);\n";
107 c += " sum += dot((float4)(1.0f), tmp[3]);\n";
108 c += " sum += dot((float4)(1.0f), tmp[4]);\n";
109 c += " sum += dot((float4)(1.0f), tmp[5]);\n";
110 c += " sum += dot((float4)(1.0f), tmp[6]);\n";
111 c += " sum += dot((float4)(1.0f), tmp[7]);\n";
112 c += " tmpx1[0] = 1.0f / sum;\n";
113 c += " }\n";
114 c += " LOCAL_MEM_BARRIER;\n";
115 c += " sum = tmpx1[0];\n";
116 c += "\n";
117 c += " int dst_s = GLOBAL_ID_0;\n";
118 c += " if (dst_s < args.dst_tensor.Slices()) {\n";
119 c += " float4 src = args.src_tensor.Read<float>(0, 0, dst_s) - "
120 "INIT_FLOAT4(maximum);\n";
121 c += " FLT4 res = TO_FLT4(exp(src) * sum);\n";
122 c += " args.dst_tensor.Write(res, 0, 0, dst_s);\n";
123 c += " }\n";
124 c += "}\n";
125 return c;
126 }
127
BindArguments(ArgumentsBinder * args)128 absl::Status Softmax1x1::BindArguments(ArgumentsBinder* args) {
129 float4 mask = GetMaskForLastPlane(src_[0]->Channels());
130 RETURN_IF_ERROR(args->SetFloat("mask_x", mask.x));
131 RETURN_IF_ERROR(args->SetFloat("mask_y", mask.y));
132 RETURN_IF_ERROR(args->SetFloat("mask_z", mask.z));
133 RETURN_IF_ERROR(args->SetFloat("mask_w", mask.w));
134 return absl::OkStatus();
135 }
136
GetGridSize() const137 int3 Softmax1x1::GetGridSize() const {
138 return int3(dst_[0]->Slices(), dst_[0]->Batch(), 1);
139 }
140
CreateSoftmax1x1(const OperationDef & definition)141 Softmax1x1 CreateSoftmax1x1(const OperationDef& definition) {
142 return Softmax1x1(definition);
143 }
144
145 } // namespace gpu
146 } // namespace tflite
147