| 61 | } |
| 62 | |
| 63 | HloComputation* HloModule::AddComputationInternal( |
| 64 | std::unique_ptr<HloComputation> computation, bool is_entry, |
| 65 | bool uniquify_identifiers) { |
| 66 | if (is_entry) { |
| 67 | CHECK_EQ(nullptr, entry_computation_); |
| 68 | entry_computation_ = computation.get(); |
| 69 | |
| 70 | // If the module configuration has no entry layout computation set, create a |
| 71 | // default one based on the program shape. |
| 72 | if (!config_.has_entry_computation_layout()) { |
| 73 | config_.SetDefaultComputationLayout( |
| 74 | entry_computation_->ComputeProgramShape()); |
| 75 | } |
| 76 | input_output_alias_config_ = HloInputOutputAliasConfig( |
| 77 | entry_computation_->root_instruction()->shape()); |
| 78 | } |
| 79 | |
| 80 | if (uniquify_identifiers) { |
| 81 | computation->UniquifyName(&computation_name_uniquer_); |
| 82 | for (auto* instruction : computation->instructions()) { |
| 83 | instruction->UniquifyName(&instruction_name_uniquer_); |
| 84 | } |
| 85 | |
| 86 | // Pick unique IDs for each instruction. |
| 87 | for (auto* instruction : computation->instructions()) { |
| 88 | instruction->SetUniqueId(NewUniqueInstructionId()); |
| 89 | } |
| 90 | // Set unique id to this computation. |
| 91 | CHECK_NE(computation->root_instruction()->unique_id(), -1) |
| 92 | << "Root has no valid id: " << computation->ToString(); |
| 93 | computation->SetUniqueId(computation->root_instruction()->unique_id()); |
| 94 | } else { |
| 95 | // Don't uniquify the names of the computation or instruction, but we must |
| 96 | // run the names through the uniquifiers to prevent future name collisions |
| 97 | // for computations and instructions created later. Also, set the |
| 98 | // next_unique_id_ to the one greater than the max unique id of any |
| 99 | // instruction (or the computation) to avoid ID collisions. |
| 100 | computation_name_uniquer_.GetUniqueName(computation->name()); |
| 101 | for (auto* instruction : computation->instructions()) { |
| 102 | instruction_name_uniquer_.GetUniqueName(instruction->name()); |
| 103 | next_unique_id_ = std::max(next_unique_id_, instruction->unique_id() + 1); |
| 104 | } |
| 105 | if (next_unique_id_ < computation->unique_id() + 1) { |
| 106 | next_unique_id_ = computation->unique_id() + 1; |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | computation->set_parent(this); |
| 111 | computations_.push_back(std::move(computation)); |
| 112 | return computations_.back().get(); |
| 113 | } |
| 114 | |
| 115 | HloComputation* HloModule::AddEntryComputation( |
| 116 | std::unique_ptr<HloComputation> computation) { |
no test coverage detected