[DAGCombine] Optimize pow(X, (2/3)) and Pow(X,(3/2))

[GGanesh]

Summary

Try to convert pow(X, 1.50) --> X * sqrt(X) and Pow(X,(2/3)) to cube root of (X*X).

[lebedev.ri]

llvm/test/CodeGen/X86/pow.150.ll

1

Please precommit the test

2–3

Why not just check the final asm?

[lebedev.ri]

This review may be stuck/dead, consider abandoning if no longer relevant.
Removing myself as reviewer in attempt to clean dashboard.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 }
 
 SDValue DAGCombiner::visitFPOW(SDNode *N) {
   ConstantFPSDNode *ExponentC = isConstOrConstSplatFP(N->getOperand(1));
   if (!ExponentC)
     return SDValue();
   SelectionDAG::FlagInserter FlagsInserter(DAG, N);
 
-  // Try to convert x ** (1/3) into cube root.
+  // Try to convert 
+  //   a. x ** (1/3) into cube root.
+  //   b. x ** (2/3) into cube root of (x*x).
   // TODO: Handle the various flavors of long double.
   // TODO: Since we're approximating, we don't need an exact 1/3 exponent.
   //       Some range near 1/3 should be fine.
   EVT VT = N->getValueType(0);
-  if ((VT == MVT::f32 && ExponentC->getValueAPF().isExactlyValue(1.0f/3.0f)) ||
-      (VT == MVT::f64 && ExponentC->getValueAPF().isExactlyValue(1.0/3.0))) {
+
+  bool Exponent1by3 = 
+      ((VT == MVT::f32 && ExponentC->getValueAPF().isExactlyValue(1.0f/3.0f)) ||
+      (VT == MVT::f64 && ExponentC->getValueAPF().isExactlyValue(1.0/3.0)));
+
+  bool Exponent2by3 = 
+       ((VT == MVT::f32 && ExponentC->getValueAPF().isExactlyValue(2.0f/3.0f)) ||
+       (VT == MVT::f64 && ExponentC->getValueAPF().isExactlyValue(2.0/3.0)));
+
+  if (Exponent1by3 || Exponent2by3) {
     // pow(-0.0, 1/3) = +0.0; cbrt(-0.0) = -0.0.
     // pow(-inf, 1/3) = +inf; cbrt(-inf) = -inf.
     // pow(-val, 1/3) =  nan; cbrt(-val) = -num.
     // For regular numbers, rounding may cause the results to differ.
     // Therefore, we require { nsz ninf nnan afn } for this transform.
     // TODO: We could select out the special cases if we don't have nsz/ninf.
     SDNodeFlags Flags = N->getFlags();
     if (!Flags.hasNoSignedZeros() || !Flags.hasNoInfs() || !Flags.hasNoNaNs() ||
         !Flags.hasApproximateFuncs())
       return SDValue();
 
     // Do not create a cbrt() libcall if the target does not have it, and do not
     // turn a pow that has lowering support into a cbrt() libcall.
     if (!DAG.getLibInfo().has(LibFunc_cbrt) ||
         (!DAG.getTargetLoweringInfo().isOperationExpand(ISD::FPOW, VT) &&
          DAG.getTargetLoweringInfo().isOperationExpand(ISD::FCBRT, VT)))
       return SDValue();
 
-    return DAG.getNode(ISD::FCBRT, SDLoc(N), VT, N->getOperand(0));
+    SDLoc DL(N);
+    SDValue cbrt =  DAG.getNode(ISD::FCBRT, DL, VT, N->getOperand(0));
+    if (Exponent1by3)
+      return cbrt;
+    return DAG.getNode(ISD::FMUL, DL, VT, cbrt, cbrt);
   }
 
-  // Try to convert x ** (1/4) and x ** (3/4) into square roots.
+  // Try to convert the following into square roots.
+  //   a. x ** (1/4) 
+  //   b. x ** (3/4) 
+  //   c. x ** (3/2)
   // x ** (1/2) is canonicalized to sqrt, so we do not bother with that case.
   // TODO: This could be extended (using a target hook) to handle smaller
   // power-of-2 fractional exponents.
   bool ExponentIs025 = ExponentC->getValueAPF().isExactlyValue(0.25);
   bool ExponentIs075 = ExponentC->getValueAPF().isExactlyValue(0.75);
-  if (ExponentIs025 || ExponentIs075) {
+  bool ExponentIs150 = ExponentC->getValueAPF().isExactlyValue(1.50);
+  if (ExponentIs025 || ExponentIs075 || ExponentIs150) {
     // pow(-0.0, 0.25) = +0.0; sqrt(sqrt(-0.0)) = -0.0.
     // pow(-inf, 0.25) = +inf; sqrt(sqrt(-inf)) =  NaN.
     // pow(-0.0, 0.75) = +0.0; sqrt(-0.0) * sqrt(sqrt(-0.0)) = +0.0.
     // pow(-inf, 0.75) = +inf; sqrt(-inf) * sqrt(sqrt(-inf)) =  NaN.
     // For regular numbers, rounding may cause the results to differ.
     // Therefore, we require { nsz ninf afn } for this transform.
     // TODO: We could select out the special cases if we don't have nsz/ninf.
     SDNodeFlags Flags = N->getFlags();
if (ExponentIs025 || ExponentIs075 || ExponentIs150) {
     // Assume that libcalls are the smallest code.
     // TODO: This restriction should probably be lifted for vectors.
     if (ForCodeSize)
       return SDValue();
 
     // pow(X, 0.25) --> sqrt(sqrt(X))
     SDLoc DL(N);
     SDValue Sqrt = DAG.getNode(ISD::FSQRT, DL, VT, N->getOperand(0));
+
+    // pow(X, 1.50) --> X * sqrt(X)
+    if (ExponentIs150)
+      return DAG.getNode(ISD::FMUL, DL, VT, N->getOperand(0), Sqrt);
+
     SDValue SqrtSqrt = DAG.getNode(ISD::FSQRT, DL, VT, Sqrt);
     if (ExponentIs025)
       return SqrtSqrt;
     // pow(X, 0.75) --> sqrt(X) * sqrt(sqrt(X))
     return DAG.getNode(ISD::FMUL, DL, VT, Sqrt, SqrtSqrt);
   }
 
   return SDValue();

llvm/test/CodeGen/X86/pow.150.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-- -debug 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+declare float @llvm.pow.f32(float, float)
+declare <4 x float> @llvm.pow.v4f32(<4 x float>, <4 x float>)
+declare double @llvm.pow.f64(double, double)
+declare <2 x double> @llvm.pow.v2f64(<2 x double>, <2 x double>)
+
+define float @pow_f32_three_halves_fmf(float %x) nounwind {
+; CHECK: Combining: {{.*}}: f32 = fpow ninf nsz afn [[X:t[0-9]+]], ConstantFP:f32<1.500000e+00>
+; CHECK-NEXT: Creating new node: [[SQRT:t[0-9]+]]: f32 = fsqrt ninf nsz afn [[X]]
+; CHECK-NEXT: Creating new node: [[R:t[0-9]+]]: f32 = fmul ninf nsz afn [[X]], [[SQRT]]
+  %r = call nsz ninf afn float @llvm.pow.f32(float %x, float 1.5e-00)
+  ret float %r
+}
+
+define double @pow_f64_three_halves_fmf(double %x) nounwind {
+; CHECK: Combining: {{.*}}: f64 = fpow ninf nsz afn [[X:t[0-9]+]], ConstantFP:f64<1.500000e+00>
+; CHECK-NEXT: Creating new node: [[SQRT:t[0-9]+]]: f64 = fsqrt ninf nsz afn [[X]]
+; CHECK-NEXT: Creating new node: [[R:t[0-9]+]]: f64 = fmul ninf nsz afn [[X]], [[SQRT]]
+  %r = call nsz ninf afn double @llvm.pow.f64(double %x, double 1.5e-00)
+  ret double %r
+}
+
+define <4 x float> @pow_v4f32_three_halves_fmf(<4 x float> %x) nounwind {
+; CHECK: Combining: {{.*}}: v4f32 = fpow nnan ninf nsz arcp contract afn reassoc [[X:t[0-9]+]], {{.*}}
+; CHECK-NEXT: Creating new node: [[SQRT:t[0-9]+]]: v4f32 = fsqrt nnan ninf nsz arcp contract afn reassoc [[X]]
+; CHECK-NEXT: Creating new node: [[R:t[0-9]+]]: v4f32 = fmul nnan ninf nsz arcp contract afn reassoc [[X]], [[SQRT]]
+; CHECK-NEXT:  ... into: [[R]]: v4f32 = fmul nnan ninf nsz arcp contract afn reassoc [[X]], [[SQRT]]
+  %r = call fast <4 x float> @llvm.pow.v4f32(<4 x float> %x, <4 x float> <float 1.5e+00, float 1.5e+00, float 1.5e+00, float 1.5e+00>)
+  ret <4 x float> %r
+}
+
+define <2 x double> @pow_v2f64_three_halves_fmf(<2 x double> %x) nounwind {
+; CHECK: Combining: {{.*}}: v2f64 = fpow nnan ninf nsz arcp contract afn reassoc [[X:t[0-9]+]], {{.*}}
+; CHECK-NEXT: Creating new node: [[SQRT:t[0-9]+]]: v2f64 = fsqrt nnan ninf nsz arcp contract afn reassoc [[X]]
+; CHECK-NEXT: Creating new node: [[R:t[0-9]+]]: v2f64 = fmul nnan ninf nsz arcp contract afn reassoc [[X]], [[SQRT]]
+; CHECK-NEXT:  ... into: [[R]]: v2f64 = fmul nnan ninf nsz arcp contract afn reassoc [[X]], [[SQRT]]
+  %r = call fast <2 x double> @llvm.pow.v2f64(<2 x double> %x, <2 x double> <double 1.5e+0, double 1.5e+0>)
+  ret <2 x double> %r
+}