Backwards dataflow analysis that finds arguments to a graph that must be compile-time constants.
| 167 | // Backwards dataflow analysis that finds arguments to a graph that must be |
| 168 | // compile-time constants. |
| 169 | Status BackwardsConstAnalysis(const Graph& g, |
| 170 | std::vector<bool>* compile_time_const_arg_indices, |
| 171 | std::vector<bool>* compile_time_const_nodes, |
| 172 | FunctionLibraryRuntime* flib_runtime, |
| 173 | std::function<bool(const Edge&)> edge_filter) { |
| 174 | std::vector<bool> compile_time_const_nodes_impl; |
| 175 | if (compile_time_const_nodes) { |
| 176 | CHECK_EQ(compile_time_const_nodes->size(), g.num_node_ids()); |
| 177 | } else { |
| 178 | compile_time_const_nodes_impl.resize(g.num_node_ids()); |
| 179 | compile_time_const_nodes = &compile_time_const_nodes_impl; |
| 180 | } |
| 181 | |
| 182 | Status status; |
| 183 | auto visit = [&](Node* node) { |
| 184 | if (!status.ok()) return; |
| 185 | |
| 186 | // If this is a metadata-only op, don't propagate the const requirement. |
| 187 | if (XlaOpRegistry::IsMetadataOp(node->type_string())) { |
| 188 | return; |
| 189 | } |
| 190 | |
| 191 | // If this node must be const, and it isn't a metadata op, then all of its |
| 192 | // parents must be const. |
| 193 | if ((*compile_time_const_nodes)[node->id()]) { |
| 194 | if (node->type_string() == "_Arg") { |
| 195 | int index; |
| 196 | status = GetNodeAttr(node->attrs(), "index", &index); |
| 197 | if (!status.ok()) return; |
| 198 | if (compile_time_const_arg_indices) { |
| 199 | (*compile_time_const_arg_indices)[index] = true; |
| 200 | } |
| 201 | return; |
| 202 | } |
| 203 | for (const Edge* pred : node->in_edges()) { |
| 204 | if (!pred->IsControlEdge() && edge_filter(*pred)) { |
| 205 | // If the src node of the `pred` is an IdentityN do not mark it as a |
| 206 | // compile-time const. Only mark the corresponding input to the |
| 207 | // IdentityN node as a const. |
| 208 | // Note: XLA IdentityN op simply forwards its inputs so this is safe. |
| 209 | while (edge_filter(*pred) && |
| 210 | pred->src()->type_string() == "IdentityN") { |
| 211 | status = pred->src()->input_edge(pred->src_output(), &pred); |
| 212 | if (!status.ok()) return; |
| 213 | } |
| 214 | if (edge_filter(*pred)) { |
| 215 | (*compile_time_const_nodes)[pred->src()->id()] = true; |
| 216 | } |
| 217 | } |
| 218 | } |
| 219 | return; |
| 220 | } |
| 221 | |
| 222 | // Mark any compile-time constant operator arguments as const. |
| 223 | std::vector<int> const_input_idxs; |
| 224 | status = GetCompileTimeConstInputs(node, &const_input_idxs, flib_runtime); |
| 225 | |
| 226 | if (!status.ok()) { |