[llvm] 82f0827 - [InstCombine] Make `FoldOpIntoSelect` handle non-constants and use condition to deduce constants.

[Noah Goldstein via llvm-commits]

Author: Noah Goldstein
Date: 2023-04-14T13:14:32-05:00
New Revision: 82f082761383907d4ceefad5ab50f0758fbe6a46

URL: https://github.com/llvm/llvm-project/commit/82f082761383907d4ceefad5ab50f0758fbe6a46
DIFF: https://github.com/llvm/llvm-project/commit/82f082761383907d4ceefad5ab50f0758fbe6a46.diff

LOG: [InstCombine] Make `FoldOpIntoSelect` handle non-constants and use condition to deduce constants.

Make the fold use the information present in the condition for deducing constants i.e:
```
%c = icmp eq i8 %x, 10
%s = select i1 %c, i8 3, i8 2
%r = mul i8 %x, %s
```

If we fold the `mul` into the select, on the true side we insert `10` for `%x` in the `mul`.

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D146349

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
    llvm/test/Transforms/InstCombine/binop-select.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 6a61b8526a16a..5af3b677b376e 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -1028,21 +1028,29 @@ Instruction *InstCombinerImpl::foldBinopOfSextBoolToSelect(BinaryOperator &BO) {
   return SelectInst::Create(X, TVal, FVal);
 }
 
-static Constant *constantFoldOperationIntoSelectOperand(
-    Instruction &I, SelectInst *SI, Value *SO) {
-  auto *ConstSO = dyn_cast<Constant>(SO);
-  if (!ConstSO)
-    return nullptr;
-
+static Constant *constantFoldOperationIntoSelectOperand(Instruction &I,
+                                                        SelectInst *SI,
+                                                        bool IsTrueArm) {
   SmallVector<Constant *> ConstOps;
   for (Value *Op : I.operands()) {
-    if (Op == SI)
-      ConstOps.push_back(ConstSO);
-    else if (auto *C = dyn_cast<Constant>(Op))
-      ConstOps.push_back(C);
-    else
+    CmpInst::Predicate Pred;
+    Constant *C = nullptr;
+    if (Op == SI) {
+      C = dyn_cast<Constant>(IsTrueArm ? SI->getTrueValue()
+                                       : SI->getFalseValue());
+    } else if (match(SI->getCondition(),
+                     m_ICmp(Pred, m_Specific(Op), m_Constant(C))) &&
+               Pred == (IsTrueArm ? ICmpInst::ICMP_EQ : ICmpInst::ICMP_NE)) {
+      // Pass
+    } else {
+      C = dyn_cast<Constant>(Op);
+    }
+    if (C == nullptr)
       return nullptr;
+
+    ConstOps.push_back(C);
   }
+
   return ConstantFoldInstOperands(&I, ConstOps, I.getModule()->getDataLayout());
 }
 
@@ -1085,8 +1093,8 @@ Instruction *InstCombinerImpl::FoldOpIntoSelect(Instruction &Op, SelectInst *SI,
   }
 
   // Make sure that one of the select arms constant folds successfully.
-  Value *NewTV = constantFoldOperationIntoSelectOperand(Op, SI, TV);
-  Value *NewFV = constantFoldOperationIntoSelectOperand(Op, SI, FV);
+  Value *NewTV = constantFoldOperationIntoSelectOperand(Op, SI, /*IsTrueArm*/ true);
+  Value *NewFV = constantFoldOperationIntoSelectOperand(Op, SI, /*IsTrueArm*/ false);
   if (!NewTV && !NewFV)
     return nullptr;
 

diff  --git a/llvm/test/Transforms/InstCombine/binop-select.ll b/llvm/test/Transforms/InstCombine/binop-select.ll
index 655a34ccdea02..da6c17984eb91 100644
--- a/llvm/test/Transforms/InstCombine/binop-select.ll
+++ b/llvm/test/Transforms/InstCombine/binop-select.ll
@@ -73,8 +73,8 @@ define i32 @test5(i1 %c, i32 %x, i32 %y) {
 define i32 @test_sub_deduce_true(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test_sub_deduce_true(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 9
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 6, i32 [[Y:%.*]]
-; CHECK-NEXT:    [[SUB:%.*]] = call i32 @llvm.sadd.sat.i32(i32 [[X]], i32 [[COND]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.sadd.sat.i32(i32 [[X]], i32 [[Y:%.*]])
+; CHECK-NEXT:    [[SUB:%.*]] = select i1 [[C]], i32 15, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[SUB]]
 ;
   %c = icmp eq i32 %x, 9
@@ -99,8 +99,8 @@ define i32 @test_sub_deduce_true_no_const_fold(i32 %x, i32 %y) {
 define i32 @test_sub_deduce_false(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test_sub_deduce_false(
 ; CHECK-NEXT:    [[C_NOT:%.*]] = icmp eq i32 [[X:%.*]], 9
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C_NOT]], i32 7, i32 [[Y:%.*]]
-; CHECK-NEXT:    [[SUB:%.*]] = call i32 @llvm.sadd.sat.i32(i32 [[X]], i32 [[COND]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.sadd.sat.i32(i32 [[X]], i32 [[Y:%.*]])
+; CHECK-NEXT:    [[SUB:%.*]] = select i1 [[C_NOT]], i32 16, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[SUB]]
 ;
   %c = icmp ne i32 %x, 9