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

[goldstein.w.n]

Summary

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.

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[nikic]

Can you please rebase over https://github.com/llvm/llvm-project/commit/d0de2c51c9a9fc0fedb97ee98f61ce08cb34972b? I hope that will clarify what the actual change here is, because making foldOperationIntoSelectOperand() fallible doesn't make a whole lot of sense.

[goldstein.w.n]

Rebase

[goldstein.w.n]

In D146349#4208910, @nikic wrote:

Can you please rebase over https://github.com/llvm/llvm-project/commit/d0de2c51c9a9fc0fedb97ee98f61ce08cb34972b? I hope that will clarify what the actual change here is, because making foldOperationIntoSelectOperand() fallible doesn't make a whole lot of sense.

Done and agreed. Updated version only uses the information from the condition
to improve the constant folding.

[nikic]

This should also relax the requirement on at least one caller so that the change becomes testable. Or maybe you can use the one in InstCombineCalls to test this, I think one of the multi-argument intrinsics should work.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1035

Why was this factored out into a separate function? Doesn't really seem to simplify things.

[goldstein.w.n]

In D146349#4262696, @nikic wrote:

This should also relax the requirement on at least one caller so that the change becomes testable. Or maybe you can use the one in InstCombineCalls to test this, I think one of the multi-argument intrinsics should work.

Thats what happens in: D146350, do you want me to merge the two?

[nikic]

In D146349#4262718, @goldstein.w.n wrote:

In D146349#4262696, @nikic wrote:

This should also relax the requirement on at least one caller so that the change becomes testable. Or maybe you can use the one in InstCombineCalls to test this, I think one of the multi-argument intrinsics should work.

Thats what happens in: D146350, do you want me to merge the two?

That patch does more than we need for this. It would be enough to relax any of the existing calls (or use the intrinsic one if it works).

[goldstein.w.n]

Remove new helper

[nikic]

LGTM

[goldstein.w.n]

Rebase

[nikic]

Sorry, I missed something in this patch as well. This transform isn't valid for undef constants: https://alive2.llvm.org/ce/z/4JjGKv We should guard it against isGuaranteedNotToBeUndefOrPoison.

[goldstein.w.n]

In D146349#4309379, @nikic wrote:

Sorry, I missed something in this patch as well. This transform isn't valid for undef constants: https://alive2.llvm.org/ce/z/4JjGKv We should guard it against isGuaranteedNotToBeUndefOrPoison.

Because its constant elements can we just do C->containsUndefOrPoisonElement()?

Edit2: Edit: NVM its the condition that needs to be guarded.

[goldstein.w.n]

In D146349#4309379, @nikic wrote:

Sorry, I missed something in this patch as well. This transform isn't valid for undef constants: https://alive2.llvm.org/ce/z/4JjGKv We should guard it against isGuaranteedNotToBeUndefOrPoison.

@nikic, I think D149592 should fix the issue.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombinerImpl::foldBinopOfSextBoolToSelect(BinaryOperator &BO) {
   // bo (sext i1 X), C --> select X, (bo -1, C), (bo 0, C)
   Constant *Ones = ConstantInt::getAllOnesValue(BO.getType());
   Constant *Zero = ConstantInt::getNullValue(BO.getType());
   Value *TVal = Builder.CreateBinOp(BO.getOpcode(), Ones, C);
   Value *FVal = Builder.CreateBinOp(BO.getOpcode(), Zero, C);
   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());
 }
 
 static Value *foldOperationIntoSelectOperand(Instruction &I, SelectInst *SI,
                                              Value *NewOp, InstCombiner &IC) {
   Instruction *Clone = I.clone();
   Clone->replaceUsesOfWith(SI, NewOp);
   IC.InsertNewInstBefore(Clone, *SI);
if ((SrcTy == nullptr) != (DestTy == nullptr))
       return nullptr;
 
     // If vectors, verify that they have the same number of elements.
     if (SrcTy && SrcTy->getElementCount() != DestTy->getElementCount())
       return nullptr;
   }
 
   // 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;
 
   // Create an instruction for the arm that did not fold.
   if (!NewTV)
     NewTV = foldOperationIntoSelectOperand(Op, SI, TV, *this);
   if (!NewFV)
     NewFV = foldOperationIntoSelectOperand(Op, SI, FV, *this);

llvm/test/Transforms/InstCombine/binop-select.ll

;
   %cond = select i1 %c, i32 %x, i32 0
   %add = add i32 %cond, %x
   ret i32 %add
 }
 
 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
   %cond = select i1 %c, i32 6, i32 %y
   %sub = call i32 @llvm.sadd.sat.i32(i32 %x, i32 %cond)
   ret i32 %sub
 }
 
 define i32 @test_sub_deduce_true_no_const_fold(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test_sub_deduce_true_no_const_fold(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 9
 ; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 [[Y:%.*]], i32 6
 ; CHECK-NEXT:    [[SUB:%.*]] = call i32 @llvm.sadd.sat.i32(i32 [[X]], i32 [[COND]])
 ; CHECK-NEXT:    ret i32 [[SUB]]
 ;
   %c = icmp eq i32 %x, 9
   %cond = select i1 %c, i32 %y, i32 6
   %sub = call i32 @llvm.sadd.sat.i32(i32 %x, i32 %cond)
   ret i32 %sub
 }
 
 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
   %cond = select i1 %c, i32 %y, i32 7
   %sub = call i32 @llvm.sadd.sat.i32(i32 %x, i32 %cond)
   ret i32 %sub
 }