[llvm] [InstCombine] Improve coverage of `foldSelectValueEquivalence` (PR #88298)

[Yingwei Zheng via llvm-commits]

================
@@ -1285,40 +1285,97 @@ Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
     Swapped = true;
   }
 
-  // In X == Y ? f(X) : Z, try to evaluate f(Y) and replace the operand.
-  // Make sure Y cannot be undef though, as we might pick different values for
-  // undef in the icmp and in f(Y). Additionally, take care to avoid replacing
-  // X == Y ? X : Z with X == Y ? Y : Z, as that would lead to an infinite
-  // replacement cycle.
   Value *CmpLHS = Cmp.getOperand(0), *CmpRHS = Cmp.getOperand(1);
-  if (TrueVal != CmpLHS &&
-      isGuaranteedNotToBeUndefOrPoison(CmpRHS, SQ.AC, &Sel, &DT)) {
-    if (Value *V = simplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, SQ,
-                                          /* AllowRefinement */ true))
-      // Require either the replacement or the simplification result to be a
-      // constant to avoid infinite loops.
-      // FIXME: Make this check more precise.
-      if (isa<Constant>(CmpRHS) || isa<Constant>(V))
-        return replaceOperand(Sel, Swapped ? 2 : 1, V);
+  auto ReplaceLHSOpWithRHSOp = [&](Value *OldOp,
+                                   Value *NewOp) -> Instruction * {
+    // In X == Y ? f(X) : Z, try to evaluate f(Y) and replace the operand.
+    // Take care to avoid replacing X == Y ? X : Z with X == Y ? Y : Z, as that
+    // would lead to an infinite replacement cycle.
+    // If we will be able to evaluate f(Y) to a constant, we can allow undef,
+    // otherwise Y cannot be undef as we might pick different values for undef
+    // in the icmp and in f(Y).
+    if (TrueVal == OldOp)
+      return nullptr;
+
+    std::optional<bool> IsNeverUndefCached;
+    auto IsNeverUndef = [&](Value *Op) {
+      if (!IsNeverUndefCached.has_value())
+        IsNeverUndefCached =
+            isGuaranteedNotToBeUndefOrPoison(Op, SQ.AC, &Sel, &DT);
+      return *IsNeverUndefCached;
+    };
 
+    if (Value *V = simplifyWithOpReplaced(TrueVal, OldOp, NewOp, SQ,
+                                          /* AllowRefinement */ true)) {
+      // Need some gurantees about the new simplified op to ensure we don't inf
+      // loop.
+      // If we simplify to a constant, replace.
+      bool ShouldReplace = match(V, m_ImmConstant());
+      bool NeedsNoUndef = !ShouldReplace;
+      // Or replace if either NewOp is a constant
+      if (!ShouldReplace && match(NewOp, m_ImmConstant()))
+        ShouldReplace = true;
+      // Or if we end up simplifying f(Y) -> Y i.e: Old & New -> New & New ->
+      // New.
+      if (!ShouldReplace && V == NewOp)
+        ShouldReplace = true;
+
+      // Finally, if we are going to create a new one-use instruction, replace.
+      if (!ShouldReplace && isa<Instruction>(OldOp) && OldOp->hasNUses(2) &&
+          (!isa<Instruction>(NewOp) || !NewOp->hasOneUse()))
+        ShouldReplace = true;
+
+      // Unless we simplify the new instruction to a constant, need to ensure Y
+      // is not undef.
+      if (NeedsNoUndef && ShouldReplace)
+        ShouldReplace = IsNeverUndef(NewOp);
+
+      if (ShouldReplace)
+        return replaceOperand(Sel, Swapped ? 2 : 1, V);
+    }
+    // If we can't simplify, but we will either:
+    //  1) Create a new binop where both ops are NewOp i.e (add x, y) is "worse"
+    //     than (add y, y) in this case, wait until the second call so we don't
+    //     miss a one-use simplification.
+    //  2) Create a new one-use instruction.
+    // proceed.
+    else if (TrueVal->hasOneUse() &&
----------------
dtcxzyw wrote:

I think this `else` can be dropped because we can fall through to this line if the above simplification failed.


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