| 229 | } |
| 230 | |
| 231 | Status FunctionalizeControlFlowPass::Run( |
| 232 | const GraphOptimizationPassOptions& options) { |
| 233 | Graph* graph = options.graph->get(); |
| 234 | if (VLOG_IS_ON(4)) { |
| 235 | DumpGraphToFile("functionalize_control_flow_before", *graph, |
| 236 | options.flib_def); |
| 237 | } |
| 238 | std::unique_ptr<ProcessFunctionLibraryRuntime> pflr( |
| 239 | new ProcessFunctionLibraryRuntime( |
| 240 | /*device_mgr=*/nullptr, options.session_options->env, |
| 241 | TF_GRAPH_DEF_VERSION, options.flib_def, OptimizerOptions())); |
| 242 | FunctionLibraryRuntime* flr = |
| 243 | pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice); |
| 244 | |
| 245 | // Find XLA compile ops and its corresponding FunctionDef. |
| 246 | // TPUCompile op is not in the map because graph rewriting might happen |
| 247 | // multiple times, and we want to avoid functionalize it again. |
| 248 | static std::map<string, string>* kNodeTypeToFunctionAttrMapping = |
| 249 | new std::map<string, string>{ |
| 250 | // _TPUReplicate ops are generated by EncapsulateTPUComputationsPass. |
| 251 | {"_TPUReplicate", "computation"}, |
| 252 | // XlaLaunch ops are generated by EncapsulateXlaComputationsPass. |
| 253 | {"XlaLaunch", "function"}, |
| 254 | }; |
| 255 | std::map<string, absl::optional<string>> canonicalized_name_to_new_name; |
| 256 | bool fld_modified = false; |
| 257 | for (Node* n : graph->nodes()) { |
| 258 | auto it = kNodeTypeToFunctionAttrMapping->find(n->type_string()); |
| 259 | if (it == kNodeTypeToFunctionAttrMapping->end()) { |
| 260 | continue; |
| 261 | } |
| 262 | const string func_attr = it->second; |
| 263 | NameAttrList func; |
| 264 | TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), func_attr, &func)); |
| 265 | VLOG(2) << "Graph has node " << n->type_string() |
| 266 | << ". Corresponding function: " << func.name(); |
| 267 | string new_func_name = options.flib_def->UniqueFunctionName( |
| 268 | absl::StrCat(func.name(), "_f15n_")); |
| 269 | bool modified; |
| 270 | TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction( |
| 271 | func.name(), new_func_name, func.attr(), options.flib_def, flr, |
| 272 | &canonicalized_name_to_new_name, &modified)); |
| 273 | if (modified) { |
| 274 | n->ClearAttr(func_attr); |
| 275 | func.set_name(new_func_name); |
| 276 | n->AddAttr(func_attr, func); |
| 277 | fld_modified = true; |
| 278 | } |
| 279 | } |
| 280 | |
| 281 | // TODO(ylc, endlessroad): Change this to "if (fld_modified")" |
| 282 | if (false) { |
| 283 | if (VLOG_IS_ON(4)) { |
| 284 | DumpGraphToFile("functionalize_control_flow_before_prune", *graph, |
| 285 | options.flib_def); |
| 286 | } |
| 287 | TF_RETURN_IF_ERROR( |
| 288 | PruneUnreachableFunctionsFromGraph(*graph, options.flib_def)); |
nothing calls this directly
no test coverage detected