1 | /* Copyright 2018 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 | #include "tensorflow/lite/graph_info.h" |
16 | |
17 | #include <algorithm> |
18 | #include <vector> |
19 | |
20 | #include "tensorflow/lite/c/common.h" |
21 | #include "tensorflow/lite/context_util.h" |
22 | |
23 | namespace tflite { |
24 | namespace { |
25 | |
26 | template <class T> |
27 | void Uniquefy(std::vector<T>* items) { |
28 | std::sort(items->begin(), items->end()); |
29 | items->erase(std::unique(items->begin(), items->end()), items->end()); |
30 | } |
31 | |
32 | // Helper class that actually performs partitioning by node sub set. |
33 | // Outputs to a provided `NodeSubset` structure. |
34 | // |
35 | // Example usage: |
36 | // PartitionGraphIntoIndependentNodeSubsetsImpl partitioner( |
37 | // info, nodes_to_part, control_edges, node_subsets); |
38 | // partitioner.Partition(); |
39 | // |
40 | // NOTE: Changing the partitioning logic would require a change to |
41 | // FP16GraphPartitionHelper. |
42 | // LINT.IfChange |
43 | class PartitionGraphIntoIndependentNodeSubsetsImpl { |
44 | public: |
45 | PartitionGraphIntoIndependentNodeSubsetsImpl( |
46 | const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, |
47 | const ControlEdges& control_edges, std::vector<NodeSubset>* node_subsets) |
48 | : info_(info), |
49 | node_subsets_(node_subsets), |
50 | node_type_(info_->num_total_nodes(), NodeSubset::kTfNonPartition), |
51 | control_edges_(control_edges), |
52 | num_incoming_control_edges_(info_->num_execution_nodes(), 0) { |
53 | // Populate the node_type_ map. |
54 | for (auto node_index : TfLiteIntArrayView(nodes_to_partition)) { |
55 | node_type_[node_index] = NodeSubset::kTfPartition; |
56 | } |
57 | Uniquefy(&control_edges_); |
58 | } |
59 | |
60 | // Actually partition the graph. |
61 | void Partition() { |
62 | // Initialize here to make Partition() re-entrant. |
63 | node_subsets_->clear(); |
64 | tensor_epochs_.clear(); |
65 | tensor_epochs_.resize(info_->num_tensors(), kEpochAlwaysReady); |
66 | node_epochs_.clear(); |
67 | node_epochs_.resize(info_->num_execution_nodes(), kEpochNotReady); |
68 | num_incoming_control_edges_.clear(); |
69 | num_incoming_control_edges_.resize(info_->num_execution_nodes(), 0); |
70 | for (const auto& edge : control_edges_) { |
71 | ++num_incoming_control_edges_[edge.second]; |
72 | } |
73 | |
74 | // Set computed tensors to be kEpochNotReady (initializer set everything to |
75 | // AlwaysReady). |
76 | for (int node_index = 0; node_index < info_->num_execution_nodes(); |
77 | node_index++) { |
78 | const TfLiteNode& node = info_->node(node_index); |
79 | for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) { |
80 | tensor_epochs_[output_tensor_index] = kEpochNotReady; |
81 | } |
82 | } |
83 | |
84 | // Do a graph traversal where each iteration in the loop is an epoch |
85 | // that corresponds to a node sub set that only contains nodes that are of |
86 | // the same node_type_. |
87 | while (true) { |
88 | BuildNodeSubset(); |
89 | if (node_subsets_->back().nodes.empty()) { |
90 | node_subsets_->pop_back(); |
91 | break; |
92 | } |
93 | } |
94 | |
95 | // Mark model outputs as node sub set outputs. All the rest have already |
96 | // been identified. |
97 | for (int output_index : info_->outputs()) { |
98 | int output_epoch = tensor_epochs_[output_index]; |
99 | if (output_epoch == kEpochAlwaysReady) { |
100 | // This happens when an input of subgraph is also an output of subgraph. |
101 | continue; |
102 | } |
103 | NodeSubset& output_subset = (*node_subsets_)[output_epoch]; |
104 | output_subset.output_tensors.push_back(output_index); |
105 | } |
106 | // Make sure every node sub set's inputs and outputs are unique, since the |
107 | // list of inputs and outputs is generated in a way that produces |
108 | // duplicates. |
109 | for (NodeSubset& node_subset : *node_subsets_) { |
110 | // Sort and uniquefy using standard library algorithms. |
111 | Uniquefy(&node_subset.input_tensors); |
112 | Uniquefy(&node_subset.output_tensors); |
113 | } |
114 | } |
115 | |
116 | private: |
117 | // Special integer values needed for tensor_epochs_ and node_epochs_. |
118 | enum { |
119 | // The node or tensor is not ready to be assigned an epoch. e.g. a node's |
120 | // inputs have not all been assigned epochs. |
121 | kEpochNotReady = -1, |
122 | // Used for tensor_epochs_. This means that the tensor is always ready. |
123 | // e.g. an input to the whole model or a constant that has no dependencies. |
124 | kEpochAlwaysReady = -2 |
125 | }; |
126 | |
127 | // Updates the node at `node_index` in the execution plan and returns true if |
128 | // it is assigned to an epoch. False is returned if the node is already set to |
129 | // an epoch, its inputs are not all assigned to epochs, or if it cannot be |
130 | // assigned to the current epoch since the epoch's node_type doesn't match. |
131 | bool UpdateNode(int node_index) { |
132 | const TfLiteNode& node = info_->node(node_index); |
133 | NodeSubset& current_subset = node_subsets_->back(); |
134 | int current_epoch = node_subsets_->size() - 1; |
135 | // Check if node is already done. |
136 | if (node_epochs_[node_index] != kEpochNotReady) { |
137 | return false; |
138 | } |
139 | // See if all dependencies of this node are already assigned to a |
140 | // node sub set. |
141 | for (int input_tensor_index : TfLiteIntArrayView(node.inputs)) { |
142 | if (input_tensor_index != kTfLiteOptionalTensor && |
143 | tensor_epochs_[input_tensor_index] == kEpochNotReady) { |
144 | return false; |
145 | } |
146 | } |
147 | // In order for the current node to be schedulable, all nodes on which it |
148 | // explicitly depends must have been scheduled. |
149 | if (num_incoming_control_edges_[node_index] != 0) { |
150 | return false; |
151 | } |
152 | |
153 | int original_node_idx = info_->node_index(node_index); |
154 | // When we are starting a new epoch, the first ready node defines |
155 | // the type of that epoch. |
156 | if (current_subset.type == NodeSubset::kTfUnexplored) { |
157 | current_subset.type = node_type_[original_node_idx]; |
158 | } |
159 | // The node gets assigned to this epoch if it is the same type as |
160 | // the epoch's assigned type. Note, if this is the current ready |
161 | // node encountered during this epoch, this condition will be |
162 | // automatically true. |
163 | if (current_subset.type == node_type_[original_node_idx]) { |
164 | node_epochs_[node_index] = current_epoch; |
165 | current_subset.nodes.push_back(original_node_idx); |
166 | // All outputs of this node now are assigned to this epoch as |
167 | // well. |
168 | for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) { |
169 | tensor_epochs_[output_tensor_index] = current_epoch; |
170 | } |
171 | // Look at our inputs one more time to update that tensor's |
172 | // epochs' outputs |
173 | for (int input_tensor_index : TfLiteIntArrayView(node.inputs)) { |
174 | if (input_tensor_index == kTfLiteOptionalTensor) { |
175 | continue; |
176 | } |
177 | int input_epoch = tensor_epochs_[input_tensor_index]; |
178 | int node_epoch = current_epoch; |
179 | if (input_epoch != node_epoch) { |
180 | current_subset.input_tensors.push_back(input_tensor_index); |
181 | // Set inputs to be outputs of the node sub set where they reside. |
182 | // the if condition makes sure inputs to the whole computation |
183 | // are not included (i.e. those initialized to -2 above). |
184 | if (input_epoch >= 0) { |
185 | NodeSubset& input_subset = (*node_subsets_)[input_epoch]; |
186 | input_subset.output_tensors.push_back(input_tensor_index); |
187 | } |
188 | } |
189 | } |
190 | |
191 | // Now that node_index is scheduled, remove it as a precondition from its |
192 | // dependent nodes. |
193 | for (auto edge_iter = |
194 | std::lower_bound(control_edges_.begin(), control_edges_.end(), |
195 | ControlEdge(node_index, 0)); |
196 | edge_iter != control_edges_.end() && edge_iter->first == node_index; |
197 | ++edge_iter) { |
198 | --num_incoming_control_edges_[edge_iter->second]; |
199 | } |
200 | return true; |
201 | } else { |
202 | return false; |
203 | } |
204 | } |
205 | |
206 | // Completely populates the current node_subset by doing graph traversal |
207 | void BuildNodeSubset() { |
208 | node_subsets_->emplace_back(NodeSubset()); |
209 | // loop until no more nodes can be updated. |
210 | while (true) { |
211 | bool did_something = false; |
212 | for (int node_index = 0; node_index < info_->num_execution_nodes(); |
213 | node_index++) { |
214 | if (UpdateNode(node_index)) { |
215 | did_something = true; |
216 | } |
217 | } |
218 | if (!did_something) return; |
219 | } |
220 | } |
221 | |
222 | // Temporary data needed for partitioning. |
223 | const GraphInfo* info_; |
224 | // List of node_subsets to populate |
225 | std::vector<NodeSubset>* node_subsets_; |
226 | // NOTE: This vector contains a place-holder for *all* nodes in the graph, not |
227 | // just ones in the execution plan. This is because nodes_to_partition is |
228 | // passed in as a list of original node indices & not execution plan indices. |
229 | std::vector<NodeSubset::Type> node_type_; |
230 | // Maps from tensor index to the epoch in which it is assigned. Also special |
231 | // negative values of kEpochNotReady if not assigned, kEpochAlwaysReady if it |
232 | // is an input to the whole model or a constant that has no dependencies. |
233 | std::vector<int> tensor_epochs_; |
234 | // Maps from tensor index to the epoch in which it is assigned. Also special |
235 | // negative values of kEpochNotReady if not assigned. |
236 | std::vector<int> node_epochs_; |
237 | |
238 | // Must be cycle-free. Before calling Partition(), must be sorted |
239 | // lexicographically. Duplicate entries are harmless. |
240 | ControlEdges control_edges_; |
241 | |
242 | // Number of incoming control edges for each node. |
243 | std::vector<int> num_incoming_control_edges_; |
244 | }; |
245 | // LINT.ThenChange(//tensorflow/lite/delegates/utils.h) |
246 | |
247 | } // namespace |
248 | |
249 | TfLiteStatus PartitionGraphIntoIndependentNodeSubsets( |
250 | const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, |
251 | const ControlEdges& control_edges, std::vector<NodeSubset>* node_subsets) { |
252 | PartitionGraphIntoIndependentNodeSubsetsImpl(info, nodes_to_partition, |
253 | control_edges, node_subsets) |
254 | .Partition(); |
255 | return kTfLiteOk; |
256 | } |
257 | |
258 | TfLiteStatus PartitionGraphIntoIndependentNodeSubsets( |
259 | const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, |
260 | std::vector<NodeSubset>* node_subsets) { |
261 | ControlEdges control_edges; |
262 | // Add a dependency chain between stateful ops. |
263 | for (int last_op_with_side_effect = -1, node_index = 0; |
264 | node_index < info->num_execution_nodes(); ++node_index) { |
265 | const auto& node = info->node(node_index); |
266 | if (node.might_have_side_effect) { |
267 | if (last_op_with_side_effect != -1) { |
268 | control_edges.emplace_back(last_op_with_side_effect, node_index); |
269 | } |
270 | last_op_with_side_effect = node_index; |
271 | } |
272 | } |
273 | return PartitionGraphIntoIndependentNodeSubsets(info, nodes_to_partition, |
274 | control_edges, node_subsets); |
275 | } |
276 | |
277 | } // namespace tflite |
278 | |