[llvm] [AMDGPU] Fix pow with a negative base (PR #222248)

[Tim Gymnich via llvm-commits]

================
@@ -3310,6 +3312,63 @@ SDValue AMDGPUTargetLowering::lowerFEXP(SDValue Op, SelectionDAG &DAG) const {
   return R;
 }
 
+// No pow instruction or libcall to fall back on. fmul_legacy returns 0 for a
+// zero operand even against an infinity or a NaN, so pow(x, 0) and pow(1, y)
+// fall out as exp2(0) = 1.
+SDValue AMDGPUTargetLowering::lowerFPOW(SDValue Op, SelectionDAG &DAG) const {
+  EVT VT = Op.getValueType();
+  assert(VT == MVT::f32);
+
+  SDLoc SL(Op);
+  SDValue X = Op.getOperand(0);
+  SDValue Y = Op.getOperand(1);
+  SDNodeFlags Flags = Op->getFlags();
+
+  // Fast expansion: ignores denormals, NaN for a negative base.
+  if (allowApproxFunc(DAG, Flags)) {
+    SDValue Log = DAG.getNode(AMDGPUISD::LOG, SL, VT, X, Flags);
+    SDValue Mul = DAG.getNode(AMDGPUISD::FMUL_LEGACY, SL, VT, Y, Log, Flags);
+    return DAG.getNode(AMDGPUISD::EXP, SL, VT, Mul, Flags);
+  }
+
+  SDValue Abs = DAG.getNode(ISD::FABS, SL, VT, X, Flags);
+  SDValue Log = DAG.getNode(ISD::FLOG2, SL, VT, Abs, Flags);
+  SDValue Mul = DAG.getNode(AMDGPUISD::FMUL_LEGACY, SL, VT, Y, Log, Flags);
+  SDValue R = DAG.getNode(ISD::FEXP2, SL, VT, Mul, Flags);
+
+  EVT SetCCVT = getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(), VT);
+  SDValue One = DAG.getConstantFP(1.0, SL, VT);
+
+  // Infinities count as integers, and every f32 >= 2^24 in magnitude is even.
+  SDValue YTrunc = DAG.getNode(ISD::FTRUNC, SL, VT, Y);
+  SDValue YIsInt = DAG.getSetCC(SL, SetCCVT, YTrunc, Y, ISD::SETOEQ);
+  SDValue YHalf =
+      DAG.getNode(ISD::FMUL, SL, VT, Y, DAG.getConstantFP(0.5, SL, VT));
+  SDValue YHalfTrunc = DAG.getNode(ISD::FTRUNC, SL, VT, YHalf);
+  SDValue YIsOdd =
+      DAG.getNode(ISD::AND, SL, SetCCVT, YIsInt,
+                  DAG.getSetCC(SL, SetCCVT, YHalfTrunc, YHalf, ISD::SETONE));
+
+  // pow(-x, odd y) = -pow(x, y).
+  R = DAG.getNode(ISD::FCOPYSIGN, SL, VT, R,
+                  DAG.getNode(ISD::SELECT, SL, VT, YIsOdd, X, One));
+
+  if (Flags.hasNoNaNs())
+    return R;
+
+  // A negative finite base to a non-integral power is NaN. -inf is excluded:
+  // the core already gives pow(+inf, y).
+  SDValue XNegFinite = DAG.getNode(
+      ISD::IS_FPCLASS, SL, SetCCVT, X,
+      DAG.getTargetConstant(fcNegNormal | fcNegSubnormal, SL, MVT::i32));
----------------
tgymnich wrote:

under `denormal_fpenv(...|preservesign)` subnormals should be treated as zero, thus we cannot include `fcNegSubnormal` here.

https://github.com/llvm/llvm-project/pull/222248