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

via llvm-commits llvm-commits at lists.llvm.org
Wed Jun 5 12:34:00 PDT 2024


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

>From c9c544c4c4dbb445a94f24e9b06947421eb08f3b Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Wed, 10 Apr 2024 12:07:37 -0500
Subject: [PATCH 1/3] [InstCombine] Add tests for expanding
 `foldSelectValueEquivalence`; NFC

---
 .../InstCombine/select-cmp-eq-op-fold.ll      | 186 ++++++++++++++++++
 1 file changed, 186 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll

diff --git a/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll b/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll
new file mode 100644
index 0000000000000..ec82b1944f723
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll
@@ -0,0 +1,186 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+declare void @use.i1(i1)
+declare void @use.i8(i8)
+define i8 @replace_with_y_noundef(i8 %x, i8 noundef %y, i8 %z) {
+; CHECK-LABEL: @replace_with_y_noundef(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[AND]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %cmp = icmp eq i8 %x, %y
+  %and = and i8 %x, %y
+  %sel = select i1 %cmp, i8 %and, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_x_noundef(i8 noundef %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @replace_with_x_noundef(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    call void @use.i1(i1 [[CMP]])
+; CHECK-NEXT:    [[AND:%.*]] = or i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Z:%.*]], i8 [[AND]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %cmp = icmp ne i8 %x, %y
+  call void @use.i1(i1 %cmp)
+  %and = or i8 %x, %y
+  %sel = select i1 %cmp, i8 %z, i8 %and
+  ret i8 %sel
+}
+
+define i8 @replace_with_x_maybe_undef_fail(i8 %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @replace_with_x_maybe_undef_fail(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    call void @use.i1(i1 [[CMP]])
+; CHECK-NEXT:    [[AND:%.*]] = or i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Z:%.*]], i8 [[AND]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %cmp = icmp ne i8 %x, %y
+  call void @use.i1(i1 %cmp)
+  %and = or i8 %x, %y
+  %sel = select i1 %cmp, i8 %z, i8 %and
+  ret i8 %sel
+}
+
+define i8 @replace_with_y_for_new_oneuse(i8 noundef %xx, i8 noundef %y, i8 %z) {
+; CHECK-LABEL: @replace_with_y_for_new_oneuse(
+; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %x = mul i8 %xx, 13
+  %cmp = icmp eq i8 %x, %y
+  %add = add nuw i8 %x, %y
+  %sel = select i1 %cmp, i8 %add, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_y_for_new_oneuse2(i8 %xx, i8 noundef %y, i8 %z, i8 %q) {
+; CHECK-LABEL: @replace_with_y_for_new_oneuse2(
+; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X]], [[Q:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %x = mul i8 %xx, 13
+  %cmp = icmp eq i8 %x, %y
+  %add = add nuw i8 %x, %q
+  %sel = select i1 %cmp, i8 %add, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_x_for_new_oneuse(i8 noundef %xx, i8 noundef %yy, i8 %z, i8 %w) {
+; CHECK-LABEL: @replace_with_x_for_new_oneuse(
+; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
+; CHECK-NEXT:    [[Y:%.*]] = add i8 [[YY:%.*]], [[W:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %x = mul i8 %xx, 13
+  %y = add i8 %yy, %w
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul i8 %x, %y
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_x_for_new_oneuse2(i8 noundef %xx, i8 %yy, i8 %z, i8 %w, i8 %q) {
+; CHECK-LABEL: @replace_with_x_for_new_oneuse2(
+; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
+; CHECK-NEXT:    [[Y:%.*]] = add i8 [[YY:%.*]], [[W:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[Y]], [[Q:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %x = mul i8 %xx, 13
+  %y = add i8 %yy, %w
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul i8 %q, %y
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_x_for_simple_binop(i8 noundef %xx, i8 %yy, i8 %z, i8 %w) {
+; CHECK-LABEL: @replace_with_x_for_simple_binop(
+; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
+; CHECK-NEXT:    [[Y:%.*]] = add i8 [[YY:%.*]], [[W:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[Y]]
+; CHECK-NEXT:    call void @use.i8(i8 [[Y]])
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %x = mul i8 %xx, 13
+  %y = add i8 %yy, %w
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul i8 %x, %y
+  call void @use.i8(i8 %y)
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_none_for_new_oneuse_fail_maybe_undef(i8 %xx, i8 %y, i8 %z) {
+; CHECK-LABEL: @replace_with_none_for_new_oneuse_fail_maybe_undef(
+; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y:%.*]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %x = mul i8 %xx, 13
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul i8 %x, %y
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_y_for_simple_binop(i8 %x, i8 noundef %y, i8 %z) {
+; CHECK-LABEL: @replace_with_y_for_simple_binop(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul nsw i8 %x, %y
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  ret i8 %sel
+}
+
+define i8 @replace_with_y_for_simple_binop_fail_multiuse(i8 %x, i8 noundef %y, i8 %z) {
+; CHECK-LABEL: @replace_with_y_for_simple_binop_fail_multiuse(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    call void @use.i8(i8 [[MUL]])
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul nsw i8 %x, %y
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  call void @use.i8(i8 %mul)
+  ret i8 %sel
+}
+
+define i8 @replace_with_y_for_simple_binop_fail(i8 %x, i8 noundef %y, i8 %z, i8 %q) {
+; CHECK-LABEL: @replace_with_y_for_simple_binop_fail(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i8 [[X]], [[Q:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
+; CHECK-NEXT:    ret i8 [[SEL]]
+;
+  %cmp = icmp eq i8 %x, %y
+  %mul = mul nsw i8 %x, %q
+  %sel = select i1 %cmp, i8 %mul, i8 %z
+  ret i8 %sel
+}

>From a3c3705e3f1c0e51cd09ae4ddde596c7273926d2 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Wed, 5 Jun 2024 11:30:53 -0500
Subject: [PATCH 2/3] [InstCombine] Improve coverage of
 `foldSelectValueEquivalence` for constants

We don't need the `noundef` check if the new simplification is a
constant.

This cleans up regressions from folding multiuse:
    `(icmp eq/ne (sub/xor x, y), 0)` -> `(icmp eq/ne x, y)`.
---
 .../InstCombine/InstCombineSelect.cpp         | 57 +++++++++++--------
 llvm/test/Transforms/InstCombine/abs-1.ll     |  7 +--
 llvm/test/Transforms/InstCombine/select.ll    |  6 +-
 3 files changed, 38 insertions(+), 32 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index a9aba550e3853..0eaec485b853f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -1288,22 +1288,33 @@ 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))
+  auto ReplaceOldOpWithNewOp = [&](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;
+
+    if (Value *V = simplifyWithOpReplaced(TrueVal, OldOp, NewOp, SQ,
+                                          /* AllowRefinement= */ true)) {
+      // Need some guarantees about the new simplified op to ensure we don't inf
+      // loop.
+      // If we simplify to a constant, replace.
+      if (match(V, m_ImmConstant()))
         return replaceOperand(Sel, Swapped ? 2 : 1, V);
 
+      // If NewOp is a constant and OldOp is not replace iff NewOp doesn't
+      // contain and undef/poison elements.
+      if (match(NewOp, m_ImmConstant()) &&
+          isGuaranteedNotToBeUndefOrPoison(NewOp, SQ.AC, &Sel, &DT))
+        return replaceOperand(Sel, Swapped ? 2 : 1, V);
+    }
+
     // Even if TrueVal does not simplify, we can directly replace a use of
     // CmpLHS with CmpRHS, as long as the instruction is not used anywhere
     // else and is safe to speculatively execute (we may end up executing it
@@ -1311,17 +1322,17 @@ Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
     // undefined behavior). Only do this if CmpRHS is a constant, as
     // profitability is not clear for other cases.
     // FIXME: Support vectors.
-    if (match(CmpRHS, m_ImmConstant()) && !match(CmpLHS, m_ImmConstant()) &&
-        !Cmp.getType()->isVectorTy())
-      if (replaceInInstruction(TrueVal, CmpLHS, CmpRHS))
+    if (OldOp == CmpLHS && match(NewOp, m_ImmConstant()) &&
+        !match(OldOp, m_ImmConstant()) && !Cmp.getType()->isVectorTy())
+      if (replaceInInstruction(TrueVal, OldOp, NewOp))
         return &Sel;
-  }
-  if (TrueVal != CmpRHS &&
-      isGuaranteedNotToBeUndefOrPoison(CmpLHS, SQ.AC, &Sel, &DT))
-    if (Value *V = simplifyWithOpReplaced(TrueVal, CmpRHS, CmpLHS, SQ,
-                                          /* AllowRefinement */ true))
-      if (isa<Constant>(CmpLHS) || isa<Constant>(V))
-        return replaceOperand(Sel, Swapped ? 2 : 1, V);
+    return nullptr;
+  };
+
+  if (Instruction *R = ReplaceOldOpWithNewOp(CmpLHS, CmpRHS))
+    return R;
+  if (Instruction *R = ReplaceOldOpWithNewOp(CmpRHS, CmpLHS))
+    return R;
 
   auto *FalseInst = dyn_cast<Instruction>(FalseVal);
   if (!FalseInst)
diff --git a/llvm/test/Transforms/InstCombine/abs-1.ll b/llvm/test/Transforms/InstCombine/abs-1.ll
index 32bd7a37053ed..0cf7cd97d8ff4 100644
--- a/llvm/test/Transforms/InstCombine/abs-1.ll
+++ b/llvm/test/Transforms/InstCombine/abs-1.ll
@@ -852,11 +852,8 @@ define i8 @abs_diff_signed_sgt_nuw_extra_use3(i8 %a, i8 %b) {
 
 define i32 @abs_diff_signed_slt_swap_wrong_pred1(i32 %a, i32 %b) {
 ; CHECK-LABEL: @abs_diff_signed_slt_swap_wrong_pred1(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[SUB_BA:%.*]] = sub nsw i32 [[B]], [[A]]
-; CHECK-NEXT:    [[SUB_AB:%.*]] = sub nsw i32 [[A]], [[B]]
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i32 [[SUB_BA]], i32 [[SUB_AB]]
-; CHECK-NEXT:    ret i32 [[COND]]
+; CHECK-NEXT:    [[SUB_AB:%.*]] = sub nsw i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    ret i32 [[SUB_AB]]
 ;
   %cmp = icmp eq i32 %a, %b
   %sub_ba = sub nsw i32 %b, %a
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index d3cf2af4be652..84e0dfa48a6dd 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -2787,8 +2787,7 @@ define <2 x i8> @select_replacement_add_eq_vec_nonuniform(<2 x i8> %x, <2 x i8>
 define <2 x i8> @select_replacement_add_eq_vec_poison(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @select_replacement_add_eq_vec_poison(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 1, i8 poison>
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 1>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i8> [[ADD]], <2 x i8> [[Y:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i8> <i8 2, i8 poison>, <2 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[SEL]]
 ;
   %cmp = icmp eq <2 x i8> %x, <i8 1, i8 poison>
@@ -2839,8 +2838,7 @@ define i8 @select_replacement_sub_noundef(i8 %x, i8 noundef %y, i8 %z) {
 define i8 @select_replacement_sub(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @select_replacement_sub(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[SUB]], i8 [[Z:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 0, i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp eq i8 %x, %y

>From e3d80650952d9f866062834ad0dcc9ec49e711b7 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Wed, 5 Jun 2024 11:35:13 -0500
Subject: [PATCH 3/3] [InstCombine] Improve coverage of
 `foldSelectValueEquivalence` for non-constants

There are a few more cases we can easily cover.
   - If f(Y) simplifies to Y
   - If replace X with Y makes X one-use (and make Y non one-use).

These all require that Y is not undef or poison
---
 .../InstCombine/InstCombineSelect.cpp         | 51 +++++++++++++++++--
 .../InstCombine/select-cmp-eq-op-fold.ll      | 16 +++---
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 0eaec485b853f..1b442aa6b8ebd 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -1300,6 +1300,8 @@ Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
     if (TrueVal == OldOp)
       return nullptr;
 
+    bool NewOpNeverUndef = false;
+
     if (Value *V = simplifyWithOpReplaced(TrueVal, OldOp, NewOp, SQ,
                                           /* AllowRefinement= */ true)) {
       // Need some guarantees about the new simplified op to ensure we don't inf
@@ -1308,11 +1310,52 @@ Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
       if (match(V, m_ImmConstant()))
         return replaceOperand(Sel, Swapped ? 2 : 1, V);
 
-      // If NewOp is a constant and OldOp is not replace iff NewOp doesn't
-      // contain and undef/poison elements.
-      if (match(NewOp, m_ImmConstant()) &&
-          isGuaranteedNotToBeUndefOrPoison(NewOp, SQ.AC, &Sel, &DT))
+      // We can't do any further replacement if NewOp may be undef/poison.
+      if (!isGuaranteedNotToBeUndefOrPoison(NewOp, SQ.AC, &Sel, &DT))
+        return nullptr;
+
+      // If NewOp is a constant, replace.
+      if (match(NewOp, m_ImmConstant()))
+        return replaceOperand(Sel, Swapped ? 2 : 1, V);
+
+      // If we simplified the TrueArm -> NewOp then replace.
+      // This handles things like `select`/`min`/`max`/`or`/`and`/etc...
+      if (NewOp == V)
         return replaceOperand(Sel, Swapped ? 2 : 1, V);
+
+      NewOpNeverUndef = true;
+    }
+
+    // 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.
+    if (TrueVal->hasOneUse() && isa<Instruction>(TrueVal)) {
+      bool BinOpMatch =
+          match(TrueVal, m_c_BinOp(m_Specific(OldOp), m_Specific(NewOp)));
+      bool NewOneUse = isa<Instruction>(OldOp) && OldOp->hasNUses(2) &&
+                       (!isa<Instruction>(NewOp) || !NewOp->hasOneUse());
+      if (BinOpMatch || NewOneUse) {
+        bool Replaced = false;
+        auto *TrueIns = cast<Instruction>(TrueVal);
+        for (unsigned OpIdx = 0; OpIdx < TrueIns->getNumOperands(); ++OpIdx) {
+          if (TrueIns->getOperand(OpIdx) != OldOp)
+            continue;
+          // Need to ensure NewOp is noundef (same reason as above). Wait
+          // until the last moment to do this check as it can be relatively
+          // expensive.
+          if (!NewOpNeverUndef &&
+              !isGuaranteedNotToBeUndefOrPoison(NewOp, SQ.AC, &Sel, &DT))
+            break;
+          NewOpNeverUndef = true;
+          TrueIns->setOperand(OpIdx, NewOp);
+          Replaced = true;
+        }
+        if (Replaced)
+          return replaceOperand(Sel, Swapped ? 2 : 1, TrueIns);
+      }
     }
 
     // Even if TrueVal does not simplify, we can directly replace a use of
diff --git a/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll b/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll
index ec82b1944f723..d2a5c71025370 100644
--- a/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll
+++ b/llvm/test/Transforms/InstCombine/select-cmp-eq-op-fold.ll
@@ -6,8 +6,7 @@ declare void @use.i8(i8)
 define i8 @replace_with_y_noundef(i8 %x, i8 noundef %y, i8 %z) {
 ; CHECK-LABEL: @replace_with_y_noundef(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[AND]], i8 [[Z:%.*]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Y]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp eq i8 %x, %y
@@ -20,8 +19,7 @@ define i8 @replace_with_x_noundef(i8 noundef %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @replace_with_x_noundef(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    call void @use.i1(i1 [[CMP]])
-; CHECK-NEXT:    [[AND:%.*]] = or i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Z:%.*]], i8 [[AND]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[Z:%.*]], i8 [[X]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
   %cmp = icmp ne i8 %x, %y
@@ -50,7 +48,7 @@ define i8 @replace_with_y_for_new_oneuse(i8 noundef %xx, i8 noundef %y, i8 %z) {
 ; CHECK-LABEL: @replace_with_y_for_new_oneuse(
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[ADD:%.*]] = shl nuw i8 [[Y]], 1
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
@@ -65,7 +63,7 @@ define i8 @replace_with_y_for_new_oneuse2(i8 %xx, i8 noundef %y, i8 %z, i8 %q) {
 ; CHECK-LABEL: @replace_with_y_for_new_oneuse2(
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[X]], [[Q:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw i8 [[Y]], [[Q:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[ADD]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
@@ -81,7 +79,7 @@ define i8 @replace_with_x_for_new_oneuse(i8 noundef %xx, i8 noundef %yy, i8 %z,
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
 ; CHECK-NEXT:    [[Y:%.*]] = add i8 [[YY:%.*]], [[W:%.*]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[X]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;
@@ -115,7 +113,7 @@ define i8 @replace_with_x_for_simple_binop(i8 noundef %xx, i8 %yy, i8 %z, i8 %w)
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[XX:%.*]], 13
 ; CHECK-NEXT:    [[Y:%.*]] = add i8 [[YY:%.*]], [[W:%.*]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], [[Y]]
-; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[X]], [[X]]
 ; CHECK-NEXT:    call void @use.i8(i8 [[Y]])
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
@@ -147,7 +145,7 @@ define i8 @replace_with_none_for_new_oneuse_fail_maybe_undef(i8 %xx, i8 %y, i8 %
 define i8 @replace_with_y_for_simple_binop(i8 %x, i8 noundef %y, i8 %z) {
 ; CHECK-LABEL: @replace_with_y_for_simple_binop(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i8 [[X]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i8 [[Y]], [[Y]]
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 [[MUL]], i8 [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[SEL]]
 ;



More information about the llvm-commits mailing list