1/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
23namespace tflite {
24namespace {
25
26template <class T>
27void 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
43class 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
249TfLiteStatus 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
258TfLiteStatus 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