| 60 | } |
| 61 | |
| 62 | static StatusOr<std::pair<HloComputation*, CallInliner::InlinedInstructionMap>> |
| 63 | WidenWhileBody(HloComputation* narrow_body, const Shape& wide_shape) { |
| 64 | const Shape& narrow_shape = narrow_body->parameter_instruction(0)->shape(); |
| 65 | |
| 66 | HloComputation* wide_while_body = [&]() { |
| 67 | HloComputation::Builder builder(StrCat("wide.", narrow_body->name())); |
| 68 | builder.AddInstruction( |
| 69 | HloInstruction::CreateParameter(0, wide_shape, "wide_param")); |
| 70 | return narrow_body->parent()->AddEmbeddedComputation(builder.Build()); |
| 71 | }(); |
| 72 | |
| 73 | HloInstruction* wide_parameter = wide_while_body->parameter_instruction(0); |
| 74 | HloInstruction* truncated_parameter = TupleUtil::ExtractPrefix( |
| 75 | wide_parameter, narrow_shape.tuple_shapes_size()); |
| 76 | HloInstruction* call_narrow_body = |
| 77 | wide_while_body->AddInstruction(HloInstruction::CreateCall( |
| 78 | narrow_shape, {truncated_parameter}, narrow_body)); |
| 79 | |
| 80 | std::vector<HloInstruction*> live_through_values; |
| 81 | for (int i = narrow_shape.tuple_shapes_size(); |
| 82 | i < wide_shape.tuple_shapes_size(); i++) { |
| 83 | live_through_values.push_back( |
| 84 | wide_while_body->AddInstruction(HloInstruction::CreateGetTupleElement( |
| 85 | wide_shape.tuple_shapes(i), wide_parameter, i))); |
| 86 | } |
| 87 | |
| 88 | wide_while_body->set_root_instruction( |
| 89 | TupleUtil::AppendSuffix(call_narrow_body, live_through_values)); |
| 90 | |
| 91 | TF_ASSIGN_OR_RETURN(auto inlined_instructions_map, |
| 92 | CallInliner::Inline(call_narrow_body)); |
| 93 | return {{wide_while_body, std::move(inlined_instructions_map)}}; |
| 94 | } |
| 95 | |
| 96 | /*static*/ StatusOr<WhileUtil::MakeInstructionsLiveInResult> |
| 97 | WhileUtil::MakeInstructionsLiveIn( |
no test coverage detected