Plot a given instruction neighborhood or computation with graphviz.
| 523 | |
| 524 | // Plot a given instruction neighborhood or computation with graphviz. |
| 525 | void DoPlotCommand(const Options& opts, const HloModule& module, |
| 526 | const std::vector<string>& tokens) { |
| 527 | string node_name = tokens[0]; |
| 528 | |
| 529 | // Find the node with the given name. |
| 530 | const HloInstruction* instr = FindInstruction(module, node_name); |
| 531 | const HloComputation* comp = FindComputation(module, node_name); |
| 532 | if (!instr && !comp) { |
| 533 | std::cerr << "Couldn't find HloInstruction or HloComputation named " |
| 534 | << node_name << "." << std::endl; |
| 535 | return; |
| 536 | } |
| 537 | |
| 538 | uint64 graph_width = kDefaultWidth; |
| 539 | absl::flat_hash_set<const HloInstruction*> boundary; |
| 540 | if (tokens.size() >= 2) { |
| 541 | if (comp) { |
| 542 | std::cerr << "Can only use graph-size parameter with instructions, but " |
| 543 | << node_name << " is a computation." << std::endl; |
| 544 | return; |
| 545 | } |
| 546 | |
| 547 | int bound_index = 1; |
| 548 | // Get the <width> if present. |
| 549 | if (absl::SimpleAtoi(tokens[bound_index], &graph_width)) { |
| 550 | bound_index++; |
| 551 | } else { |
| 552 | // <width> not found, need to reset graph_width. |
| 553 | graph_width = kDefaultWidth; |
| 554 | } |
| 555 | // Get the '/'. |
| 556 | if (bound_index < tokens.size()) { |
| 557 | // This token must be a '/'. |
| 558 | if (tokens[bound_index] != "/") { |
| 559 | std::cerr << "Expect a /, but get a '" << tokens[bound_index] << "'." |
| 560 | << std::endl; |
| 561 | return; |
| 562 | } |
| 563 | bound_index++; |
| 564 | } |
| 565 | // Get the boundary nodes. |
| 566 | while (bound_index < tokens.size()) { |
| 567 | string bnode_name = tokens[bound_index]; |
| 568 | const HloInstruction* binstr = FindInstruction(module, bnode_name); |
| 569 | if (!binstr) { |
| 570 | std::cerr << "Couldn't find HloInstruction named " << bnode_name << "." |
| 571 | << std::endl; |
| 572 | return; |
| 573 | } |
| 574 | boundary.insert(binstr); |
| 575 | bound_index++; |
| 576 | } |
| 577 | } |
| 578 | |
| 579 | // Generate the graph and print the resulting string, which should be a |
| 580 | // graphviz url. |
| 581 | if (comp) { |
| 582 | RenderAndDisplayGraph(opts, [&](RenderedGraphFormat format) { |
no test coverage detected