[InstCombine] Generalize sub of selects optimization to all BinaryOperators

This exposes further optimization opportunities if the selects are
correlated.

llvm-svn: 242235

3 files changed, 37 insertions, 26 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index 2d2c109f324..a8d01725517 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -1611,32 +1611,6 @@ Instruction *InstCombiner::visitSub(BinaryOperator &I) {
       return BinaryOperator::CreateAnd(A, B);
   }
 
-  // (sub (select (a, c, b)), (select (a, d, b))) -> (select (a, (sub c, d), 0))
-  // (sub (select (a, b, c)), (select (a, b, d))) -> (select (a, 0, (sub c, d)))
-  if (auto *SI0 = dyn_cast<SelectInst>(Op0)) {
-    if (auto *SI1 = dyn_cast<SelectInst>(Op1)) {
-      if (SI0->getCondition() == SI1->getCondition()) {
-        if (Value *V = SimplifySubInst(
-                SI0->getFalseValue(), SI1->getFalseValue(), I.hasNoSignedWrap(),
-                I.hasNoUnsignedWrap(), DL, TLI, DT, AC))
-          return SelectInst::Create(
-              SI0->getCondition(),
-              Builder->CreateSub(SI0->getTrueValue(), SI1->getTrueValue(), "",
-                                 /*HasNUW=*/I.hasNoUnsignedWrap(),
-                                 /*HasNSW=*/I.hasNoSignedWrap()),
-              V);
-        if (Value *V = SimplifySubInst(SI0->getTrueValue(), SI1->getTrueValue(),
-                                       I.hasNoSignedWrap(),
-                                       I.hasNoUnsignedWrap(), DL, TLI, DT, AC))
-          return SelectInst::Create(
-              SI0->getCondition(), V,
-              Builder->CreateSub(SI0->getFalseValue(), SI1->getFalseValue(), "",
-                                 /*HasNUW=*/I.hasNoUnsignedWrap(),
-                                 /*HasNSW=*/I.hasNoSignedWrap()));
-      }
-    }
-  }
-
   if (Op0->hasOneUse()) {
     Value *Y = nullptr;
     // ((X | Y) - X) --> (~X & Y)
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index fd34a244f27..d04ae0702a6 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -623,6 +623,33 @@ Value *InstCombiner::SimplifyUsingDistributiveLaws(BinaryOperator &I) {
       }
   }
 
+  // (op (select (a, c, b)), (select (a, d, b))) -> (select (a, (op c, d), 0))
+  // (op (select (a, b, c)), (select (a, b, d))) -> (select (a, 0, (op c, d)))
+  if (auto *SI0 = dyn_cast<SelectInst>(LHS)) {
+    if (auto *SI1 = dyn_cast<SelectInst>(RHS)) {
+      if (SI0->getCondition() == SI1->getCondition()) {
+        Value *SI = nullptr;
+        if (Value *V = SimplifyBinOp(TopLevelOpcode, SI0->getFalseValue(),
+                                     SI1->getFalseValue(), DL, TLI, DT, AC))
+          SI = Builder->CreateSelect(SI0->getCondition(),
+                                     Builder->CreateBinOp(TopLevelOpcode,
+                                                          SI0->getTrueValue(),
+                                                          SI1->getTrueValue()),
+                                     V);
+        if (Value *V = SimplifyBinOp(TopLevelOpcode, SI0->getTrueValue(),
+                                     SI1->getTrueValue(), DL, TLI, DT, AC))
+          SI = Builder->CreateSelect(
+              SI0->getCondition(), V,
+              Builder->CreateBinOp(TopLevelOpcode, SI0->getFalseValue(),
+                                   SI1->getFalseValue()));
+        if (SI) {
+          SI->takeName(&I);
+          return SI;
+        }
+      }
+    }
+  }
+
   return nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/add2.ll b/llvm/test/Transforms/InstCombine/add2.ll
index 1af2b0ffbf0..2fe9e8cadeb 100644
--- a/llvm/test/Transforms/InstCombine/add2.ll
+++ b/llvm/test/Transforms/InstCombine/add2.ll
@@ -395,3 +395,13 @@ define i8 @add_of_mul(i8 %x, i8 %y, i8 %z) {
   %sum = add nsw i8 %mA, %mB
   ret i8 %sum
 }
+
+define i32 @add_of_selects(i1 %A, i32 %B) {
+  %sel0 = select i1 %A, i32 0, i32 -2
+  %sel1 = select i1 %A, i32 %B, i32 2
+  %add = add i32 %sel0, %sel1
+  ret i32 %add
+; CHECK-LABEL: @add_of_selects(
+; CHECK-NEXT: %[[sel:.*]] = select i1 %A, i32 %B, i32 0
+; CHECK-NEXT: ret i32 %[[sel]]
+}