| 468 | } |
| 469 | |
| 470 | StatusOr<bool> InstructionFusion::Run(HloModule* module) { |
| 471 | bool changed = false; |
| 472 | module_ = module; |
| 473 | int64 fuse_count = 0; |
| 474 | std::vector<std::vector<bool>>* fusion_config = nullptr; |
| 475 | HloModuleConfig module_config; |
| 476 | if (config_collection_mode_ != FusionConfigCollection::kOff) { |
| 477 | module_config = module->config(); |
| 478 | fusion_config = module_config.mutable_fusion_config(); |
| 479 | fusion_config->clear(); |
| 480 | } |
| 481 | |
| 482 | // Use sorted computations because fusion configuration is order-sensitive. |
| 483 | for (auto* computation : module->MakeNonfusionComputationsSorted()) { |
| 484 | CHECK(!computation->IsFusionComputation()); |
| 485 | computation_ = computation; |
| 486 | reachability_ = HloReachabilityMap::Build(computation_); |
| 487 | |
| 488 | HloInstructionSet do_not_duplicate; |
| 489 | // If we allow duplications, we need to compute which instructions we do not |
| 490 | // want to duplicate based on a global analysis of the graph. |
| 491 | if (may_duplicate_) { |
| 492 | do_not_duplicate = |
| 493 | ComputeGloballyUnfusible(computation_->MakeInstructionPostOrder()); |
| 494 | } |
| 495 | auto fusion_queue = GetFusionQueue(computation_); |
| 496 | |
| 497 | // Instruction fusion effectively fuses edges in the computation graph |
| 498 | // (producer instruction -> consumer instruction) so we iterate over all |
| 499 | // edges. When we fuse an edge, we create a copy of the producer inside the |
| 500 | // fusion instruction. |
| 501 | while (true) { |
| 502 | auto next_entry = |
| 503 | fusion_queue->DequeueNextInstructionAndOperandsToFuseInOrder(); |
| 504 | auto instruction = next_entry.first; |
| 505 | if (instruction == nullptr) { |
| 506 | break; |
| 507 | } |
| 508 | |
| 509 | if (!instruction->IsFusible() && |
| 510 | instruction->opcode() != HloOpcode::kFusion) { |
| 511 | continue; |
| 512 | } |
| 513 | |
| 514 | std::vector<int64>& sorted_operand_numbers = next_entry.second; |
| 515 | |
| 516 | for (int64 i : sorted_operand_numbers) { |
| 517 | HloInstruction* operand = instruction->mutable_operand(i); |
| 518 | |
| 519 | if (!operand->IsFusible()) { |
| 520 | continue; |
| 521 | } |
| 522 | |
| 523 | // Consumes a unit of compiler fuel and returns true if we should |
| 524 | // continue with the transformation. |
| 525 | auto consume_fuel = [&] { |
| 526 | return ConsumeFuel(name(), /*ran_out_of_fuel_msg=*/[&] { |
| 527 | return absl::StrFormat("Not fusing operand %d of %s, namely, %s", i, |
nothing calls this directly
no test coverage detected