| 34 | : exclude_entry_computation_(exclude_entry_computation) {} |
| 35 | |
| 36 | StatusOr<bool> TupleSimplifier::Run(HloModule* module) { |
| 37 | // Initially add all GTE and Tuple instructions to the worklist. |
| 38 | bool changed = false; |
| 39 | for (auto* computation : module->computations()) { |
| 40 | if (exclude_entry_computation_ && |
| 41 | computation == module->entry_computation()) { |
| 42 | continue; |
| 43 | } |
| 44 | for (auto* instruction : computation->MakeInstructionPostOrder()) { |
| 45 | if (instruction->opcode() == HloOpcode::kTuple) { |
| 46 | // Collapse the following structure into just 'Tuple-shaped Op': |
| 47 | // |
| 48 | // Tuple-shaped Op |
| 49 | // | |
| 50 | // +-----+-----+ |
| 51 | // | | | |
| 52 | // GTE GTE GTE |
| 53 | // | | | |
| 54 | // +-----+-----+ |
| 55 | // | |
| 56 | // Tuple |
| 57 | // |
| 58 | HloInstruction* top_tuple = nullptr; |
| 59 | bool can_simplify = true; |
| 60 | for (int64 operand_number = 0; |
| 61 | operand_number < instruction->operand_count(); ++operand_number) { |
| 62 | HloInstruction* operand = |
| 63 | instruction->mutable_operand(operand_number); |
| 64 | if (operand->opcode() != HloOpcode::kGetTupleElement || |
| 65 | operand->tuple_index() != operand_number) { |
| 66 | can_simplify = false; |
| 67 | break; |
| 68 | } |
| 69 | if (top_tuple == nullptr) { |
| 70 | top_tuple = operand->mutable_operand(0); |
| 71 | if (!ShapeUtil::Compatible(top_tuple->shape(), |
| 72 | instruction->shape())) { |
| 73 | can_simplify = false; |
| 74 | break; |
| 75 | } |
| 76 | } else if (top_tuple != operand->operand(0)) { |
| 77 | can_simplify = false; |
| 78 | break; |
| 79 | } |
| 80 | } |
| 81 | if (can_simplify && top_tuple != nullptr) { |
| 82 | changed = true; |
| 83 | TF_RETURN_IF_ERROR( |
| 84 | computation->ReplaceInstruction(instruction, top_tuple)); |
| 85 | } |
| 86 | } else { |
| 87 | auto ancestor = instruction->LatestNonGteAncestorAndIndex(); |
| 88 | if (ancestor.first == instruction) { |
| 89 | continue; |
| 90 | } |
| 91 | // If possible replace a chain of GTE with the operation which produces |
| 92 | // the element. For example, replace uses of GTE with below with just |
| 93 | // 'Op' (assuming 'Op' is at the index of the GTE instruction): |
nothing calls this directly
no test coverage detected