static*/
| 235 | } |
| 236 | |
| 237 | /*static*/ StatusOr<WhileUtil::LoopStateTy> WhileUtil::MakeCountedLoop( |
| 238 | HloComputation* computation, int32 trip_count, |
| 239 | const WhileUtil::LoopStateTy& init_values, |
| 240 | const WhileUtil::LoopBodyGeneratorTy& loop_body_generator, |
| 241 | const OpMetadata& metadata) { |
| 242 | CHECK_GE(trip_count, 0); |
| 243 | |
| 244 | // Both MakeCountedLoopConditionComputation and MakeCountedLoopBodyComputation |
| 245 | // use loop_state_shape to create a literal, which requires loop_state_shape |
| 246 | // to have a layout. |
| 247 | Shape loop_state_shape = MakeLoopStateShapeWithLayout(init_values); |
| 248 | TF_ASSIGN_OR_RETURN( |
| 249 | std::unique_ptr<HloComputation> cond, |
| 250 | MakeCountedLoopConditionComputation(loop_state_shape, trip_count)); |
| 251 | TF_ASSIGN_OR_RETURN( |
| 252 | std::unique_ptr<HloComputation> body, |
| 253 | MakeCountedLoopBodyComputation(loop_state_shape, loop_body_generator)); |
| 254 | TF_ASSIGN_OR_RETURN(HloInstruction * init_tuple, |
| 255 | MakeInitTupleFromInitValues(computation, init_values)); |
| 256 | HloModule* module = computation->parent(); |
| 257 | HloInstruction* while_instr = |
| 258 | computation->AddInstruction(HloInstruction::CreateWhile( |
| 259 | loop_state_shape, module->AddEmbeddedComputation(std::move(cond)), |
| 260 | module->AddEmbeddedComputation(std::move(body)), init_tuple)); |
| 261 | while_instr->set_metadata(metadata); |
| 262 | |
| 263 | std::vector<HloInstruction*> result; |
| 264 | for (int64 i = 0, e = init_values.size(); i < e; i++) { |
| 265 | TF_ASSIGN_OR_RETURN(HloInstruction * user_state, |
| 266 | MakeGetTupleElementHlo(while_instr, i + 1)); |
| 267 | result.push_back(user_state); |
| 268 | } |
| 269 | return result; |
| 270 | } |
| 271 | |
| 272 | /*static*/ std::vector<HloInstruction*> WhileUtil::GetInvariantGTEsForWhileBody( |
| 273 | const HloComputation& while_body) { |
nothing calls this directly
no test coverage detected