| 4086 | } // namespace |
| 4087 | |
| 4088 | Status AlgebraicSimplifierVisitor::HandleTranspose(HloInstruction* transpose) { |
| 4089 | auto operand = transpose->mutable_operand(0); |
| 4090 | if (std::is_sorted(transpose->dimensions().begin(), |
| 4091 | transpose->dimensions().end())) { |
| 4092 | VLOG(10) << "deleting no-op transpose"; |
| 4093 | return ReplaceInstruction(transpose, operand); |
| 4094 | } |
| 4095 | |
| 4096 | if (HloOpcode::kTranspose == operand->opcode()) { |
| 4097 | return ReplaceWithNewInstruction( |
| 4098 | transpose, HloInstruction::CreateTranspose( |
| 4099 | transpose->shape(), operand->mutable_operand(0), |
| 4100 | ComposePermutations(operand->dimensions(), |
| 4101 | transpose->dimensions()))); |
| 4102 | } |
| 4103 | |
| 4104 | // Convert transpose(dot(a,b)) to dot(b,a). |
| 4105 | if (operand->opcode() == HloOpcode::kDot && operand->user_count() == 1 && |
| 4106 | operand->shape().rank() == 2) { |
| 4107 | TF_ASSIGN_OR_RETURN(bool did_transform, [&]() -> StatusOr<bool> { |
| 4108 | const auto& dnums = operand->dot_dimension_numbers(); |
| 4109 | if (dnums.lhs_batch_dimensions_size() != 0) { |
| 4110 | return false; |
| 4111 | } |
| 4112 | HloInstruction* lhs = operand->mutable_operand(0); |
| 4113 | if (lhs->shape().rank() != 1 + dnums.lhs_contracting_dimensions_size()) { |
| 4114 | return false; |
| 4115 | } |
| 4116 | HloInstruction* rhs = operand->mutable_operand(1); |
| 4117 | if (rhs->shape().rank() != 1 + dnums.rhs_contracting_dimensions_size()) { |
| 4118 | return false; |
| 4119 | } |
| 4120 | DotDimensionNumbers new_dnums; |
| 4121 | *new_dnums.mutable_lhs_contracting_dimensions() = |
| 4122 | dnums.rhs_contracting_dimensions(); |
| 4123 | *new_dnums.mutable_rhs_contracting_dimensions() = |
| 4124 | dnums.lhs_contracting_dimensions(); |
| 4125 | TF_RETURN_IF_ERROR(ReplaceWithNewInstruction( |
| 4126 | transpose, HloInstruction::CreateDot(transpose->shape(), /*lhs=*/rhs, |
| 4127 | /*rhs=*/lhs, new_dnums, |
| 4128 | operand->precision_config()))); |
| 4129 | return true; |
| 4130 | }()); |
| 4131 | if (did_transform) { |
| 4132 | return Status::OK(); |
| 4133 | } |
| 4134 | } |
| 4135 | |
| 4136 | // Replace transpose with a reshape if more than one degenerate method is |
| 4137 | // permuted. |
| 4138 | if (OnlyPermutesDegenerateDims(transpose->shape(), transpose->dimensions())) { |
| 4139 | return ReplaceWithNewInstruction( |
| 4140 | transpose, HloInstruction::CreateReshape( |
| 4141 | transpose->shape(), transpose->mutable_operand(0))); |
| 4142 | } |
| 4143 | |
| 4144 | if (operand->opcode() == HloOpcode::kRng && operand->user_count() == 1) { |
| 4145 | *operand->mutable_shape() = transpose->shape(); |
nothing calls this directly
no test coverage detected