| 2796 | } |
| 2797 | |
| 2798 | Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) { |
| 2799 | VLOG(10) << "trying transform [pow(A, 0) => 1]: " << power->ToString(); |
| 2800 | HloInstruction *lhs, *rhs; |
| 2801 | CHECK(Match(power, m::Power(m::Op(&lhs), m::Op(&rhs)))); |
| 2802 | if (IsAll(rhs, 0)) { |
| 2803 | return ReplaceInstruction(power, MakeScalarLike(power, 1)); |
| 2804 | } |
| 2805 | |
| 2806 | VLOG(10) << "trying transform [pow(A, 1) => A]: " << power->ToString(); |
| 2807 | if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(power, lhs)) { |
| 2808 | return Status::OK(); |
| 2809 | } |
| 2810 | |
| 2811 | // pow(exp(A),B) => exp(A*B) |
| 2812 | HloInstruction *a, *b; |
| 2813 | if (Match(power, m::Power(m::Exp(m::Op(&a)), m::Op(&b)))) { |
| 2814 | auto a_times_b = computation_->AddInstruction(HloInstruction::CreateBinary( |
| 2815 | power->shape(), HloOpcode::kMultiply, a, b)); |
| 2816 | return ReplaceWithNewInstruction( |
| 2817 | power, HloInstruction::CreateUnary(power->shape(), HloOpcode::kExp, |
| 2818 | a_times_b)); |
| 2819 | } |
| 2820 | |
| 2821 | VLOG(10) << "trying transform [pow(A, 2) => A*A]: " << power->ToString(); |
| 2822 | if (IsAll(rhs, 2)) { |
| 2823 | return ReplaceWithNewInstruction( |
| 2824 | power, HloInstruction::CreateBinary(power->shape(), |
| 2825 | HloOpcode::kMultiply, lhs, lhs)); |
| 2826 | } |
| 2827 | |
| 2828 | // Pow(A, 3) is used in GELU. |
| 2829 | VLOG(10) << "trying transform [pow(A, 3) => A*A*A]: " << power->ToString(); |
| 2830 | if (IsAll(rhs, 3)) { |
| 2831 | HloInstruction * tmp = computation_->AddInstruction( |
| 2832 | HloInstruction::CreateBinary( |
| 2833 | power->shape(), HloOpcode::kMultiply, lhs, lhs)); |
| 2834 | return ReplaceWithNewInstruction( |
| 2835 | power, HloInstruction::CreateBinary(power->shape(), |
| 2836 | HloOpcode::kMultiply, lhs, tmp)); |
| 2837 | } |
| 2838 | |
| 2839 | VLOG(10) << "trying transform [pow(A, -1) => 1/A]: " << power->ToString(); |
| 2840 | if (IsAll(rhs, -1)) { |
| 2841 | return ReplaceWithNewInstruction( |
| 2842 | power, HloInstruction::CreateBinary(power->shape(), HloOpcode::kDivide, |
| 2843 | MakeScalarLike(lhs, 1), lhs)); |
| 2844 | } |
| 2845 | |
| 2846 | VLOG(10) << "trying transform [pow(pow(A, X), Y) => pow(A, X*Y)]: " |
| 2847 | << power->ToString(); |
| 2848 | |
| 2849 | // Don't perform this optimization if either of the exponents is complex; this |
| 2850 | // identity is true only for real-valued exponents. In addition, we cowardly |
| 2851 | // refuse to do this transformation if the two exponents have different |
| 2852 | // element types. |
| 2853 | if (lhs->opcode() == HloOpcode::kPower && |
| 2854 | !ShapeUtil::ElementIsComplex(lhs->operand(1)->shape()) && |
| 2855 | !ShapeUtil::ElementIsComplex(rhs->shape()) && |
no test coverage detected