| 161 | } |
| 162 | |
| 163 | bool OptimizerCSE::Optimize( |
| 164 | const std::function<bool(const Node*)>& consider_fn) { |
| 165 | // This very simple implementation works if the whole graph is one |
| 166 | // giant basic block (because we just traverse nodes in a |
| 167 | // topological order). This simple implementation works well |
| 168 | // with control flow/loops/etc. But we need to be careful about |
| 169 | // control flow if we want to add more sophisticated CSE optimizations. |
| 170 | |
| 171 | // TODO(jeff): We need to handle Update nodes specially, but dealing |
| 172 | // with more general control flow will also solve this issue, and for |
| 173 | // now, our updates are almost always the most downstream nodes in |
| 174 | // the graph. |
| 175 | std::vector<Node*> order; |
| 176 | GetReversePostOrder(*g_, &order); |
| 177 | |
| 178 | // Our value is just a single Node*, meaning we keep just a single |
| 179 | // candidate for a given node hash value. This may cause us to |
| 180 | // (rarely) lose some optimization opportunities if there are |
| 181 | // hash collisions, but it allows us to avoid having the value |
| 182 | // be a set<Node*> (or equivalent). |
| 183 | std::unordered_map<size_t, Node*> available; |
| 184 | |
| 185 | // Scratch space for Equivalent calls. Allocated here and passed in to |
| 186 | // Equivalent to avoid allocation inside the loop below. |
| 187 | bool changed = false; |
| 188 | AttrSlice::Scratch scratch; |
| 189 | for (Node* n : order) { |
| 190 | if (!n->IsOp()) continue; |
| 191 | |
| 192 | // Don't prune placeholder nodes. |
| 193 | if (n->type_string() == "Placeholder" || |
| 194 | n->type_string() == "PlaceholderV2" || |
| 195 | n->type_string() == "PlaceholderWithDefault") { |
| 196 | continue; |
| 197 | } |
| 198 | |
| 199 | // See if we should consider this node at all |
| 200 | if (consider_fn != nullptr && !consider_fn(n)) continue; |
| 201 | |
| 202 | size_t h = NodeHash(n); |
| 203 | Node** candidate = &available[h]; |
| 204 | if (*candidate == nullptr) { |
| 205 | // No existing match: insert "n" into the hash table under "h" |
| 206 | *candidate = n; |
| 207 | } else if (Equivalent(*candidate, n, &scratch)) { |
| 208 | VLOG(1) << "CSE: equivalent: " << (*candidate)->name() << " and " |
| 209 | << n->name(); |
| 210 | // *candidate and n are equivalent. Therefore, we can replace |
| 211 | // n with *candidate by fixing up outgoing edges from "n" to instead |
| 212 | // come from "*candidate", and then delete n from the graph |
| 213 | for (const Edge* e : n->out_edges()) { |
| 214 | g_->AddEdge(*candidate, e->src_output(), e->dst(), e->dst_input()); |
| 215 | } |
| 216 | |
| 217 | MergeDebugInfo(NodeDebugInfo(*n), *candidate); |
| 218 | g_->RemoveNode(n); |
| 219 | changed = true; |
| 220 | } |