The algorithm first does a forward pass (parameters to root) to determine a set of instructions to consider using bfloat16, then does a backward pass to determine the precisions of those instructions according to the need of their users. During the backward pass, the potential changes are stored in changes_to_bf16_ which are subject to further adjustments then applied to the HLOs.
| 731 | // changes_to_bf16_ which are subject to further adjustments then applied to the |
| 732 | // HLOs. |
| 733 | StatusOr<bool> BFloat16Propagation::Run(HloModule* module) { |
| 734 | consider_using_bfloat16_.clear(); |
| 735 | instructions_visited_in_backward_pass_.clear(); |
| 736 | computations_visited_in_backward_pass_.clear(); |
| 737 | values_that_must_be_kept_as_f32_.clear(); |
| 738 | caller_counts_.clear(); |
| 739 | changes_to_bf16_.clear(); |
| 740 | changed_ = false; |
| 741 | |
| 742 | auto computations_topological_order = module->MakeComputationPostOrder(); |
| 743 | |
| 744 | // Before running the propagation pass, we insert copies (kConvert to the same |
| 745 | // type) of F32 inputs to while loops. This prevents other uses of the same |
| 746 | // input from aliasing the while loop input/output, so that there's greater |
| 747 | // chance to use BF16 inside the loop. If some of these added copies do not |
| 748 | // help, they will remain F32 after BF16 propagation and will be removed since |
| 749 | // they are no-ops. |
| 750 | for (auto computation : computations_topological_order) { |
| 751 | for (auto inst : computation->MakeInstructionPostOrder()) { |
| 752 | if (inst->opcode() != HloOpcode::kWhile) { |
| 753 | continue; |
| 754 | } |
| 755 | |
| 756 | auto operand = inst->mutable_operand(0); |
| 757 | TF_ASSIGN_OR_RETURN( |
| 758 | HloInstruction * copy, |
| 759 | computation->DeepCopyInstructionWithCustomCopier( |
| 760 | operand, [](HloInstruction* leaf, const ShapeIndex& leaf_index, |
| 761 | HloComputation* comp) { |
| 762 | if (leaf->shape().element_type() != F32) { |
| 763 | return leaf; |
| 764 | } |
| 765 | return comp->AddInstruction( |
| 766 | HloInstruction::CreateConvert(leaf->shape(), leaf)); |
| 767 | })); |
| 768 | TF_RETURN_IF_ERROR(operand->ReplaceUseWith(inst, copy)); |
| 769 | } |
| 770 | } |
| 771 | |
| 772 | TF_ASSIGN_OR_RETURN(dataflow_, HloDataflowAnalysis::Run(*module)); |
| 773 | |
| 774 | // The first step is a forward pass (parameters to root), where we determine |
| 775 | // the potential candidate instructions to use bfloat16 in the outputs that |
| 776 | // are not likely to cause overhead from extra explicit conversions. This is |
| 777 | // done forwardly because we determine whether an HLO is a candidate partially |
| 778 | // based on whether its operands are candidates. |
| 779 | for (auto computation : computations_topological_order) { |
| 780 | for (auto inst : computation->MakeInstructionPostOrder()) { |
| 781 | if (InstructionIsCandidateForBF16Output(inst)) { |
| 782 | consider_using_bfloat16_.insert(inst); |
| 783 | } |
| 784 | } |
| 785 | } |
| 786 | |
| 787 | // The second step is a backward pass (root to parameters), where we modify |
| 788 | // the precisions of the instructions identified in the first step when |
| 789 | // feasible. This is done backwardly because we determine the precision of an |
| 790 | // HLO's output based on how it is later used. |
nothing calls this directly
no test coverage detected