static*/
| 94 | } |
| 95 | |
| 96 | /*static*/ StatusOr<WhileUtil::MakeInstructionsLiveInResult> |
| 97 | WhileUtil::MakeInstructionsLiveIn( |
| 98 | HloInstruction* while_instr, |
| 99 | absl::Span<HloInstruction* const> instructions) { |
| 100 | CHECK(while_instr->shape().IsTuple()); |
| 101 | |
| 102 | int64 elements_in_old_while_shape = while_instr->shape().tuple_shapes_size(); |
| 103 | Shape new_while_shape = while_instr->shape(); |
| 104 | for (auto* instruction : instructions) { |
| 105 | *new_while_shape.add_tuple_shapes() = instruction->shape(); |
| 106 | } |
| 107 | |
| 108 | TF_ASSIGN_OR_RETURN( |
| 109 | HloComputation * new_while_condition, |
| 110 | WidenWhileCondition(while_instr->while_condition(), new_while_shape)); |
| 111 | |
| 112 | HloComputation* new_while_body; |
| 113 | CallInliner::InlinedInstructionMap inlined_instructions_map; |
| 114 | TF_ASSIGN_OR_RETURN( |
| 115 | std::tie(new_while_body, inlined_instructions_map), |
| 116 | WidenWhileBody(while_instr->while_body(), new_while_shape)); |
| 117 | |
| 118 | HloInstruction* new_while_init = |
| 119 | TupleUtil::AppendSuffix(while_instr->mutable_operand(0), instructions); |
| 120 | HloComputation* containing_computation = while_instr->parent(); |
| 121 | HloInstruction* new_while = containing_computation->AddInstruction( |
| 122 | HloInstruction::CreateWhile(new_while_shape, new_while_condition, |
| 123 | new_while_body, new_while_init)); |
| 124 | |
| 125 | // We want to get rid of the old while instruction even if it has side |
| 126 | // effecting operations so we do a manual HloComputation::RemoveInstruction |
| 127 | // instead of relying on HloComputation::ReplaceInstruction. |
| 128 | TF_RETURN_IF_ERROR(while_instr->ReplaceAllUsesWith(TupleUtil::ExtractPrefix( |
| 129 | new_while, while_instr->shape().tuple_shapes_size()))); |
| 130 | TF_RETURN_IF_ERROR(containing_computation->RemoveInstruction(while_instr)); |
| 131 | |
| 132 | HloInstruction* while_body_param = new_while_body->parameter_instruction(0); |
| 133 | std::vector<HloInstruction*> live_in_instructions; |
| 134 | for (int64 i = elements_in_old_while_shape; |
| 135 | i < new_while_shape.tuple_shapes_size(); i++) { |
| 136 | live_in_instructions.push_back( |
| 137 | new_while_body->AddInstruction(HloInstruction::CreateGetTupleElement( |
| 138 | instructions[i - elements_in_old_while_shape]->shape(), |
| 139 | while_body_param, i))); |
| 140 | } |
| 141 | |
| 142 | WhileUtil::MakeInstructionsLiveInResult result; |
| 143 | |
| 144 | result.new_while_instr = new_while; |
| 145 | result.while_body_live_in_values = std::move(live_in_instructions); |
| 146 | result.while_body_instruction_map = std::move(inlined_instructions_map); |
| 147 | |
| 148 | return std::move(result); |
| 149 | } |
| 150 | |
| 151 | static StatusOr<std::unique_ptr<HloComputation>> |
| 152 | MakeCountedLoopConditionComputation(const Shape& loop_state_shape, |
nothing calls this directly
no test coverage detected