[llvm] [AMDGPU] Avoid fract fold when result is observable at +/-inf (PR #219515)

[Matt Arsenault via llvm-commits]

================
@@ -2112,7 +2121,71 @@ Value *AMDGPUCodeGenPrepareImpl::matchFractPatNanAvoidant(Value &V) {
                          m_FMinimum(m_Value(Arg0), m_APFloatAllowPoison(C)))))
     return nullptr;
 
-  return matchFractPatImpl(*Arg0, *C);
+  Value *FractArg = matchFractPatImpl(*Arg0, *C);
+  if (!FractArg)
+    return nullptr;
+
+  if (RootToReplace &&
+      !isFractNanAvoidantFoldLegalAtInf(cast<IntrinsicInst>(&V), Arg0, FractArg,
+                                        RootToReplace))
+    return nullptr;
+
+  return FractArg;
+}
+
+/// True if \p U is a select that never picks \p V when \p X is infinity, e.g.
+/// fract's library inf-clamp: select (fcmp oeq (fabs X), +inf), 0.0, V
+static bool isInfDiscardingSelect(User *U, const Value *X, const Value *V) {
+  auto *Sel = dyn_cast<SelectInst>(U);
+  if (!Sel)
+    return false;
+
+  bool IsTrueArm = Sel->getTrueValue() == V;
+  if (!IsTrueArm && Sel->getFalseValue() != V)
+    return false;
+  if (IsTrueArm && Sel->getFalseValue() == V)
+    return false;
+
+  auto *Cmp = dyn_cast<FCmpInst>(Sel->getCondition());
+  if (!Cmp)
+    return false;
+
+  auto [ClassVal, ClassesIfTrue, ClassesIfFalse] =
+      fcmpImpliesClass(Cmp->getPredicate(), *Sel->getFunction(),
+                       Cmp->getOperand(0), Cmp->getOperand(1));
+  if (ClassVal != X)
+    return false;
+
+  return !((IsTrueArm ? ClassesIfTrue : ClassesIfFalse) & fcInf);
+}
+
+/// At +/-inf \p Sub is nan, so \p MinI clamps while fract would return nan;
+/// true if replacing \p Root with fract of \p X is unobservable there. Callers
+/// must not pass a \p Root that can select around \p MinI at +/-inf, since the
+/// poison cases assume \p Root is poison whenever \p MinI is.
+bool AMDGPUCodeGenPrepareImpl::isFractNanAvoidantFoldLegalAtInf(
+    const IntrinsicInst *MinI, const Value *Sub, const Value *X,
+    Instruction *Root) const {
+  // minimum propagates the nan from x - floor(x), exactly like fract.
+  if (MinI->getIntrinsicID() == Intrinsic::minimum)
+    return true;
+
+  // nnan poisons the min at +/-inf; ninf proves nothing here.
+  if (MinI->hasNoNaNs())
+    return true;
+
+  // At +/-inf the fsub is nan with infinite operands, so nnan/ninf poison it.
+  const auto *FPSub = cast<FPMathOperator>(Sub);
+  if (FPSub->hasNoNaNs() || FPSub->hasNoInfs())
+    return true;
+
+  // The result is only observed where an explicit infinity check discards it:
+  //   isinf(x) ? 0.0 : (isnan(x) ? x : minnum(x - floor(x), C))
+  if (all_of(Root->users(),
----------------
arsenm wrote:

Shouldn't need to do a user scan, the patterns are matched from the root to start?

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