| 28 | using absl::StrCat; |
| 29 | |
| 30 | static StatusOr<HloComputation*> WidenWhileCondition( |
| 31 | HloComputation* narrow_condition, const Shape& wide_shape) { |
| 32 | const Shape& narrow_shape = |
| 33 | narrow_condition->parameter_instruction(0)->shape(); |
| 34 | |
| 35 | HloComputation* wide_while_cond = [&]() { |
| 36 | HloComputation::Builder builder(StrCat("wide.", narrow_condition->name())); |
| 37 | builder.AddInstruction( |
| 38 | HloInstruction::CreateParameter(0, wide_shape, "wide_param")); |
| 39 | |
| 40 | // This is needed so that the root instruction is shaped as a PRED[] -- we |
| 41 | // need to get this right to begin with since we can't mutate the type of |
| 42 | // the root instruction later. We later change the root instruction to |
| 43 | // something more appropriate. |
| 44 | builder.AddInstruction( |
| 45 | HloInstruction::CreateConstant(LiteralUtil::CreateR0<bool>(false))); |
| 46 | return narrow_condition->parent()->AddEmbeddedComputation(builder.Build()); |
| 47 | }(); |
| 48 | |
| 49 | HloInstruction* truncated_parameter = |
| 50 | TupleUtil::ExtractPrefix(wide_while_cond->parameter_instruction(0), |
| 51 | narrow_shape.tuple_shapes_size()); |
| 52 | HloInstruction* call_narrow_cond = wide_while_cond->AddInstruction( |
| 53 | HloInstruction::CreateCall(ShapeUtil::MakeShape(PRED, {}), |
| 54 | {truncated_parameter}, narrow_condition)); |
| 55 | |
| 56 | wide_while_cond->set_root_instruction(call_narrow_cond); |
| 57 | |
| 58 | TF_RETURN_IF_ERROR(CallInliner::Inline(call_narrow_cond).status()); |
| 59 | return wide_while_cond; |
| 60 | } |
| 61 | |
| 62 | static StatusOr<std::pair<HloComputation*, CallInliner::InlinedInstructionMap>> |
| 63 | WidenWhileBody(HloComputation* narrow_body, const Shape& wide_shape) { |
no test coverage detected