Method HandleFusion

tensorflow/compiler/xla/service/gpu/fusion_merger.cc:181–293

}

Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) {
  ++total_visited_;
  // Skip 'fusion' instruction if there are no users into which we can merge.
  if (fusion->users().empty()) {
    VLOG(3) << "Not merging " << fusion->name() << ": Has no users.";
    ++num_fail_no_users_;
    return Status::OK();
  }

  // Skip 'fusion' instruction if it is not a loop fusion. Library fusion
  // instructions match specific patterns, so they shouldn't be further fused.
  // Input fusion instructions need to be rooted at a particular HLO (e.g.
  // kReduce), so they shouldn't be further fused either.
  if (!fusion->IsLoopFusion()) {
    VLOG(3) << "Not merging " << fusion->name() << ": Is not loop fusion.";
    ++num_fail_not_loop_fusion_;
    return Status::OK();
  }

  // Skip 'fusion' instruction if we cannot merge into all of its users.
  // Merging into all users enables the removal of 'fusion' from the
  // computation.
  if (!absl::c_all_of(fusion->users(), [&](const HloInstruction* user) {
        return user->opcode() == HloOpcode::kFusion &&
               IsProducerConsumerFusible(*fusion, *user);
      })) {
    VLOG(3) << "Not merging " << fusion->name()
            << ": Some of its users are not loop/input fusion kernels.";
    ++num_fail_merge_all_users_;
    return Status::OK();
  }

  // Skip 'fusion' instruction if any of its fused instructions are expensive.
  // This is done to avoid the duplication of expensive instructions, which
  // would occur if 'fusion' were merged into multiple users.
  //
  // If 'fusion' has just one user, then an earlier fusion pass chose not to
  // fuse this producer/consumer pair (likely because of expensive instruction
  // re-use by the consumer), and so we honor that choice here as well.
  if (absl::c_any_of(fusion->fused_instructions(),
                     [](const HloInstruction* instruction) {
                       return instruction->opcode() != HloOpcode::kParameter &&
                              GpuInstructionFusion::IsExpensive(*instruction);
                     })) {
    VLOG(3) << "Not merging " << fusion->name()
            << ": Contains one or more expensive instructions.";
    ++num_fail_expensive_fused_instruction_;
    return Status::OK();
  }

  // Skip 'fusion' instruction if merging it into all users would result in a
  // net increase in bytes transferred (currently allowing the net bytes
  // transferred to be exceeded up to ~10% in exchange for eliminating the
  // overhead from a GPU kernel launch).
  const double current_bytes_transferred = GetCurrentBytesTransferred(fusion);
  const double merged_bytes_transferred = GetMergedBytesTransferred(fusion);
  const double merged_to_current_bytes_ratio =
      merged_bytes_transferred / std::max(1.0, current_bytes_transferred);