| 37 | namespace { |
| 38 | |
| 39 | StatusOr<bool> RunWhileDCE(HloModule* module, HloLivenessAnalysis* liveness) { |
| 40 | bool changed = false; |
| 41 | for (auto* computation : module->computations()) { |
| 42 | for (auto* instruction : computation->instructions()) { |
| 43 | if (instruction->opcode() != HloOpcode::kWhile) { |
| 44 | continue; |
| 45 | } |
| 46 | |
| 47 | const auto* xla_while = instruction; |
| 48 | auto* while_body_comp = xla_while->while_body(); |
| 49 | auto* while_body_param = while_body_comp->parameter_instruction(0); |
| 50 | auto* while_body_root = while_body_comp->root_instruction(); |
| 51 | |
| 52 | if (!xla_while->shape().IsTuple() || |
| 53 | while_body_root->opcode() != HloOpcode::kTuple) { |
| 54 | // Only run DCE on tuple-shaped while loops where body root is Tuple, |
| 55 | // with no I/O instructions. |
| 56 | VLOG(1) << "WhileDCE SKIP while: " << xla_while->ToString(); |
| 57 | continue; |
| 58 | } |
| 59 | |
| 60 | // Remove dead tuple elements. |
| 61 | const int64 tuple_element_count = |
| 62 | ShapeUtil::TupleElementCount(xla_while->shape()); |
| 63 | for (int64 i = 0; i < tuple_element_count; ++i) { |
| 64 | if (liveness->IsLive(xla_while, {i})) { |
| 65 | continue; |
| 66 | } |
| 67 | VLOG(1) << "WhileDCE Dead while tuple element." |
| 68 | << " while: " << xla_while->name() << " tuple_index: " << i; |
| 69 | // Transform while.body computation to make tuple element at |
| 70 | // 'shape_index' as simple pass-through parameter (which candidate |
| 71 | // be removed later by simplification pass). |
| 72 | HloInstruction* pass_thru_gte = while_body_comp->AddInstruction( |
| 73 | HloInstruction::CreateGetTupleElement( |
| 74 | while_body_param->shape().tuple_shapes(i), while_body_param, |
| 75 | i)); |
| 76 | // Replace while.body.root Tuple operand at 'tuple_index' with |
| 77 | // 'pass_thru_gte', making prior operand a dead root (to be cleaned |
| 78 | // up with a subsequent DCE pass). |
| 79 | TF_RETURN_IF_ERROR( |
| 80 | while_body_root->ReplaceOperandWith(i, pass_thru_gte)); |
| 81 | changed = true; |
| 82 | } |
| 83 | } |
| 84 | } |
| 85 | return changed; |
| 86 | } |
| 87 | |
| 88 | } // namespace |
| 89 |
nothing calls this directly
no test coverage detected