[llvm] [AMDGPU] Fold redundant inf/nan checks into frexp instructions (PR #214936)

[Matt Arsenault via llvm-commits]

================
@@ -18829,6 +18829,131 @@ SDValue SITargetLowering::performClampCombine(SDNode *N,
   return getCanonicalConstantFP(DCI.DAG, SDLoc(N), N->getValueType(0), F);
 }
 
+/// Check if a value is a positive or negative infinity constant.
+static bool isInfinityFPConstant(SDValue V) {
+  auto *CFP = dyn_cast<ConstantFPSDNode>(V);
+  return CFP && CFP->getValueAPF().isInfinity();
+}
+
+SDValue
+SITargetLowering::performFrexpSelectCombine(SDNode *N,
+                                            DAGCombinerInfo &DCI) const {
+  // This optimization only applies when the hardware handles inf/nan correctly.
+  if (Subtarget->hasFractBug())
+    return SDValue();
+
+  SDValue Cond = N->getOperand(0);
+  SDValue TrueVal = N->getOperand(1);
+  SDValue FalseVal = N->getOperand(2);
+
+  // Determine which value is 0 and which might be the frexp result.
+  // Pattern 1: select cond, 0, frexp_result (cond true -> return 0)
+  // Pattern 2: select cond, frexp_result, 0 (cond false -> return 0)
+  // Only check FP zero - frexp returns FP or integer, handle separately.
+  SDValue FrexpVal;
+  bool CondSelectsZero; // If true, condition=true selects zero
+
+  bool TrueIsFPZero = isNullFPConstant(TrueVal);
+  bool FalseIsFPZero = isNullFPConstant(FalseVal);
+  bool TrueIsIntZero = isNullConstant(TrueVal);
+  bool FalseIsIntZero = isNullConstant(FalseVal);
+
+  if (TrueIsFPZero || TrueIsIntZero) {
+    FrexpVal = FalseVal;
+    CondSelectsZero = true;
+  } else if (FalseIsFPZero || FalseIsIntZero) {
+    FrexpVal = TrueVal;
+    CondSelectsZero = false;
+  } else {
+    return SDValue();
+  }
+
+  // Check if FrexpVal comes from amdgcn_frexp_exp or amdgcn_frexp_mant.
+  SDValue FrexpInput;
+  if (!sd_match(FrexpVal, m_IntrinsicWOChain<Intrinsic::amdgcn_frexp_exp>(
+                              m_Value(FrexpInput))) &&
+      !sd_match(FrexpVal, m_IntrinsicWOChain<Intrinsic::amdgcn_frexp_mant>(
+                              m_Value(FrexpInput))))
+    return SDValue();
+
+  // The frexp intrinsics ignore sign, so we can strip sign ops when comparing.
+  SDValue FrexpInputStripped = peekFPSignOps(FrexpInput);
+
+  bool IsNonFiniteTest = false;
+
+  // Handle AMDGPUISD::FP_CLASS or ISD::IS_FPCLASS conditions.
+  // These test specific floating-point classes using a bitmask.
+  if (Cond.getOpcode() == AMDGPUISD::FP_CLASS ||
+      Cond.getOpcode() == ISD::IS_FPCLASS) {
+    SDValue ClassInput = Cond.getOperand(0);
+    SDValue ClassInputStripped = peekFPSignOps(ClassInput);
+
+    if (ClassInputStripped != FrexpInputStripped)
+      return SDValue();
+
+    auto *MaskNode = dyn_cast<ConstantSDNode>(Cond.getOperand(1));
+    if (!MaskNode)
+      return SDValue();
+
+    unsigned Mask = MaskNode->getZExtValue();
+
+    // fcFinite = all finite classes (not inf, not nan)
+    // If the mask tests for finite values and selects frexp when true,
+    // we can fold away the select since frexp returns 0 for non-finite.
+    constexpr unsigned fcFinite =
+        0x1F8; // fcPosNormal|fcNegNormal|fcPosSubnormal|fcNegSubnormal|fcPosZero|fcNegZero
+    constexpr unsigned fcInfNan = 0x207; // fcPosInf|fcNegInf|fcSNan|fcQNan
+
+    if (Mask == fcFinite) {
+      // is_fpclass(x, finite) selects frexp when x is finite
+      // frexp already returns 0 for non-finite, so select frexp, 0 -> frexp
+      IsNonFiniteTest = !CondSelectsZero;
+    } else if (Mask == fcInfNan || Mask == 0x3 || Mask == 0x204) {
+      // is_fpclass(x, inf|nan) or is_fpclass(x, nan) or is_fpclass(x, inf)
+      // selects 0 when x is non-finite
+      IsNonFiniteTest = CondSelectsZero;
+    }
+  } else if (Cond.getOpcode() == ISD::SETCC) {
+    // Handle SETCC conditions for inf/nan tests.
+    ISD::CondCode CC = cast<CondCodeSDNode>(Cond.getOperand(2))->get();
+    SDValue CondLHS = Cond.getOperand(0);
+    SDValue CondRHS = Cond.getOperand(1);
+    SDValue CondLHSStripped = peekFPSignOps(CondLHS);
+
+    if (CC == ISD::SETUO) {
+      // fcmp uno x, y - true if either x or y is NaN
+      SDValue CondRHSStripped = peekFPSignOps(CondRHS);
+      if (CondLHSStripped == FrexpInputStripped ||
+          CondRHSStripped == FrexpInputStripped) {
+        IsNonFiniteTest = CondSelectsZero;
+      }
+    } else if ((CC == ISD::SETOEQ || CC == ISD::SETUEQ) &&
+               isInfinityFPConstant(CondRHS)) {
+      // fcmp oeq/ueq |x|, inf - true if x is inf (or inf/nan for ueq)
+      if (CondLHSStripped == FrexpInputStripped)
+        IsNonFiniteTest = CondSelectsZero;
+    } else if ((CC == ISD::SETONE || CC == ISD::SETUNE) &&
+               isInfinityFPConstant(CondRHS)) {
+      // fcmp one/une |x|, inf - true if x is NOT inf
+      if (CondLHSStripped == FrexpInputStripped)
+        IsNonFiniteTest = !CondSelectsZero;
----------------
arsenm wrote:

This is too confusing of a way to check the compare is exactly fabs 

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