xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/gpu/gl/runtime.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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/gl/runtime.h"
17 
18 #include <algorithm>
19 #include <cstdint>
20 #include <functional>
21 #include <memory>
22 #include <utility>
23 #include <variant>
24 #include <vector>
25 
26 #include "absl/status/status.h"
27 #include "absl/strings/str_cat.h"
28 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
29 #include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
30 #include "tensorflow/lite/delegates/gpu/common/memory_management.h"
31 #include "tensorflow/lite/delegates/gpu/common/memory_management/types.h"
32 #include "tensorflow/lite/delegates/gpu/common/status.h"
33 #include "tensorflow/lite/delegates/gpu/common/types.h"
34 #include "tensorflow/lite/delegates/gpu/gl/gl_call.h"
35 #include "tensorflow/lite/delegates/gpu/gl/gl_errors.h"
36 #include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
37 #include "tensorflow/lite/delegates/gpu/gl/gl_texture.h"
38 #include "tensorflow/lite/delegates/gpu/gl/object.h"
39 #include "tensorflow/lite/delegates/gpu/gl/portable_gl31.h"
40 #include "tensorflow/lite/delegates/gpu/gl/variable.h"
41 
42 namespace tflite {
43 namespace gpu {
44 namespace gl {
45 namespace {
46 
47 struct TextureF16Maker {
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureF16Maker48   absl::Status operator()(const uint3& size) const {
49     return CreateReadOnlyImageTextureF16(size, data, gl_texture);
50   }
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureF16Maker51   absl::Status operator()(const uint2& size) const {
52     return CreateReadOnlyImageTextureF16(size, data, gl_texture);
53   }
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureF16Maker54   absl::Status operator()(const size_t& size) const {
55     return CreateReadOnlyImageTextureF16(uint2(static_cast<uint32_t>(size), 1U),
56                                          data, gl_texture);
57   }
58   absl::Span<const uint16_t> data;
59   GlTexture* gl_texture;
60 };
61 
62 struct TextureF32Maker {
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureF32Maker63   absl::Status operator()(const uint3& size) const {
64     return CreateReadOnlyImageTexture(size, data, gl_texture);
65   }
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureF32Maker66   absl::Status operator()(const uint2& size) const {
67     return CreateReadOnlyImageTexture(size, data, gl_texture);
68   }
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureF32Maker69   absl::Status operator()(const size_t& size) const {
70     return CreateReadOnlyImageTexture(uint2(static_cast<uint32_t>(size), 1U),
71                                       data, gl_texture);
72   }
73   absl::Span<const float> data;
74   GlTexture* gl_texture;
75 };
76 
MakeGlTexture(const Object & object,const ObjectData & data,GlTexture * gl_texture)77 absl::Status MakeGlTexture(const Object& object, const ObjectData& data,
78                            GlTexture* gl_texture) {
79   if (object.access == AccessType::READ_WRITE ||
80       object.access == AccessType::WRITE) {
81     return absl::InvalidArgumentError("Read-write textures are not supported");
82   }
83   if (object.data_type != DataType::FLOAT16 &&
84       object.data_type != DataType::FLOAT32) {
85     return absl::InvalidArgumentError(
86         "Textures support float16 or float32 only.");
87   }
88   switch (object.data_type) {
89     case DataType::FLOAT16: {
90       if (data.size() % 2 != 0) {
91         return absl::InvalidArgumentError("Texture size is not aligned");
92       }
93       return std::visit(
94           TextureF16Maker{
95               .data = absl::MakeConstSpan(
96                   reinterpret_cast<const uint16_t*>(data.data()),
97                   data.size() / 2),
98               .gl_texture = gl_texture,
99           },
100           object.size);
101     }
102     case DataType::FLOAT32: {
103       if (data.size() % sizeof(float) != 0) {
104         return absl::InvalidArgumentError("Texture size is not aligned");
105       }
106       return std::visit(
107           TextureF32Maker{
108               .data = absl::MakeConstSpan(
109                   reinterpret_cast<const float*>(data.data()),
110                   data.size() / sizeof(float)),
111               .gl_texture = gl_texture,
112           },
113           object.size);
114     }
115     default:
116       return absl::InvalidArgumentError("Unsupported textures data type.");
117   }
118 }
119 
120 struct TextureRefMaker {
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureRefMaker121   absl::Status operator()(const uint3& size) const {
122     return CreateReadWriteRgbaImageTexture(type, size, gl_texture);
123   }
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureRefMaker124   absl::Status operator()(const uint2& size) const {
125     return CreateReadWriteRgbaImageTexture(type, size, gl_texture);
126   }
operator ()tflite::gpu::gl::__anon0bd63f3a0111::TextureRefMaker127   absl::Status operator()(const size_t& size) const {
128     return CreateReadWriteRgbaImageTexture(
129         type, uint2(static_cast<uint32_t>(size), 1U), gl_texture);
130   }
131   DataType type;
132   GlTexture* gl_texture;
133 };
134 
135 // Makes read-write gl texture
MakeGlTextureRef(const Object & object,GlTexture * gl_texture)136 absl::Status MakeGlTextureRef(const Object& object, GlTexture* gl_texture) {
137   return std::visit(TextureRefMaker{object.data_type, gl_texture}, object.size);
138 }
139 
MakeGlBuffer(const Object & object,const ObjectData & data,GlBuffer * gl_buffer)140 absl::Status MakeGlBuffer(const Object& object, const ObjectData& data,
141                           GlBuffer* gl_buffer) {
142   if (data.size() % SizeOf(object.data_type) != 0) {
143     return absl::InvalidArgumentError("Buffer size is not aligned");
144   }
145   return CreateReadOnlyShaderStorageBuffer(absl::MakeConstSpan(data),
146                                            gl_buffer);
147 }
148 
MakeBindingFunc(const Object & object,uint32_t id,const ObjectManager * objects,std::function<absl::Status ()> * binding_func)149 absl::Status MakeBindingFunc(const Object& object, uint32_t id,
150                              const ObjectManager* objects,
151                              std::function<absl::Status()>* binding_func) {
152   const uint32_t binding = object.binding;
153   switch (object.object_type) {
154     case ObjectType::BUFFER: {
155       auto ptr = objects->FindBuffer(id);
156       if (!ptr) {
157         return absl::NotFoundError(
158             absl::StrCat("Buffer ", id, " is not found"));
159       }
160       size_t size_in_bytes = ByteSizeOf(object);
161       // TODO(akulik): make comparison != instead of <
162       if (ptr->bytes_size() < size_in_bytes) {
163         return absl::FailedPreconditionError(
164             absl::StrCat("Buffer ", id, " size in bytes ", ptr->bytes_size(),
165                          " < requested size_in_bytes ", size_in_bytes));
166       }
167       *binding_func = [=]() { return ptr->BindToIndex(binding); };
168       break;
169     }
170     case ObjectType::TEXTURE: {
171       auto ptr = objects->FindTexture(id);
172       if (!ptr) {
173         return absl::NotFoundError(
174             absl::StrCat("Texture ", id, " is not found"));
175       }
176       *binding_func = [=]() { return ptr->BindAsReadWriteImage(binding); };
177       break;
178     }
179     case ObjectType::UNKNOWN:
180       return absl::InvalidArgumentError("Unknown object type");
181   }
182   return absl::OkStatus();
183 }
184 
185 // TODO(b/147771327): think about merging this function with MakeBindingFunc()
MakeLateBindingFunc(const Object & object,uint32_t id,const ObjectManager * objects,std::function<absl::Status ()> * binding_func)186 absl::Status MakeLateBindingFunc(const Object& object, uint32_t id,
187                                  const ObjectManager* objects,
188                                  std::function<absl::Status()>* binding_func) {
189   const uint32_t binding = object.binding;
190   switch (object.object_type) {
191     case ObjectType::BUFFER: {
192       auto ptr = objects->FindBuffer(id);
193       if (!ptr) {
194         return absl::NotFoundError(
195             absl::StrCat("Buffer ", id, " is not found"));
196       }
197       *binding_func = [=]() {
198         auto ptr = objects->FindBuffer(id);
199         if (!ptr) {
200           return absl::NotFoundError(
201               absl::StrCat("Buffer ", id, " is not found"));
202         }
203         if (!ptr->is_valid()) {
204           return absl::InvalidArgumentError("Buffer is not initialized.");
205         }
206         size_t size_in_bytes = ByteSizeOf(object);
207         if (ptr->bytes_size() < size_in_bytes) {
208           return absl::FailedPreconditionError(
209               absl::StrCat("Buffer ", id, " size in bytes ", ptr->bytes_size(),
210                            " < requested size_in_bytes ", size_in_bytes));
211         }
212         return ptr->BindToIndex(binding);
213       };
214       break;
215     }
216     case ObjectType::TEXTURE: {
217       auto ptr = objects->FindTexture(id);
218       if (!ptr) {
219         return absl::NotFoundError(
220             absl::StrCat("Texture ", id, " is not found"));
221       }
222       *binding_func = [=]() {
223         auto ptr = objects->FindTexture(id);
224         if (!ptr) {
225           return absl::NotFoundError(
226               absl::StrCat("Texture ", id, " is not found"));
227         }
228         if (!ptr->is_valid()) {
229           return absl::InvalidArgumentError("Texture is not initialized.");
230         }
231         return ptr->BindAsReadWriteImage(binding);
232       };
233       break;
234     }
235     case ObjectType::UNKNOWN:
236       return absl::InvalidArgumentError("Unknown object type");
237   }
238   return absl::OkStatus();
239 }
240 
241 }  // namespace
242 
Runtime(const RuntimeOptions & options,const GpuInfo & gpu_info,CommandQueue * command_queue,const ObjectManager * external_objects)243 Runtime::Runtime(const RuntimeOptions& options, const GpuInfo& gpu_info,
244                  CommandQueue* command_queue,
245                  const ObjectManager* external_objects)
246     : options_(options),
247       gpu_info_(gpu_info),
248       external_objects_(external_objects),
249       command_queue_(command_queue) {
250   programs_.reserve(256);
251   if (options_.bundle_readonly_objects) {
252     shared_readonly_buffer_ = std::make_unique<SharedBufferData>();
253   }
254 }
255 
AddProgram(const GlShader & shader,const std::vector<Variable> & parameters,const std::vector<Object> & objects,const uint3 & num_workgroups)256 absl::Status Runtime::AddProgram(const GlShader& shader,
257                                  const std::vector<Variable>& parameters,
258                                  const std::vector<Object>& objects,
259                                  const uint3& num_workgroups) {
260   GlProgram program;
261   RETURN_IF_ERROR(GlProgram::CreateWithShader(shader, &program));
262 
263   for (auto& parameter : parameters) {
264     RETURN_IF_ERROR(program.SetParameter(parameter));
265   }
266 
267   programs_.emplace_back(
268       CompiledProgramDescriptor{std::move(program), num_workgroups, {}});
269 
270   // Create const buffers, resolve external references and collect internal
271   // buffer references.
272   for (auto& object : objects) {
273     auto& program = programs_.back();
274     BindFunc binding_func;
275     if (IsRef(object)) {
276       // Reference object could be provided externally as a model input/output
277       // but also for debugging purposes. Otherwise all references are collected
278       // and allocated later.
279       absl::Status status = MakeLateBindingFunc(
280           object, GetRef(object), external_objects_, &binding_func);
281       if (!status.ok()) {
282         if (absl::IsNotFound(status)) {
283           program.refs.push_back(object);
284           continue;  // don't add to binding.
285         }
286         return status;
287       }
288     } else {
289       // Allocate const object.
290       uint32_t id;
291       RETURN_IF_ERROR(AllocateConstObject(object, &id));
292       RETURN_IF_ERROR(
293           MakeBindingFunc(object, id, &const_objects_, &binding_func));
294     }
295     program.bindings.push_back(std::move(binding_func));
296   }
297 
298   // All parameters once set stay with program, therefore, we only need to keep
299   // program and bindings for execution.
300   return absl::OkStatus();
301 }
302 
AllocateInternalObject(const Object & object)303 absl::Status Runtime::AllocateInternalObject(const Object& object) {
304   const ObjectRef ref = GetRef(object);
305   switch (object.object_type) {
306     case ObjectType::BUFFER: {
307       GlBuffer gl_buffer;
308       RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<uint8_t>(
309           ByteSizeOf(object), &gl_buffer));
310       RETURN_IF_ERROR(
311           internal_objects_.RegisterBuffer(ref, std::move(gl_buffer)));
312       break;
313     }
314     case ObjectType::TEXTURE: {
315       GlTexture gl_texture;
316       RETURN_IF_ERROR(MakeGlTextureRef(object, &gl_texture));
317       RETURN_IF_ERROR(
318           internal_objects_.RegisterTexture(ref, std::move(gl_texture)));
319       break;
320     }
321     default:
322       return absl::InternalError("Unexpected internal object type");
323   }
324   return absl::OkStatus();
325 }
326 
AllocateConstObject(const Object & object,uint32_t * id)327 absl::Status Runtime::AllocateConstObject(const Object& object, uint32_t* id) {
328   const ObjectData* data = GetData(object);
329   if (data == nullptr) {
330     return absl::InternalError(
331         "Unable to allocate reference as a const object");
332   }
333   *id = next_const_id_++;
334   switch (object.object_type) {
335     case ObjectType::BUFFER: {
336       GlBuffer gl_buffer;
337       if (!shared_readonly_buffer_ ||
338           !shared_readonly_buffer_->Add(*data, &gl_buffer)) {
339         RETURN_IF_ERROR(MakeGlBuffer(object, *data, &gl_buffer));
340       }
341       RETURN_IF_ERROR(const_objects_.RegisterBuffer(*id, std::move(gl_buffer)));
342       break;
343     }
344     case ObjectType::TEXTURE: {
345       GlTexture gl_texture;
346       RETURN_IF_ERROR(MakeGlTexture(object, *data, &gl_texture));
347       RETURN_IF_ERROR(
348           const_objects_.RegisterTexture(*id, std::move(gl_texture)));
349       break;
350     }
351     case ObjectType::UNKNOWN:
352       return absl::InternalError("Unknown object type");
353   }
354   return absl::OkStatus();
355 }
356 
PrepareForExecution()357 absl::Status Runtime::PrepareForExecution() {
358   if (shared_readonly_buffer_ && !shared_readonly_buffer_->empty()) {
359     GlBuffer shared_buffer;
360     RETURN_IF_ERROR(
361         shared_readonly_buffer_->CreateSharedGlBuffer(&shared_buffer));
362     shared_readonly_buffer_.reset(nullptr);
363     RETURN_IF_ERROR(const_objects_.RegisterBuffer(next_const_id_++,
364                                                   std::move(shared_buffer)));
365   }
366 
367   if (options_.reuse_internal_objects) {
368     // Analyze internal objects and make a pool of shared objects to be re-used
369     // by them. These shared objects need to be allocated upfront.
370     std::vector<Object> shared_objects;
371     RETURN_IF_ERROR(AssignInternalObjects(&shared_objects));
372     for (const Object& object : shared_objects) {
373       RETURN_IF_ERROR(AllocateInternalObject(object));
374     }
375   }
376 
377   // Allocate all internal objects and create bindings for them.
378   for (auto& program : programs_) {
379     for (auto& object : program.refs) {
380       // Check whether it is created already.
381       BindFunc binding;
382       ObjectRef ref = GetRef(object);
383       absl::Status status =
384           MakeBindingFunc(object, ref, &internal_objects_, &binding);
385       if (!status.ok()) {
386         if (absl::IsNotFound(status)) {
387           RETURN_IF_ERROR(AllocateInternalObject(object));
388           RETURN_IF_ERROR(
389               MakeBindingFunc(object, ref, &internal_objects_, &binding));
390         } else {
391           return status;
392         }
393       }
394       program.bindings.push_back(std::move(binding));
395     }
396     program.refs.clear();
397   }
398   return absl::OkStatus();
399 }
400 
401 namespace {
402 
403 const size_t kNotAssigned = std::numeric_limits<size_t>::max();
404 
405 struct CombinedUsageRecords {
406   std::vector<TensorUsageRecord<size_t>> buffers;
407   std::vector<TensorUsageRecord<size_t>> textures_1d;
408   std::vector<TensorUsageRecord<uint2>> textures_2d;
409   std::vector<TensorUsageRecord<uint3>> textures_3d;
410   std::vector<size_t> usage_refs;
411 };
412 
413 template <typename TensorSizeT>
UpdateUsageRecord(TensorUsageRecord<TensorSizeT> * usage_rec,size_t task_id)414 void UpdateUsageRecord(TensorUsageRecord<TensorSizeT>* usage_rec,
415                        size_t task_id) {
416   usage_rec->first_task = std::min(usage_rec->first_task, task_id);
417   usage_rec->last_task = std::max(usage_rec->last_task, task_id);
418 }
419 
420 struct AddUsageRecordForTextureFunc {
operator ()tflite::gpu::gl::__anon0bd63f3a0611::AddUsageRecordForTextureFunc421   void operator()(const uint3& size) const {
422     auto& usage_ref = usage_records->usage_refs[object_ref];
423     if (usage_ref == kNotAssigned) {
424       usage_ref = usage_records->textures_3d.size();
425       usage_records->textures_3d.emplace_back(/*tensor_size=*/size,
426                                               /*first_task=*/program_id,
427                                               /*last_task=*/program_id);
428     } else {
429       UpdateUsageRecord(&usage_records->textures_3d[usage_ref], program_id);
430     }
431   }
432 
operator ()tflite::gpu::gl::__anon0bd63f3a0611::AddUsageRecordForTextureFunc433   void operator()(const uint2& size) const {
434     auto& usage_ref = usage_records->usage_refs[object_ref];
435     if (usage_ref == kNotAssigned) {
436       usage_ref = usage_records->textures_2d.size();
437       usage_records->textures_2d.emplace_back(/*tensor_size=*/size,
438                                               /*first_task=*/program_id,
439                                               /*last_task=*/program_id);
440     } else {
441       UpdateUsageRecord(&usage_records->textures_2d[usage_ref], program_id);
442     }
443   }
444 
operator ()tflite::gpu::gl::__anon0bd63f3a0611::AddUsageRecordForTextureFunc445   void operator()(size_t size) const {
446     auto& usage_ref = usage_records->usage_refs[object_ref];
447     if (usage_ref == kNotAssigned) {
448       usage_ref = usage_records->textures_1d.size();
449       usage_records->textures_1d.emplace_back(/*tensor_size=*/size,
450                                               /*first_task=*/program_id,
451                                               /*last_task=*/program_id);
452     } else {
453       UpdateUsageRecord(&usage_records->textures_1d[usage_ref], program_id);
454     }
455   }
456 
457   CombinedUsageRecords* usage_records;
458   const ObjectRef& object_ref;
459   const size_t program_id;
460 };
461 
462 // We assume that AddUsageRecord for different objects is called in order of
463 // program_id.
AddUsageRecord(CombinedUsageRecords * usage_records,const Object & object,const size_t program_id)464 absl::Status AddUsageRecord(CombinedUsageRecords* usage_records,
465                             const Object& object, const size_t program_id) {
466   auto ref = GetRef(object);
467   if (ref >= usage_records->usage_refs.size()) {
468     usage_records->usage_refs.resize(ref + 1, kNotAssigned);
469   }
470   auto& usage_ref = usage_records->usage_refs[ref];
471   if (object.object_type == ObjectType::BUFFER) {
472     if (usage_ref == kNotAssigned) {
473       usage_ref = usage_records->buffers.size();
474       usage_records->buffers.emplace_back(
475           /*tensor_size=*/NumElements(object.size),
476           /*first_task=*/program_id,
477           /*last_task=*/program_id);
478     } else {
479       UpdateUsageRecord(&usage_records->buffers[usage_ref], program_id);
480     }
481     return absl::OkStatus();
482   }
483   if (object.object_type == ObjectType::TEXTURE) {
484     std::visit(AddUsageRecordForTextureFunc{usage_records, ref, program_id},
485                object.size);
486     return absl::OkStatus();
487   }
488   return absl::InternalError("Unexpected object type");
489 }
490 
ApplyBuffersAssignment(const ObjectsAssignment<size_t> & assignment,const std::vector<size_t> & global_ref_to_usage_rec,const std::vector<Object * > & global_ref_to_object_ptr,std::vector<ObjectRef> * global_ref_to_shared_ref,std::vector<Object> * shared_objects)491 absl::Status ApplyBuffersAssignment(
492     const ObjectsAssignment<size_t>& assignment,
493     const std::vector<size_t>& global_ref_to_usage_rec,
494     const std::vector<Object*>& global_ref_to_object_ptr,
495     std::vector<ObjectRef>* global_ref_to_shared_ref,
496     std::vector<Object>* shared_objects) {
497   std::vector<ObjectRef> assigned_id_to_shared_ref(
498       assignment.object_sizes.size(), kInvalidObjectRef);
499   for (size_t global_ref = 0; global_ref < global_ref_to_usage_rec.size();
500        ++global_ref) {
501     const auto& usage_rec_id = global_ref_to_usage_rec[global_ref];
502     Object* object = global_ref_to_object_ptr[global_ref];
503     if (usage_rec_id == kNotAssigned || object == nullptr ||
504         object->object_type != ObjectType::BUFFER) {
505       // Skip objects with other data type and non-buffers.
506       continue;
507     }
508 
509     // id of shared object, returned by memory allocation algorithm.
510     size_t assigned_id = assignment.object_ids[usage_rec_id];
511 
512     // id of corresponding shared object in vector share_objects.
513     ObjectRef shared_ref = assigned_id_to_shared_ref[assigned_id];
514 
515     if (shared_ref == kInvalidObjectRef) {
516       // We need to create new shared object for current buffer.
517       shared_ref = shared_objects->size();
518       Object shared_object = *object;
519       shared_object.access = AccessType::READ_WRITE;
520       shared_object.object = shared_ref;
521       shared_object.size = assignment.object_sizes[assigned_id];
522       shared_objects->push_back(std::move(shared_object));
523       assigned_id_to_shared_ref[assigned_id] = shared_ref;
524     }
525     (*global_ref_to_shared_ref)[global_ref] = shared_ref;
526   }
527   return absl::OkStatus();
528 }
529 
530 template <typename ObjectSizeT>
ApplyTexturesAssignment(const ObjectsAssignment<ObjectSizeT> & assignment,const std::vector<size_t> & global_ref_to_usage_rec,const std::vector<Object * > & global_ref_to_object_ptr,std::vector<ObjectRef> * global_ref_to_shared_ref,std::vector<Object> * shared_objects)531 absl::Status ApplyTexturesAssignment(
532     const ObjectsAssignment<ObjectSizeT>& assignment,
533     const std::vector<size_t>& global_ref_to_usage_rec,
534     const std::vector<Object*>& global_ref_to_object_ptr,
535     std::vector<ObjectRef>* global_ref_to_shared_ref,
536     std::vector<Object>* shared_objects) {
537   std::vector<ObjectRef> assigned_id_to_shared_ref(
538       assignment.object_sizes.size(), kInvalidObjectRef);
539   for (size_t global_ref = 0; global_ref < global_ref_to_usage_rec.size();
540        ++global_ref) {
541     const auto& usage_rec_id = global_ref_to_usage_rec[global_ref];
542     Object* object = global_ref_to_object_ptr[global_ref];
543     if (usage_rec_id == kNotAssigned || object == nullptr ||
544         object->object_type != ObjectType::TEXTURE ||
545         !std::holds_alternative<ObjectSizeT>(object->size)) {
546       // Skip objects with other data type, non-textures and textures with wrong
547       // number of dimensions.
548       continue;
549     }
550 
551     // id of shared object, returned by memory allocation algorithm.
552     size_t assigned_id = assignment.object_ids[usage_rec_id];
553 
554     // id of corresponding shared object in vector share_objects.
555     ObjectRef shared_ref = assigned_id_to_shared_ref[assigned_id];
556 
557     if (shared_ref == kInvalidObjectRef) {
558       // We need to create new shared object for current texture.
559       shared_ref = shared_objects->size();
560       Object shared_object = *object;
561       shared_object.access = AccessType::READ_WRITE;
562       shared_object.object = shared_ref;
563       shared_object.size = assignment.object_sizes[assigned_id];
564       shared_objects->push_back(std::move(shared_object));
565       assigned_id_to_shared_ref[assigned_id] = shared_ref;
566     }
567     (*global_ref_to_shared_ref)[global_ref] = shared_ref;
568   }
569   return absl::OkStatus();
570 }
571 
572 }  // namespace
573 
574 // Assign shared objects to internal objects, using memory allocation
575 // algorithms. Usage records for the algorithms are calculated separately for
576 // each data type and object type.
AssignInternalObjects(std::vector<Object> * shared_objects)577 absl::Status Runtime::AssignInternalObjects(
578     std::vector<Object>* shared_objects) {
579   // Build tensor usage records, clusterized by object type and data type.
580   std::map<DataType, CombinedUsageRecords> usage_records_by_data_type;
581   std::vector<Object*> global_ref_to_object_ptr;
582   for (size_t i = 0; i < programs_.size(); ++i) {
583     for (auto& object : programs_[i].refs) {
584       auto ref = GetRef(object);
585       if (ref >= global_ref_to_object_ptr.size()) {
586         global_ref_to_object_ptr.resize(ref + 1, nullptr);
587       }
588       if (global_ref_to_object_ptr[ref] == nullptr) {
589         global_ref_to_object_ptr[ref] = &object;
590       }
591       RETURN_IF_ERROR(AddUsageRecord(
592           &usage_records_by_data_type[object.data_type], object, i));
593     }
594   }
595 
596   std::vector<ObjectRef> global_ref_to_shared_ref(
597       global_ref_to_object_ptr.size(), kInvalidObjectRef);
598 
599   // Calculate and apply shared objects assignment for each data type.
600   for (const auto& it : usage_records_by_data_type) {
601     const CombinedUsageRecords& usage_records = it.second;
602     if (!usage_records.buffers.empty()) {
603       ObjectsAssignment<size_t> buffer_assignment;
604       RETURN_IF_ERROR(AssignObjectsToTensors(usage_records.buffers,
605                                              MemoryStrategy::GREEDY_BEST,
606                                              &buffer_assignment));
607       RETURN_IF_ERROR(ApplyBuffersAssignment(
608           buffer_assignment, usage_records.usage_refs, global_ref_to_object_ptr,
609           &global_ref_to_shared_ref, shared_objects));
610     }
611     if (!usage_records.textures_1d.empty()) {
612       ObjectsAssignment<size_t> texture_1d_assignment;
613       RETURN_IF_ERROR(AssignObjectsToTensors(usage_records.textures_1d,
614                                              MemoryStrategy::GREEDY_BEST,
615                                              &texture_1d_assignment));
616       RETURN_IF_ERROR(ApplyTexturesAssignment(
617           texture_1d_assignment, usage_records.usage_refs,
618           global_ref_to_object_ptr, &global_ref_to_shared_ref, shared_objects));
619     }
620     if (!usage_records.textures_2d.empty()) {
621       ObjectsAssignment<uint2> texture_2d_assignment;
622       RETURN_IF_ERROR(AssignObjectsToTensors(usage_records.textures_2d,
623                                              MemoryStrategy::GREEDY_IN_ORDER,
624                                              &texture_2d_assignment));
625       RETURN_IF_ERROR(ApplyTexturesAssignment(
626           texture_2d_assignment, usage_records.usage_refs,
627           global_ref_to_object_ptr, &global_ref_to_shared_ref, shared_objects));
628     }
629     if (!usage_records.textures_3d.empty()) {
630       ObjectsAssignment<uint3> texture_3d_assignment;
631       RETURN_IF_ERROR(AssignObjectsToTensors(usage_records.textures_3d,
632                                              MemoryStrategy::GREEDY_IN_ORDER,
633                                              &texture_3d_assignment));
634       RETURN_IF_ERROR(ApplyTexturesAssignment(
635           texture_3d_assignment, usage_records.usage_refs,
636           global_ref_to_object_ptr, &global_ref_to_shared_ref, shared_objects));
637     }
638   }
639 
640   for (size_t i = 0; i < programs_.size(); ++i) {
641     for (auto& object : programs_[i].refs) {
642       object.object = global_ref_to_shared_ref[GetRef(object)];
643     }
644   }
645   return absl::OkStatus();
646 }
647 
Execute()648 absl::Status Runtime::Execute() {
649   for (const auto& descriptor : programs_) {
650     for (auto& b : descriptor.bindings) {
651       RETURN_IF_ERROR(b());
652     }
653     RETURN_IF_ERROR(command_queue_->Dispatch(descriptor.program,
654                                              descriptor.num_workgroups));
655   }
656   return absl::OkStatus();
657 }
658 
659 }  // namespace gl
660 }  // namespace gpu
661 }  // namespace tflite
662