xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
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_TASK_GPU_OBJECT_DESC_H_
17 #define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OBJECT_DESC_H_
18 
19 #include <map>
20 #include <memory>
21 #include <string>
22 #include <vector>
23 
24 #include "tensorflow/lite/delegates/gpu/common/access_type.h"
25 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
26 #include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
27 #include "tensorflow/lite/delegates/gpu/common/status.h"
28 #include "tensorflow/lite/delegates/gpu/common/task/serialization_base_generated.h"
29 
30 namespace tflite {
31 namespace gpu {
32 
33 struct GPUImage2DDescriptor {
34   DataType data_type;
35   bool normalized = false;   // used with INT data types, if normalized, we read
36                              // in kernel float data.
37   DataType normalized_type;  // can be FLOAT32 or FLOAT16, using with normalized
38                              // = true
39   AccessType access_type;
40 };
41 
42 struct GPUImage3DDescriptor {
43   DataType data_type;
44   AccessType access_type;
45 };
46 
47 struct GPUImage2DArrayDescriptor {
48   DataType data_type;
49   AccessType access_type;
50 };
51 
52 struct GPUImageBufferDescriptor {
53   DataType data_type;
54   AccessType access_type;
55 };
56 
57 struct GPUCustomMemoryDescriptor {
58   std::string type_name;
59 };
60 
61 enum class MemoryType { GLOBAL, CONSTANT, LOCAL };
62 
63 struct GPUBufferDescriptor {
64   DataType data_type;
65   AccessType access_type;
66   int element_size;
67   MemoryType memory_type = MemoryType::GLOBAL;
68   std::vector<std::string> attributes;
69 };
70 
71 struct GPUResources {
72   std::vector<std::string> ints;
73   std::vector<std::string> floats;
74   std::vector<std::pair<std::string, GPUBufferDescriptor>> buffers;
75   std::vector<std::pair<std::string, GPUImage2DDescriptor>> images2d;
76   std::vector<std::pair<std::string, GPUImage2DArrayDescriptor>> image2d_arrays;
77   std::vector<std::pair<std::string, GPUImage3DDescriptor>> images3d;
78   std::vector<std::pair<std::string, GPUImageBufferDescriptor>> image_buffers;
79   std::vector<std::pair<std::string, GPUCustomMemoryDescriptor>>
80       custom_memories;
81 
GetNamesGPUResources82   std::vector<std::string> GetNames() const {
83     std::vector<std::string> names = ints;
84     names.insert(names.end(), floats.begin(), floats.end());
85     for (const auto& obj : buffers) {
86       names.push_back(obj.first);
87     }
88     for (const auto& obj : images2d) {
89       names.push_back(obj.first);
90     }
91     for (const auto& obj : image2d_arrays) {
92       names.push_back(obj.first);
93     }
94     for (const auto& obj : images3d) {
95       names.push_back(obj.first);
96     }
97     for (const auto& obj : image_buffers) {
98       names.push_back(obj.first);
99     }
100     for (const auto& obj : custom_memories) {
101       names.push_back(obj.first);
102     }
103     return names;
104   }
105 
GetReadImagesCountGPUResources106   int GetReadImagesCount() const {
107     int counter = 0;
108     for (const auto& t : images2d) {
109       if (t.second.access_type == tflite::gpu::AccessType::READ) {
110         counter++;
111       }
112     }
113     for (const auto& t : image2d_arrays) {
114       if (t.second.access_type == tflite::gpu::AccessType::READ) {
115         counter++;
116       }
117     }
118     for (const auto& t : images3d) {
119       if (t.second.access_type == tflite::gpu::AccessType::READ) {
120         counter++;
121       }
122     }
123     for (const auto& t : image_buffers) {
124       if (t.second.access_type == tflite::gpu::AccessType::READ) {
125         counter++;
126       }
127     }
128     return counter;
129   }
130 
GetWriteImagesCountGPUResources131   int GetWriteImagesCount() const {
132     int counter = 0;
133     for (const auto& t : images2d) {
134       if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
135         counter++;
136       }
137     }
138     for (const auto& t : image2d_arrays) {
139       if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
140         counter++;
141       }
142     }
143     for (const auto& t : images3d) {
144       if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
145         counter++;
146       }
147     }
148     for (const auto& t : image_buffers) {
149       if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
150         counter++;
151       }
152     }
153     return counter;
154   }
155 };
156 
157 struct GenericGPUResourcesWithValue {
158   std::vector<std::pair<std::string, int>> ints;
159   std::vector<std::pair<std::string, float>> floats;
160 
AddFloatGenericGPUResourcesWithValue161   void AddFloat(const std::string& name, float value) {
162     floats.push_back({name, value});
163   }
AddIntGenericGPUResourcesWithValue164   void AddInt(const std::string& name, int value) {
165     ints.push_back({name, value});
166   }
167 };
168 
169 class GPUObjectDescriptor {
170  public:
171   GPUObjectDescriptor() = default;
172   GPUObjectDescriptor(const GPUObjectDescriptor&) = default;
173   GPUObjectDescriptor& operator=(const GPUObjectDescriptor&) = default;
174   GPUObjectDescriptor(GPUObjectDescriptor&& obj_desc) = default;
175   GPUObjectDescriptor& operator=(GPUObjectDescriptor&& obj_desc) = default;
176   virtual ~GPUObjectDescriptor() = default;
177 
SetStateVar(const std::string & key,const std::string & value)178   void SetStateVar(const std::string& key, const std::string& value) const {
179     state_vars_[key] = value;
180   }
181 
PerformConstExpr(const tflite::gpu::GpuInfo & gpu_info,const std::string & const_expr,std::string * result)182   virtual absl::Status PerformConstExpr(const tflite::gpu::GpuInfo& gpu_info,
183                                         const std::string& const_expr,
184                                         std::string* result) const {
185     return absl::UnimplementedError(
186         "No implementation of perform const expression");
187   }
188 
PerformSelector(const GpuInfo & gpu_info,const std::string & selector,const std::vector<std::string> & args,const std::vector<std::string> & template_args,std::string * result)189   virtual absl::Status PerformSelector(
190       const GpuInfo& gpu_info, const std::string& selector,
191       const std::vector<std::string>& args,
192       const std::vector<std::string>& template_args,
193       std::string* result) const {
194     return absl::UnimplementedError("No implementation of perform selector");
195   }
GetGPUResources(const GpuInfo & gpu_info)196   virtual GPUResources GetGPUResources(const GpuInfo& gpu_info) const {
197     return GPUResources();
198   }
199 
Release()200   virtual void Release() {}
201 
202   // For internal use, will work correct only for const objects and before
203   // Release() call.
GetSizeInBytes()204   virtual uint64_t GetSizeInBytes() const { return 0; }
205 
SetAccess(AccessType access_type)206   void SetAccess(AccessType access_type) { access_type_ = access_type; }
GetAccess()207   AccessType GetAccess() const { return access_type_; }
208 
209  protected:
210   friend flatbuffers::Offset<tflite::gpu::data::GPUObjectDescriptor> Encode(
211       const GPUObjectDescriptor& desc, flatbuffers::FlatBufferBuilder* builder);
212   friend void Decode(const tflite::gpu::data::GPUObjectDescriptor* fb_obj,
213                      GPUObjectDescriptor* obj);
214   mutable std::map<std::string, std::string> state_vars_;
215   AccessType access_type_;
216 };
217 
218 using GPUObjectDescriptorPtr = std::unique_ptr<GPUObjectDescriptor>;
219 
220 }  // namespace gpu
221 }  // namespace tflite
222 
223 #endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OBJECT_DESC_H_
224