| 554 | } |
| 555 | |
| 556 | std::vector<HloComputation*> HloModule::MakeComputationPostOrder() const { |
| 557 | // First determine all root computations by building a set of nonroot |
| 558 | // computations (computations which are called by an instruction in the |
| 559 | // module). |
| 560 | absl::flat_hash_set<HloComputation*> nonroot_computations; |
| 561 | for (auto& computation : computations_) { |
| 562 | for (auto* instruction : computation->instructions()) { |
| 563 | for (HloComputation* called_computation : |
| 564 | instruction->called_computations()) { |
| 565 | nonroot_computations.insert(called_computation); |
| 566 | } |
| 567 | } |
| 568 | } |
| 569 | |
| 570 | // Keep track of computations which have already been added to the post |
| 571 | // order. This prevents duplication as an embedded computation may be called |
| 572 | // from two different root computations. |
| 573 | absl::flat_hash_set<HloComputation*> added_computations; |
| 574 | std::vector<HloComputation*> post_order; |
| 575 | for (auto& computation : computations_) { |
| 576 | if (!nonroot_computations.contains(computation.get())) { |
| 577 | for (HloComputation* embedded_computation : |
| 578 | computation->MakeEmbeddedComputationsList()) { |
| 579 | if (!added_computations.contains(embedded_computation)) { |
| 580 | post_order.push_back(embedded_computation); |
| 581 | added_computations.insert(embedded_computation); |
| 582 | } |
| 583 | } |
| 584 | // Root computations should only be encountered once. |
| 585 | CHECK(!added_computations.contains(computation.get())); |
| 586 | post_order.push_back(computation.get()); |
| 587 | added_computations.insert(computation.get()); |
| 588 | } |
| 589 | } |
| 590 | if (post_order.size() != computations_.size()) { |
| 591 | for (HloComputation* computation : post_order) { |
| 592 | LOG(ERROR) << "Post Order: " << computation->name() << " (" |
| 593 | << computation->parent()->name() << ")"; |
| 594 | } |
| 595 | for (auto& computation : computations_) { |
| 596 | LOG(ERROR) << "Computations: " << computation->name() << " (" |
| 597 | << computation->parent()->name() << ")"; |
| 598 | } |
| 599 | LOG(FATAL) << "Mismatch computation count: post_order=" << post_order.size() |
| 600 | << " computation_count=" << computations_.size(); |
| 601 | } |
| 602 | return post_order; |
| 603 | } |
| 604 | |
| 605 | std::vector<HloComputation*> HloModule::MakeComputationSortedByContent() const { |
| 606 | auto result = MakeComputationPostOrder(); |