| 374 | } // namespace |
| 375 | |
| 376 | Status BuildCgmodeOpsPass::Run(const GraphOptimizationPassOptions& options) { |
| 377 | if (!options.session_options->config.gpu_options().cuda_graph_enable_jit()) { |
| 378 | return Status::OK(); |
| 379 | } |
| 380 | Graph* graph = options.graph->get(); |
| 381 | // Copy out the nodes we want to rewrite to avoid modifying the graph while we |
| 382 | // iterate on graph->op_nodes(). |
| 383 | std::vector<Node*> cgmode_compiled_kernels; |
| 384 | absl::c_copy_if(graph->op_nodes(), |
| 385 | std::back_inserter(cgmode_compiled_kernels), |
| 386 | [](const Node* n) { |
| 387 | if (n->IsSend() || n->IsRecv() || n->IsControlFlow()) { |
| 388 | return false; |
| 389 | } |
| 390 | |
| 391 | // Only compile nodes that are marked for compilation by the |
| 392 | // compilation-marking pass (via 'attr_name'). |
| 393 | return IsCgmodeCompiledKernel(*n); |
| 394 | }); |
| 395 | |
| 396 | jit::DeviceInfoCache device_info_cache; |
| 397 | |
| 398 | for (Node* n : cgmode_compiled_kernels) { |
| 399 | TF_RETURN_IF_ERROR(ReplaceNodeWithCgmodeCompileAndCgmodeRun( |
| 400 | &device_info_cache, options, *options.flib_def, graph, n)); |
| 401 | } |
| 402 | |
| 403 | if (VLOG_IS_ON(1)) { |
| 404 | DumpGraphToFile("build_cgmode_ops", *graph, options.flib_def); |
| 405 | } |
| 406 | // debug |
| 407 | VLOG(1) << "dump graph after build cuda graph " << DebugString(graph); |
| 408 | GraphCycles cycles_graph; |
| 409 | TF_ASSIGN_OR_RETURN(bool cycle_detection_graph_ok, |
| 410 | CreateCycleDetectionGraph(graph, &cycles_graph)); |
| 411 | if (!cycle_detection_graph_ok) { |
| 412 | return errors::Internal("Could not form cycle detection graph"); |
| 413 | } |
| 414 | std::vector<Node*> compile_nodes; |
| 415 | std::vector<Node*> gpu_nodes; |
| 416 | for (Node* n : graph->op_nodes()) { |
| 417 | if (n->type_string() == "_CgmodeCompile") { |
| 418 | compile_nodes.emplace_back(n); |
| 419 | continue; |
| 420 | } |
| 421 | const string& device_name_str = !n->assigned_device_name().empty() |
| 422 | ? n->assigned_device_name() |
| 423 | : n->requested_device(); |
| 424 | DeviceNameUtils::ParsedName full_device_name; |
| 425 | DeviceNameUtils::ParseFullName(device_name_str, &full_device_name); |
| 426 | if (full_device_name.type == DEVICE_GPU) { |
| 427 | gpu_nodes.emplace_back(n); |
| 428 | } |
| 429 | } |
| 430 | std::vector<Node*> sorted_compiled_nodes; |
| 431 | for (Node* n : compile_nodes) { |
| 432 | bool added = false; |
| 433 | for (int i = 0; i < sorted_compiled_nodes.size(); i++) { |
nothing calls this directly
no test coverage detected