MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / HandleConst

Method HandleConst

tensorflow/core/grappler/optimizers/loop_optimizer.cc:117–188  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

115}
116
117Status LoopInvariantNodeMotionOptimizer::HandleConst(NodeDef* node,
118 const int num_outputs,
119 const int frame_id) {
120 NodeDef* const_node = nullptr;
121 if (num_outputs == 0) {
122 // all successor nodes are invariant
123 // Remove the control inputs from this frame to the const node,
124 // when moving it out of the frame (in parent frame)
125 const_node = node;
126 node_map_->RemoveInputs(node->name());
127 node->clear_input();
128 } else {
129 // some successor nodes are variant
130 // Have to keep the const node in the frame,
131 // so create a new one outside the frame (in parent frame)
132 const string const_node_name =
133 AddPrefixToNodeName(node->name(), kLoopOptimizer);
134 const_node = node_map_->GetNode(const_node_name);
135 if (const_node == nullptr) {
136 const_node = optimized_graph_->add_node();
137 const_node->set_name(const_node_name);
138 const_node->set_op("Const");
139 const_node->set_device(node->device());
140 *const_node->mutable_attr() = node->attr();
141 node_map_->AddNode(const_node->name(), const_node);
142 }
143 auto consumers = node_map_->GetOutputs(node->name());
144 for (auto* consumer : consumers) {
145 if (invariant_nodes_.count(consumer)) {
146 for (int i = 0; i < consumer->input_size(); ++i) {
147 if (NodeName(consumer->input(i)) == node->name()) {
148 if (IsControlInput(consumer->input(i))) {
149 *consumer->mutable_input(i) = AsControlDependency(*const_node);
150 } else {
151 *consumer->mutable_input(i) = const_node->name();
152 }
153 node_map_->AddOutput(const_node->name(), consumer->name());
154 node_map_->RemoveOutput(node->name(), consumer->name());
155 }
156 }
157 }
158 }
159 }
160 // add a control input from the parent frame
161 auto parent_it = frame_parent_.find(frame_id);
162 if (parent_it != frame_parent_.end()) {
163 int parent_id = parent_it->second;
164 auto loop_cond_it = loop_cond_.find(parent_id);
165 if (loop_cond_it == loop_cond_.end()) {
166 return errors::InvalidArgument("Frame ", frame_id,
167 " doesn't have a LoopCond node");
168 }
169 auto& loop_cond_name = loop_cond_it->second->name();
170 NodeDef* switch_node = nullptr;
171 for (auto* node : node_map_->GetOutputs(loop_cond_name)) {
172 if (node->op() == "Switch") {
173 switch_node = node;
174 break;

Callers

nothing calls this directly

Calls 15

AddPrefixToNodeNameFunction · 0.85
NodeNameFunction · 0.85
IsControlInputFunction · 0.85
AsControlDependencyFunction · 0.85
InvalidArgumentFunction · 0.85
RemoveInputsMethod · 0.80
set_opMethod · 0.80
attrMethod · 0.80
GetOutputsMethod · 0.80
RemoveOutputMethod · 0.80
nameMethod · 0.65
StrCatFunction · 0.50

Tested by

no test coverage detected