This function is called once we've decided to sink reshape/transpose operands across an instruction. It returns an updated `operand` with a shape that plays nicely with `new_operand_shape`; it has the same shape (of the correct type).
| 135 | // plays nicely with `new_operand_shape`; it has the same shape (of the |
| 136 | // correct type). |
| 137 | HloInstruction* UpdateOperand(const HloInstruction* first_reshape_operand, |
| 138 | const Shape& new_operand_shape, |
| 139 | HloInstruction* operand) { |
| 140 | HloComputation* computation = operand->parent(); |
| 141 | const PrimitiveType element_type = operand->shape().element_type(); |
| 142 | const Shape new_shape = |
| 143 | ShapeUtil::ChangeElementType(new_operand_shape, element_type); |
| 144 | |
| 145 | switch (operand->opcode()) { |
| 146 | case HloOpcode::kConstant: { |
| 147 | if (first_reshape_operand->opcode() == HloOpcode::kReshape) { |
| 148 | VLOG(5) << "Adding reshape to kConstant operand"; |
| 149 | return computation->AddInstruction( |
| 150 | HloInstruction::CreateReshape(new_shape, operand)); |
| 151 | } else { |
| 152 | CHECK(first_reshape_operand->opcode() == HloOpcode::kTranspose); |
| 153 | VLOG(5) << "Adding transpose to kConstant operand"; |
| 154 | std::vector<int64> inverse_permutation = |
| 155 | InversePermutation(first_reshape_operand->dimensions()); |
| 156 | return computation->AddInstruction(HloInstruction::CreateTranspose( |
| 157 | new_shape, operand, inverse_permutation)); |
| 158 | } |
| 159 | } |
| 160 | case HloOpcode::kRng: { |
| 161 | CHECK_EQ(operand->user_count(), 1); |
| 162 | VLOG(5) << "Cloning kRng operand with new shape"; |
| 163 | return computation->AddInstruction( |
| 164 | operand->CloneWithNewOperands(new_shape, operand->operands())); |
| 165 | } |
| 166 | case HloOpcode::kReshape: |
| 167 | case HloOpcode::kTranspose: { |
| 168 | VLOG(5) << "Using existing operand of kReshape or kTranspose"; |
| 169 | return operand->mutable_operand(0); |
| 170 | } |
| 171 | case HloOpcode::kBroadcast: { |
| 172 | CHECK(ShapeUtil::IsScalar(operand->operand(0)->shape())); |
| 173 | HloInstruction* inst = computation->AddInstruction( |
| 174 | operand->CloneWithNewOperands(new_shape, operand->operands())); |
| 175 | VLOG(5) << "Changing broadcast from " << operand->ToString() << " to " |
| 176 | << inst->ToString(); |
| 177 | return inst; |
| 178 | } |
| 179 | |
| 180 | default: |
| 181 | LOG(FATAL) << "Unexpected operand opcode during update: " << operand; |
| 182 | } |
| 183 | } |
| 184 | |
| 185 | // Actually performs the reshape-move transformation -- that is, sinks the |
| 186 | // reshape or transpose operands of `instruction` across it. |
no test coverage detected