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   AccessType access_type;
36 };
37 
38 struct GPUImage3DDescriptor {
39   DataType data_type;
40   AccessType access_type;
41 };
42 
43 struct GPUImage2DArrayDescriptor {
44   DataType data_type;
45   AccessType access_type;
46 };
47 
48 struct GPUImageBufferDescriptor {
49   DataType data_type;
50   AccessType access_type;
51 };
52 
53 struct GPUCustomMemoryDescriptor {
54   std::string type_name;
55 };
56 
57 enum class MemoryType { GLOBAL, CONSTANT, LOCAL };
58 
59 struct GPUBufferDescriptor {
60   DataType data_type;
61   AccessType access_type;
62   int element_size;
63   MemoryType memory_type = MemoryType::GLOBAL;
64   std::vector<std::string> attributes;
65 };
66 
67 struct GPUResources {
68   std::vector<std::string> ints;
69   std::vector<std::string> floats;
70   std::vector<std::pair<std::string, GPUBufferDescriptor>> buffers;
71   std::vector<std::pair<std::string, GPUImage2DDescriptor>> images2d;
72   std::vector<std::pair<std::string, GPUImage2DArrayDescriptor>> image2d_arrays;
73   std::vector<std::pair<std::string, GPUImage3DDescriptor>> images3d;
74   std::vector<std::pair<std::string, GPUImageBufferDescriptor>> image_buffers;
75   std::vector<std::pair<std::string, GPUCustomMemoryDescriptor>>
76       custom_memories;
77 
GetNamesGPUResources78   std::vector<std::string> GetNames() const {
79     std::vector<std::string> names = ints;
80     names.insert(names.end(), floats.begin(), floats.end());
81     for (const auto& obj : buffers) {
82       names.push_back(obj.first);
83     }
84     for (const auto& obj : images2d) {
85       names.push_back(obj.first);
86     }
87     for (const auto& obj : image2d_arrays) {
88       names.push_back(obj.first);
89     }
90     for (const auto& obj : images3d) {
91       names.push_back(obj.first);
92     }
93     for (const auto& obj : image_buffers) {
94       names.push_back(obj.first);
95     }
96     for (const auto& obj : custom_memories) {
97       names.push_back(obj.first);
98     }
99     return names;
100   }
101 };
102 
103 class GPUObjectDescriptor {
104  public:
105   GPUObjectDescriptor() = default;
106   GPUObjectDescriptor(const GPUObjectDescriptor&) = default;
107   GPUObjectDescriptor& operator=(const GPUObjectDescriptor&) = default;
108   GPUObjectDescriptor(GPUObjectDescriptor&& obj_desc) = default;
109   GPUObjectDescriptor& operator=(GPUObjectDescriptor&& obj_desc) = default;
110   virtual ~GPUObjectDescriptor() = default;
111 
SetStateVar(const std::string & key,const std::string & value)112   void SetStateVar(const std::string& key, const std::string& value) const {
113     state_vars_[key] = value;
114   }
115 
PerformConstExpr(const std::string & const_expr)116   virtual std::string PerformConstExpr(const std::string& const_expr) const {
117     return "";
118   }
119 
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)120   virtual absl::Status PerformSelector(
121       const GpuInfo& gpu_info, const std::string& selector,
122       const std::vector<std::string>& args,
123       const std::vector<std::string>& template_args,
124       std::string* result) const {
125     *result = "";
126     return absl::OkStatus();
127   }
GetGPUResources()128   virtual GPUResources GetGPUResources() const { return GPUResources(); }
129 
Release()130   virtual void Release() {}
131 
SetAccess(AccessType access_type)132   void SetAccess(AccessType access_type) { access_type_ = access_type; }
GetAccess()133   AccessType GetAccess() const { return access_type_; }
134 
135  protected:
136   friend flatbuffers::Offset<tflite::gpu::data::GPUObjectDescriptor> Encode(
137       const GPUObjectDescriptor& desc, flatbuffers::FlatBufferBuilder* builder);
138   friend void Decode(const tflite::gpu::data::GPUObjectDescriptor* fb_obj,
139                      GPUObjectDescriptor* obj);
140   mutable std::map<std::string, std::string> state_vars_;
141   AccessType access_type_;
142 };
143 
144 using GPUObjectDescriptorPtr = std::unique_ptr<GPUObjectDescriptor>;
145 
146 }  // namespace gpu
147 }  // namespace tflite
148 
149 #endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OBJECT_DESC_H_
150