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/kernels/test_util.h"
17
18 #include <memory>
19 #include <string>
20 #include <vector>
21
22 #include "absl/container/flat_hash_map.h"
23 #include "absl/container/flat_hash_set.h"
24 #include "tensorflow/lite/delegates/gpu/common/model.h"
25 #include "tensorflow/lite/delegates/gpu/common/operations.h"
26 #include "tensorflow/lite/delegates/gpu/common/status.h"
27 #include "tensorflow/lite/delegates/gpu/common/tensor.h"
28 #include "tensorflow/lite/delegates/gpu/gl/api.h"
29 #include "tensorflow/lite/delegates/gpu/gl/egl_environment.h"
30 #include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
31 #include "tensorflow/lite/delegates/gpu/gl/object_manager.h"
32 #include "tensorflow/lite/delegates/gpu/gl/request_gpu_info.h"
33 #include "tensorflow/lite/delegates/gpu/gl/workgroups/default_calculator.h"
34
35 namespace tflite {
36 namespace gpu {
37 namespace gl {
38
SingleOpModel(Operation && operation,const std::vector<TensorRef<BHWC>> & inputs,const std::vector<TensorRef<BHWC>> & outputs)39 SingleOpModel::SingleOpModel(Operation&& operation,
40 const std::vector<TensorRef<BHWC>>& inputs,
41 const std::vector<TensorRef<BHWC>>& outputs) {
42 auto node = graph_.NewNode();
43 node->operation = std::move(operation);
44
45 for (int i = 0; i < inputs.size(); ++i) {
46 auto input = graph_.NewValue();
47 input->tensor = inputs[i];
48 graph_.AddConsumer(node->id, input->id).IgnoreError();
49 TensorFloat32 tensor;
50 tensor.id = input->tensor.ref;
51 tensor.shape = input->tensor.shape;
52 inputs_.emplace_back(std::move(tensor));
53 }
54
55 for (int i = 0; i < outputs.size(); ++i) {
56 auto output = graph_.NewValue();
57 output->tensor = outputs[i];
58 graph_.SetProducer(node->id, output->id).IgnoreError();
59 }
60 }
61
PopulateTensor(int index,std::vector<float> && data)62 bool SingleOpModel::PopulateTensor(int index, std::vector<float>&& data) {
63 if (index >= inputs_.size() ||
64 inputs_[index].shape.DimensionsProduct() != data.size()) {
65 return false;
66 }
67 inputs_[index].data = std::move(data);
68 return true;
69 }
70
Invoke(const CompilationOptions & compile_options,const RuntimeOptions & runtime_options,const NodeShader & shader)71 absl::Status SingleOpModel::Invoke(const CompilationOptions& compile_options,
72 const RuntimeOptions& runtime_options,
73 const NodeShader& shader) {
74 std::unique_ptr<EglEnvironment> env;
75 RETURN_IF_ERROR(EglEnvironment::NewEglEnvironment(&env));
76
77 ObjectManager objects;
78
79 // Create buffers for input tensors.
80 {
81 absl::flat_hash_map<int, uint32_t> tensor_to_id;
82 for (const auto* input : graph_.inputs()) {
83 tensor_to_id[input->tensor.ref] = input->id;
84 }
85 for (const auto& input : inputs_) {
86 GlBuffer buffer;
87 RETURN_IF_ERROR(CreatePHWC4BufferFromTensor(input, &buffer));
88 RETURN_IF_ERROR(
89 objects.RegisterBuffer(tensor_to_id[input.id], std::move(buffer)));
90 }
91 }
92
93 // Create buffers for output tensors.
94 for (const auto* output : graph_.outputs()) {
95 GlBuffer buffer;
96 RETURN_IF_ERROR(CreatePHWC4BufferFromTensorRef(output->tensor, &buffer));
97 RETURN_IF_ERROR(objects.RegisterBuffer(output->id, std::move(buffer)));
98 }
99
100 // Compile model.
101 GpuInfo gpu_info;
102 RETURN_IF_ERROR(RequestGpuInfo(&gpu_info));
103 std::unique_ptr<CompiledModel> compiled_model;
104 RETURN_IF_ERROR(Compile(compile_options, graph_, /*tflite_graph_io=*/{},
105 shader, *NewDefaultWorkgroupsCalculator(gpu_info),
106 &compiled_model));
107
108 // Get inference context.
109 auto command_queue = NewCommandQueue(gpu_info);
110 std::unique_ptr<InferenceContext> inference_context;
111 RETURN_IF_ERROR(compiled_model->NewRun(
112 runtime_options, &objects, command_queue.get(), &inference_context));
113 RETURN_IF_ERROR(inference_context->Reset());
114
115 // Run inference.
116 RETURN_IF_ERROR(inference_context->Execute());
117
118 // Copy output tensors to `output_`.
119 for (const auto* output : graph_.outputs()) {
120 TensorFloat32 tensor;
121 tensor.id = output->tensor.ref;
122 tensor.shape = output->tensor.shape;
123 tensor.data.reserve(output->tensor.shape.DimensionsProduct());
124 RETURN_IF_ERROR(
125 CopyFromPHWC4Buffer(*objects.FindBuffer(output->id), &tensor));
126 outputs_.push_back(std::move(tensor));
127 }
128 return absl::OkStatus();
129 }
130
Invoke(const NodeShader & shader)131 absl::Status SingleOpModel::Invoke(const NodeShader& shader) {
132 return Invoke(CompilationOptions(), RuntimeOptions(), shader);
133 }
134
135 } // namespace gl
136 } // namespace gpu
137 } // namespace tflite
138