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/memory_management.h"
17
18 #include <cstddef>
19 #include <numeric>
20 #include <utility>
21 #include <vector>
22
23 #include "tensorflow/lite/delegates/gpu/common/memory_management/equality_assignment.h"
24 #include "tensorflow/lite/delegates/gpu/common/memory_management/greedy_by_breadth_assignment.h"
25 #include "tensorflow/lite/delegates/gpu/common/memory_management/greedy_by_size_assignment.h"
26 #include "tensorflow/lite/delegates/gpu/common/memory_management/greedy_in_order_assignment.h"
27 #include "tensorflow/lite/delegates/gpu/common/memory_management/min_cost_flow_assignment.h"
28 #include "tensorflow/lite/delegates/gpu/common/memory_management/naive_assignment.h"
29 #include "tensorflow/lite/delegates/gpu/common/memory_management/types.h"
30 #include "tensorflow/lite/delegates/gpu/common/shape.h"
31 #include "tensorflow/lite/delegates/gpu/common/status.h"
32 #include "tensorflow/lite/delegates/gpu/common/types.h"
33
34 namespace tflite {
35 namespace gpu {
36 namespace {
37
TotalSize(const ObjectsAssignment<size_t> & assignment)38 size_t TotalSize(const ObjectsAssignment<size_t>& assignment) {
39 return std::accumulate(assignment.object_sizes.begin(),
40 assignment.object_sizes.end(), static_cast<size_t>(0));
41 }
42
43 } // namespace
44
ObjectsToOffsets(const ObjectsAssignment<size_t> & obj_assignment)45 OffsetsAssignment ObjectsToOffsets(
46 const ObjectsAssignment<size_t>& obj_assignment) {
47 size_t num_tensors = obj_assignment.object_ids.size();
48 size_t num_objects = obj_assignment.object_sizes.size();
49 OffsetsAssignment result = {/*offsets=*/std::vector<size_t>(num_tensors),
50 /*total_size=*/0};
51 std::vector<size_t> ids_to_offset(num_objects);
52 for (size_t i = 0; i < num_objects; ++i) {
53 ids_to_offset[i] = result.total_size;
54 result.total_size += obj_assignment.object_sizes[i];
55 }
56 for (size_t i = 0; i < num_tensors; ++i) {
57 result.offsets[i] = ids_to_offset[obj_assignment.object_ids[i]];
58 }
59 return result;
60 }
61
BestGreedy(const std::vector<TensorUsageRecord<size_t>> & usage_records,ObjectsAssignment<size_t> * assignment)62 absl::Status BestGreedy(
63 const std::vector<TensorUsageRecord<size_t>>& usage_records,
64 ObjectsAssignment<size_t>* assignment) {
65 RETURN_IF_ERROR(
66 GreedyBySizeDistPriorityAssignment(usage_records, assignment));
67 ObjectsAssignment<size_t> assignment_by_breadth;
68 if (GreedyByBreadthAssignment(usage_records, &assignment_by_breadth).ok() &&
69 TotalSize(assignment_by_breadth) < TotalSize(*assignment)) {
70 std::swap(*assignment, assignment_by_breadth);
71 }
72 return absl::OkStatus();
73 }
74
75 template <>
AssignObjectsToTensors(const std::vector<TensorUsageRecord<size_t>> & usage_records,MemoryStrategy strategy,ObjectsAssignment<size_t> * assignment,const UsageGraph * reallocation_graph)76 absl::Status AssignObjectsToTensors(
77 const std::vector<TensorUsageRecord<size_t>>& usage_records,
78 MemoryStrategy strategy, ObjectsAssignment<size_t>* assignment,
79 const UsageGraph* reallocation_graph) {
80 switch (strategy) {
81 case MemoryStrategy::NAIVE:
82 return NaiveAssignment(usage_records, assignment);
83 case MemoryStrategy::EQUALITY:
84 return EqualityAssignmentWithHash(usage_records, assignment);
85 case MemoryStrategy::GREEDY_IN_ORDER:
86 return GreedyInOrderAssignment(usage_records, assignment,
87 reallocation_graph);
88 case MemoryStrategy::GREEDY_BY_BREADTH:
89 return GreedyByBreadthAssignment(usage_records, assignment);
90 case MemoryStrategy::GREEDY_BY_SIZE:
91 return GreedyBySizeDistPriorityAssignment(usage_records, assignment);
92 case MemoryStrategy::GREEDY_BEST:
93 return BestGreedy(usage_records, assignment);
94 case MemoryStrategy::MINCOSTFLOW:
95 return MinCostFlowAssignment(usage_records, assignment);
96 default:
97 return absl::InternalError(
98 "MemoryStrategy is not supported with current tensor size type.");
99 }
100 return absl::OkStatus();
101 }
102
103 template <>
AssignObjectsToTensors(const std::vector<TensorUsageRecord<BHWC>> & usage_records,MemoryStrategy strategy,ObjectsAssignment<BHWC> * assignment,const UsageGraph * reallocation_graph)104 absl::Status AssignObjectsToTensors(
105 const std::vector<TensorUsageRecord<BHWC>>& usage_records,
106 MemoryStrategy strategy, ObjectsAssignment<BHWC>* assignment,
107 const UsageGraph* reallocation_graph) {
108 switch (strategy) {
109 case MemoryStrategy::NAIVE:
110 return NaiveAssignment(usage_records, assignment);
111 case MemoryStrategy::EQUALITY:
112 return EqualityAssignmentWithHash(usage_records, assignment);
113 default:
114 return absl::InternalError(
115 "MemoryStrategy is not supported with current tensor size type.");
116 }
117 return absl::OkStatus();
118 }
119
120 template <>
AssignObjectsToTensors(const std::vector<TensorUsageRecord<uint2>> & usage_records,MemoryStrategy strategy,ObjectsAssignment<uint2> * assignment,const UsageGraph * reallocation_graph)121 absl::Status AssignObjectsToTensors(
122 const std::vector<TensorUsageRecord<uint2>>& usage_records,
123 MemoryStrategy strategy, ObjectsAssignment<uint2>* assignment,
124 const UsageGraph* reallocation_graph) {
125 switch (strategy) {
126 case MemoryStrategy::NAIVE:
127 return NaiveAssignment(usage_records, assignment);
128 case MemoryStrategy::EQUALITY:
129 return EqualityAssignment(usage_records, assignment);
130 case MemoryStrategy::GREEDY_IN_ORDER:
131 return GreedyInOrderAssignmentMultidimensional(usage_records, assignment);
132 default:
133 return absl::InternalError(
134 "MemoryStrategy is not supported with current tensor size type.");
135 }
136 return absl::OkStatus();
137 }
138
139 template <>
AssignObjectsToTensors(const std::vector<TensorUsageRecord<uint3>> & usage_records,MemoryStrategy strategy,ObjectsAssignment<uint3> * assignment,const UsageGraph * reallocation_graph)140 absl::Status AssignObjectsToTensors(
141 const std::vector<TensorUsageRecord<uint3>>& usage_records,
142 MemoryStrategy strategy, ObjectsAssignment<uint3>* assignment,
143 const UsageGraph* reallocation_graph) {
144 switch (strategy) {
145 case MemoryStrategy::NAIVE:
146 return NaiveAssignment(usage_records, assignment);
147 case MemoryStrategy::EQUALITY:
148 return EqualityAssignment(usage_records, assignment);
149 case MemoryStrategy::GREEDY_IN_ORDER:
150 return GreedyInOrderAssignmentMultidimensional(usage_records, assignment);
151 default:
152 return absl::InternalError(
153 "MemoryStrategy is not supported with current tensor size type.");
154 }
155 return absl::OkStatus();
156 }
157
AssignOffsetsToTensors(const std::vector<TensorUsageRecord<size_t>> & usage_records,const MemoryStrategy & strategy,OffsetsAssignment * assignment,const UsageGraph * reallocation_graph)158 absl::Status AssignOffsetsToTensors(
159 const std::vector<TensorUsageRecord<size_t>>& usage_records,
160 const MemoryStrategy& strategy, OffsetsAssignment* assignment,
161 const UsageGraph* reallocation_graph) {
162 if (strategy == MemoryStrategy::GREEDY_BY_SIZE) {
163 return GreedyBySizeAssignment(usage_records, assignment);
164 }
165 ObjectsAssignment<size_t> objects_assignment;
166 RETURN_IF_ERROR(AssignObjectsToTensors(
167 usage_records, strategy, &objects_assignment, reallocation_graph));
168 *assignment = ObjectsToOffsets(objects_assignment);
169 return absl::OkStatus();
170 }
171
172 } // namespace gpu
173 } // namespace tflite
174