[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

1030

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;
-
-  SmallVector<Constant *> ConstOps;
+static SmallVector<Value *>
+getInstSelectArmOperands(Instruction &I, SelectInst *SI, bool IsTrueArm) {
+  SmallVector<Value *> Ops;
   for (Value *Op : I.operands()) {
+    CmpInst::Predicate Pred;
+    Constant *CondC;
     if (Op == SI)
-      ConstOps.push_back(ConstSO);
-    else if (auto *C = dyn_cast<Constant>(Op))
+      Ops.push_back(IsTrueArm ? SI->getTrueValue() : SI->getFalseValue());
+    else if (match(SI->getCondition(),
+                   m_ICmp(Pred, m_Specific(Op), m_Constant(CondC))) &&
+             Pred == (IsTrueArm ? ICmpInst::ICMP_EQ : ICmpInst::ICMP_NE))
+      Ops.push_back(CondC);
+    else
+      Ops.push_back(Op);
+  }
+
+  return Ops;
+}
+
+static Constant *constantFoldOperationIntoSelectOperand(Instruction &I,
+                                                        SelectInst *SI,
+                                                        bool IsTrueArm) {
+  SmallVector<Constant *> ConstOps;
+  SmallVector<Value *> Ops = getInstSelectArmOperands(I, SI, IsTrueArm);
+  for (Value *Op : Ops) {
+    if (auto *C = dyn_cast<Constant>(Op))
       ConstOps.push_back(C);
     else
-      llvm_unreachable("Operands should be select or constant");
+      return nullptr;
   }
-  return ConstantFoldInstOperands(&I, ConstOps, I.getModule()->getDataLayout());
+  return ConstantFoldInstOperands(&I, ConstOps,
+                                  I.getModule()->getDataLayout());
 }
 
-static Value *foldOperationIntoSelectOperand(Instruction &I, Value *SO,
-                                             InstCombiner::BuilderTy &Builder) {
+static Value *foldOperationIntoSelectOperand(Instruction &I, SelectInst *SI,
+                                             InstCombiner::BuilderTy &Builder,
+                                             bool IsTrueArm) {
+
+  Value *SO = IsTrueArm ? SI->getTrueValue() : SI->getFalseValue();
   if (auto *Cast = dyn_cast<CastInst>(&I))
     return Builder.CreateCast(Cast->getOpcode(), SO, I.getType());
 
+  SmallVector<Value *> Ops = getInstSelectArmOperands(I, SI, IsTrueArm);
+
   if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
-    assert(canConstantFoldCallTo(II, cast<Function>(II->getCalledOperand())) &&
-           "Expected constant-foldable intrinsic");
+      //    assert(Ops.size() == II->arg_size() &&
+      //           "Failed to collect IntrinsicArguments");
+    // TODO: Use the select condition to make II->getCalledOperand() constant
+    // here.
+    if (!canConstantFoldCallTo(II, cast<Function>(II->getCalledOperand())))
+      return nullptr;
     Intrinsic::ID IID = II->getIntrinsicID();
     if (II->arg_size() == 1)
       return Builder.CreateUnaryIntrinsic(IID, SO);
 
+    if (!(isa<Constant>(Ops[0]) || isa<Constant>(Ops[1])))
+      return nullptr;
+
     // This works for real binary ops like min/max (where we always expect the
     // constant operand to be canonicalized as op1) and unary ops with a bonus
     // constant argument like ctlz/cttz.
-    // TODO: Handle non-commutative binary intrinsics as below for binops.
     assert(II->arg_size() == 2 && "Expected binary intrinsic");
-    assert(isa<Constant>(II->getArgOperand(1)) && "Expected constant operand");
-    return Builder.CreateBinaryIntrinsic(IID, SO, II->getArgOperand(1));
+    return Builder.CreateBinaryIntrinsic(IID, Ops[0], Ops[1]);
   }
 
   if (auto *EI = dyn_cast<ExtractElementInst>(&I))
     return Builder.CreateExtractElement(SO, EI->getIndexOperand());
 
-  assert(I.isBinaryOp() && "Unexpected opcode for select folding");
-
-  // Figure out if the constant is the left or the right argument.
-  bool ConstIsRHS = isa<Constant>(I.getOperand(1));
-  Constant *ConstOperand = cast<Constant>(I.getOperand(ConstIsRHS));
+  if (!(isa<Constant>(Ops[0]) || isa<Constant>(Ops[1])))
+    return nullptr;
 
-  Value *Op0 = SO, *Op1 = ConstOperand;
-  if (!ConstIsRHS)
-    std::swap(Op0, Op1);
+  assert(I.isBinaryOp() && Ops.size() == 2 &&
+         "Unexpected opcode for select folding");
 
-  Value *NewBO = Builder.CreateBinOp(cast<BinaryOperator>(&I)->getOpcode(), Op0,
-                                     Op1, SO->getName() + ".op");
+  Value *NewBO = Builder.CreateBinOp(cast<BinaryOperator>(&I)->getOpcode(),
+                                     Ops[0], Ops[1], SO->getName() + ".op");
   if (auto *NewBOI = dyn_cast<Instruction>(NewBO))
     NewBOI->copyIRFlags(&I);
   return NewBO;
 }
 
 Instruction *InstCombinerImpl::FoldOpIntoSelect(Instruction &Op, SelectInst *SI,
                                                 bool FoldWithMultiUse) {
   // Don't modify shared select instructions unless set FoldWithMultiUse
if (CI->hasOneUse()) {
 
       if ((areLooselyEqual(TV, Op0) && areLooselyEqual(FV, Op1)) ||
           (areLooselyEqual(FV, Op0) && areLooselyEqual(TV, Op1)))
         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, TV, Builder);
+    if (!(NewTV = foldOperationIntoSelectOperand(Op, SI, Builder,
+                                                 /*IsTrueArm*/ true)))
+      return nullptr;
   if (!NewFV)
-    NewFV = foldOperationIntoSelectOperand(Op, FV, Builder);
+    if (!(NewFV = foldOperationIntoSelectOperand(Op, SI, Builder,
+                                                 /*IsTrueArm*/ false)))
+      return nullptr;
+
   return SelectInst::Create(SI->getCondition(), NewTV, NewFV, "", nullptr, SI);
 }
 
 Instruction *InstCombinerImpl::foldOpIntoPhi(Instruction &I, PHINode *PN) {
   unsigned NumPHIValues = PN->getNumIncomingValues();
   if (NumPHIValues == 0)
     return nullptr;