| 1608 | } |
| 1609 | |
| 1610 | StatusOr<std::unique_ptr<BufferAssignment>> BufferAssigner::CreateAssignment( |
| 1611 | const HloModule* module, std::unique_ptr<HloOrdering> hlo_ordering, |
| 1612 | BufferValue::SizeFunction buffer_size, |
| 1613 | LogicalBuffer::AlignmentFunction color_alignment, |
| 1614 | HloDataflowAnalysis::CanShareBuffer can_share_buffer) { |
| 1615 | TF_ASSIGN_OR_RETURN(std::unique_ptr<HloAliasAnalysis> alias_analysis, |
| 1616 | HloAliasAnalysis::Run(module, can_share_buffer)); |
| 1617 | |
| 1618 | // Set up a schedule for each computation. |
| 1619 | HloSchedule schedule(module); |
| 1620 | for (const HloComputation* computation : module->computations()) { |
| 1621 | const HloInstructionSequence* instruction_sequence = |
| 1622 | hlo_ordering->SequentialOrder(*computation); |
| 1623 | const bool has_sequential_order = instruction_sequence != nullptr; |
| 1624 | if (has_sequential_order) { |
| 1625 | schedule.set_sequence(computation, *instruction_sequence); |
| 1626 | } |
| 1627 | } |
| 1628 | |
| 1629 | TF_ASSIGN_OR_RETURN(std::unique_ptr<HloLiveRange> hlo_live_range, |
| 1630 | HloLiveRange::Run(schedule, *alias_analysis, |
| 1631 | module->entry_computation(), true)); |
| 1632 | |
| 1633 | VLOG(1) << "Assigning buffers to module " << module->name(); |
| 1634 | XLA_VLOG_LINES(3, module->ToString()); |
| 1635 | XLA_VLOG_LINES(3, alias_analysis->ToString()); |
| 1636 | XLA_VLOG_LINES(3, alias_analysis->dataflow_analysis().ToString()); |
| 1637 | VLOG(1) << "Number of buffers to assign: " |
| 1638 | << alias_analysis->buffers().size(); |
| 1639 | |
| 1640 | // Can't use absl::make_unique because BufferAssignment constructor is |
| 1641 | // private. |
| 1642 | std::unique_ptr<BufferAssignment> assignment(new BufferAssignment( |
| 1643 | module, std::move(hlo_ordering), std::move(buffer_size), |
| 1644 | std::move(color_alignment), std::move(alias_analysis), |
| 1645 | std::move(hlo_live_range))); |
| 1646 | |
| 1647 | TF_RETURN_IF_ERROR( |
| 1648 | colorer_(&assignment->alias_analysis(), assignment->hlo_ordering())); |
| 1649 | VLOG(3) << "After coloring:"; |
| 1650 | XLA_VLOG_LINES(3, |
| 1651 | assignment->alias_analysis().dataflow_analysis().ToString()); |
| 1652 | TF_RETURN_IF_ERROR(MergeInplaceOpBuffers(assignment.get())); |
| 1653 | |
| 1654 | std::vector<const HloComputation*> thread_local_computations; |
| 1655 | std::vector<const HloComputation*> global_computations; |
| 1656 | TF_RETURN_IF_ERROR(GatherComputationsByAllocationType( |
| 1657 | module, &thread_local_computations, &global_computations)); |
| 1658 | |
| 1659 | // First assign buffers for global computations. Temporary buffers for |
| 1660 | // sequential computations are collected in |
| 1661 | // 'buffers_to_assign_sequentially'. |
| 1662 | flat_hash_map<const HloComputation*, flat_hash_set<const HloValue*>> |
| 1663 | buffers_to_assign_sequentially; |
| 1664 | TF_RETURN_IF_ERROR(AssignBuffersForComputations( |
| 1665 | global_computations, |
| 1666 | /*is_thread_local=*/false, &buffers_to_assign_sequentially, |
| 1667 | assignment.get())); |
no test coverage detected