| 1587 | } |
| 1588 | |
| 1589 | Status MemorySpaceAssignment::FixSchedule() { |
| 1590 | CHECK(module_->has_schedule()); |
| 1591 | HloSchedule& schedule = module_->schedule(); |
| 1592 | for (const HloComputation* computation : |
| 1593 | module_->MakeNonfusionComputations()) { |
| 1594 | // Parallel computations aren't in the schedule and don't need to be |
| 1595 | // modified. |
| 1596 | if (!computations_in_schedule_.contains(computation)) { |
| 1597 | VLOG(4) << "Not scheduling " << computation->name() |
| 1598 | << " because it's not in the schedule."; |
| 1599 | continue; |
| 1600 | } |
| 1601 | CHECK(schedule.is_computation_scheduled(computation)); |
| 1602 | HloInstructionSequence new_sequence; |
| 1603 | |
| 1604 | absl::flat_hash_set<HloInstruction*> inserted_instructions; |
| 1605 | |
| 1606 | VLOG(4) << "Scheduling: " << computation->ToString(); |
| 1607 | |
| 1608 | for (int64 instruction_index = 0; |
| 1609 | instruction_index < flattened_instructions_.size(); |
| 1610 | ++instruction_index) { |
| 1611 | auto insts_before_iter = schedule_before_.find(instruction_index); |
| 1612 | if (insts_before_iter != schedule_before_.end()) { |
| 1613 | for (HloInstruction* new_instruction : insts_before_iter->second) { |
| 1614 | if (new_instruction->parent() == computation) { |
| 1615 | VLOG(4) << "before " << instruction_index << ": " |
| 1616 | << new_instruction->name(); |
| 1617 | EnsureInstructionAndOperandsInserted(new_instruction, &new_sequence, |
| 1618 | &inserted_instructions); |
| 1619 | } |
| 1620 | } |
| 1621 | } |
| 1622 | HloInstruction* instruction = flattened_instructions_[instruction_index]; |
| 1623 | // Insert only if it is not deleted (SimplifyGraph sets it to nullptr if |
| 1624 | // it was deleted) and not previously inserted. Also bitcasts and tuples |
| 1625 | // are treated specially and only inserted as a result of operand |
| 1626 | // dependencies. |
| 1627 | if (instruction != nullptr && |
| 1628 | !inserted_instructions.contains(instruction) && |
| 1629 | instruction->parent() == computation && |
| 1630 | instruction->opcode() != HloOpcode::kBitcast && |
| 1631 | instruction->opcode() != HloOpcode::kTuple) { |
| 1632 | VLOG(4) << "inst " << instruction_index << ": " << instruction->name(); |
| 1633 | EnsureInstructionAndOperandsInserted(instruction, &new_sequence, |
| 1634 | &inserted_instructions); |
| 1635 | } |
| 1636 | auto insts_after_iter = schedule_after_.find(instruction_index); |
| 1637 | if (insts_after_iter != schedule_after_.end()) { |
| 1638 | for (HloInstruction* new_instruction : insts_after_iter->second) { |
| 1639 | if (new_instruction->parent() == computation) { |
| 1640 | VLOG(4) << "after " << instruction_index << ": " |
| 1641 | << new_instruction->name(); |
| 1642 | EnsureInstructionAndOperandsInserted(new_instruction, &new_sequence, |
| 1643 | &inserted_instructions); |
| 1644 | } |
| 1645 | } |
| 1646 | } |
no test coverage detected