| 27 | namespace xla { |
| 28 | |
| 29 | StatusOr<bool> MultiOutputFusion::Run(HloModule* module) { |
| 30 | bool changed = false; |
| 31 | |
| 32 | for (auto* computation : module->MakeNonfusionComputations()) { |
| 33 | computation_ = computation; |
| 34 | candidates_.clear(); |
| 35 | candidates_index_.clear(); |
| 36 | all_fusion_candidates_.clear(); |
| 37 | RecomputeReachability(); |
| 38 | |
| 39 | int64 index = 0; |
| 40 | for (auto it : computation_->MakeInstructionPostOrder()) { |
| 41 | candidates_.emplace_back(it); |
| 42 | InsertOrDie(&candidates_index_, it, index++); |
| 43 | } |
| 44 | |
| 45 | // Create the initial candidate list for each Node. |
| 46 | for (auto& node : candidates_) { |
| 47 | HloInstruction* instruction = node.hlo; |
| 48 | int64 instruction_id = get_candidate_id(instruction); |
| 49 | FusionCandidate& instr_node = candidates_[instruction_id]; |
| 50 | if (!IsFusible(instruction)) { |
| 51 | continue; |
| 52 | } |
| 53 | all_fusion_candidates_.push_back(instruction); |
| 54 | |
| 55 | std::vector<HloInstruction*> candidates; |
| 56 | absl::flat_hash_set<HloInstruction*> candidates_set; |
| 57 | VLOG(10) << "Looking at instruction: " << instruction->name(); |
| 58 | for (auto operand : instruction->operands()) { |
| 59 | // Filter out the non-interesting instructions -- they |
| 60 | // will not generate the savings. |
| 61 | if (!IsProfitableOperand(operand)) { |
| 62 | VLOG(10) << "Operand not profitable: " << operand->name(); |
| 63 | continue; |
| 64 | } |
| 65 | VLOG(10) << "Operand profitable: " << operand->name(); |
| 66 | // We don't look at all users of operands as it's quadratic. Only look |
| 67 | // at one slice of users. |
| 68 | const int64 kUserSliceSize = 128; |
| 69 | |
| 70 | const int64 user_slice_begin = |
| 71 | RoundDownToNearest(operand->UserId(instruction), kUserSliceSize); |
| 72 | |
| 73 | const int64 user_slice_end = |
| 74 | std::min(static_cast<int64>(operand->users().size()), |
| 75 | user_slice_begin + kUserSliceSize); |
| 76 | |
| 77 | for (int64 i = user_slice_begin; i < user_slice_end; ++i) { |
| 78 | HloInstruction* user = operand->users()[i]; |
| 79 | VLOG(10) << "User: " << user->name(); |
| 80 | if (user == instruction || !IsFusible(user)) { |
| 81 | VLOG(10) << "User is not fusible, or is the instruction itself: " |
| 82 | << user->name(); |
| 83 | continue; |
| 84 | } |
| 85 | int64 user_id = get_candidate_id(user); |
| 86 | if (is_connected(instruction, user)) { |
nothing calls this directly
no test coverage detected