[InstSimplify] Support vectors in simplifyWithOpReplaced()

[nikic]

Summary

We can handle vectors inside simplifyWithOpReplaced(), as long as cross-lane operations are excluded. The equality can hold (or not hold) for each vector lane independently, so we shouldn't use the replacement value from other lanes.

I believe the only operations relevant here are shufflevector (where all previous bugs were seen) and calls (which might use shuffle-like intrinsics and would require more careful classification).

[spatel]

Seems right, but do we have tests that exercise the call restriction? Maybe something like these:

declare <2 x i32> @llvm.sadd.sat.v2i32(<2 x i32>, <2 x i32>)
declare <2 x i32> @arbitrary_call(<2 x i32>)

define <2 x i32> @could_fold(<2 x i32> %x, <2 x i32> %y) {
  %a = call <2 x i32> @llvm.sadd.sat.v2i32(<2 x i32> %x, <2 x i32> <i32 42, i32 42>)
  %cond = icmp eq <2 x i32> %x, zeroinitializer
  %s = select <2 x i1> %cond, <2 x i32> %a, <2 x i32> %y
  ret <2 x i32> %s
}

define <2 x i32> @must_not_fold(<2 x i32> %x, <2 x i32> %y) {
  %a = call <2 x i32> @arbitrary_call(<2 x i32> %x)
  %cond = icmp eq <2 x i32> %x, zeroinitializer
  %s = select <2 x i1> %cond, <2 x i32> %a, <2 x i32> %y
  ret <2 x i32> %s
}

[nikic]

Rebase over additional tests with calls.

[spatel]

LGTM

llvm/lib/Analysis/InstructionSimplify.cpp

Value *llvm::simplifyFCmpInst(unsigned Predicate, Value *LHS, Value *RHS,
                               FastMathFlags FMF, const SimplifyQuery &Q) {
   return ::simplifyFCmpInst(Predicate, LHS, RHS, FMF, Q, RecursionLimit);
 }
 
 static Value *simplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
                                      const SimplifyQuery &Q,
                                      bool AllowRefinement,
                                      unsigned MaxRecurse) {
-  assert(!Op->getType()->isVectorTy() && "This is not safe for vectors");
-
   // Trivial replacement.
   if (V == Op)
     return RepOp;
 
   // We cannot replace a constant, and shouldn't even try.
   if (isa<Constant>(Op))
     return nullptr;
 
   auto *I = dyn_cast<Instruction>(V);
   if (!I || !is_contained(I->operands(), Op))
     return nullptr;
 
+  if (Op->getType()->isVectorTy()) {
+    // For vector types, the simplification must hold per-lane, so forbid
+    // potentially cross-lane operations like shufflevector.
+    assert(I->getType()->isVectorTy() && "Vector type mismatch");
+    if (isa<ShuffleVectorInst>(I) || isa<CallBase>(I))
+      return nullptr;
+  }
+
   // Replace Op with RepOp in instruction operands.
   SmallVector<Value *, 8> NewOps(I->getNumOperands());
   transform(I->operands(), NewOps.begin(),
             [&](Value *V) { return V == Op ? RepOp : V; });
 
   if (!AllowRefinement) {
     // General InstSimplify functions may refine the result, e.g. by returning
     // a constant for a potentially poison value. To avoid this, implement only
static Value *simplifySelectWithICmpCond(Value *CondVal, Value *TrueVal,
   // Check for other compares that behave like bit test.
   if (Value *V =
           simplifySelectWithFakeICmpEq(CmpLHS, CmpRHS, Pred, TrueVal, FalseVal))
     return V;
 
   // If we have a scalar equality comparison, then we know the value in one of
   // the arms of the select. See if substituting this value into the arm and
   // simplifying the result yields the same value as the other arm.
-  // Note that the equivalence/replacement opportunity does not hold for vectors
-  // because each element of a vector select is chosen independently.
-  if (Pred == ICmpInst::ICMP_EQ && !CondVal->getType()->isVectorTy()) {
+  if (Pred == ICmpInst::ICMP_EQ) {
     if (simplifyWithOpReplaced(FalseVal, CmpLHS, CmpRHS, Q,
                                /* AllowRefinement */ false,
                                MaxRecurse) == TrueVal ||
         simplifyWithOpReplaced(FalseVal, CmpRHS, CmpLHS, Q,
                                /* AllowRefinement */ false,
                                MaxRecurse) == TrueVal)
       return FalseVal;
     if (simplifyWithOpReplaced(TrueVal, CmpLHS, CmpRHS, Q,

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

 ///   %cmp = icmp eq i32 %x, 2147483647
 ///   %add = add nsw i32 %x, 1
 ///   %sel = select i1 %cmp, i32 -2147483648, i32 %add
 ///
 /// We can't replace %sel with %add unless we strip away the flags.
 /// TODO: Wrapping flags could be preserved in some cases with better analysis.
 Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
                                                           ICmpInst &Cmp) {
-  // Value equivalence substitution requires an all-or-nothing replacement.
-  // It does not make sense for a vector compare where each lane is chosen
-  // independently.
-  if (!Cmp.isEquality() || Cmp.getType()->isVectorTy())
+  if (!Cmp.isEquality())
     return nullptr;
 
   // Canonicalize the pattern to ICMP_EQ by swapping the select operands.
   Value *TrueVal = Sel.getTrueValue(), *FalseVal = Sel.getFalseValue();
   bool Swapped = false;
   if (Cmp.getPredicate() == ICmpInst::ICMP_NE) {
     std::swap(TrueVal, FalseVal);
     Swapped = true;
if (TrueVal != CmpLHS &&
 
     // 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
     // with different operands, which should not cause side-effects or trigger
     // undefined behavior). Only do this if CmpRHS is a constant, as
     // profitability is not clear for other cases.
     // FIXME: The replacement could be performed recursively.
-    if (match(CmpRHS, m_ImmConstant()) && !match(CmpLHS, m_ImmConstant()))
+    // FIXME: Support vectors.
+    if (match(CmpRHS, m_ImmConstant()) && !match(CmpLHS, m_ImmConstant()) &&
+        !Cmp.getType()->isVectorTy())
       if (auto *I = dyn_cast<Instruction>(TrueVal))
         if (I->hasOneUse() && isSafeToSpeculativelyExecute(I))
           for (Use &U : I->operands())
             if (U == CmpLHS) {
               replaceUse(U, CmpRHS);
               return &Sel;
             }
   }
if (SelType->isIntOrIntVectorTy(1) && !isa<Constant>(CondVal) &&
     }
     // select (~c & b), a, b -> and b, (or freeze(a), c)
     if (match(CondVal, m_c_And(m_Not(m_Value(C)), m_Specific(FalseVal))) &&
         CondVal->hasOneUse()) {
       TrueVal = Builder.CreateFreeze(TrueVal);
       return BinaryOperator::CreateAnd(FalseVal, Builder.CreateOr(C, TrueVal));
     }
 
-    if (!SelType->isVectorTy()) {
-      if (Value *S = simplifyWithOpReplaced(TrueVal, CondVal, One, SQ,
-                                            /* AllowRefinement */ true))
-        return replaceOperand(SI, 1, S);
-      if (Value *S = simplifyWithOpReplaced(FalseVal, CondVal, Zero, SQ,
-                                            /* AllowRefinement */ true))
-        return replaceOperand(SI, 2, S);
-    }
 
     if (match(FalseVal, m_Zero()) || match(TrueVal, m_One())) {
       Use *Y = nullptr;
       bool IsAnd = match(FalseVal, m_Zero()) ? true : false;
       Value *Op1 = IsAnd ? TrueVal : FalseVal;
       if (isCheckForZeroAndMulWithOverflow(CondVal, Op1, IsAnd, Y)) {
         auto *FI = new FreezeInst(*Y, (*Y)->getName() + ".fr");
         InsertNewInstBefore(FI, *cast<Instruction>(Y->getUser()));

llvm/test/Transforms/InstCombine/select-safe-bool-transforms.ll

 ; CHECK-NEXT:    ret <2 x i1> [[A:%.*]]
 ;
   %c = select <2 x i1> %A, <2 x i1> %B, <2 x i1> zeroinitializer
   %res = select <2 x i1> %A, <2 x i1> <i1 true, i1 true>, <2 x i1> %c
   ret <2 x i1> %res
 }
 define <2 x i1> @land_lor_right2_vec(<2 x i1> %A, <2 x i1> %B) {
 ; CHECK-LABEL: @land_lor_right2_vec(
-; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[B:%.*]], <2 x i1> [[A:%.*]], <2 x i1> zeroinitializer
-; CHECK-NEXT:    [[RES:%.*]] = select <2 x i1> [[A]], <2 x i1> <i1 true, i1 true>, <2 x i1> [[C]]
-; CHECK-NEXT:    ret <2 x i1> [[RES]]
+; CHECK-NEXT:    ret <2 x i1> [[A:%.*]]
 ;
   %c = select <2 x i1> %B, <2 x i1> %A, <2 x i1> zeroinitializer
   %res = select <2 x i1> %A, <2 x i1> <i1 true, i1 true>, <2 x i1> %c
   ret <2 x i1> %res
 }
 
 ; A bor (A land B)
 define i1 @land_bor_right1(i1 %A, i1 %B) {

llvm/test/Transforms/InstCombine/select.ll

;
   %cmp = icmp eq i32 %x, -2147483648
   %sub = sub nsw i32 0, %x
   %sel = select i1 %cmp, i32 -2147483648, i32 %sub
   ret i32 %sel
 }
 
 define <2 x i32> @PR27817_nsw_vec(<2 x i32> %x) {
 ; CHECK-LABEL: @PR27817_nsw_vec(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[X:%.*]], <i32 -2147483648, i32 -2147483648>
-; CHECK-NEXT:    [[SUB:%.*]] = sub nsw <2 x i32> zeroinitializer, [[X]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i32> <i32 -2147483648, i32 -2147483648>, <2 x i32> [[SUB]]
-; CHECK-NEXT:    ret <2 x i32> [[SEL]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub <2 x i32> zeroinitializer, [[X:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[SUB]]
 ;
   %cmp = icmp eq <2 x i32> %x, <i32 -2147483648, i32 -2147483648>
   %sub = sub nsw <2 x i32> zeroinitializer, %x
   %sel = select <2 x i1> %cmp, <2 x i32> <i32 -2147483648, i32 -2147483648>, <2 x i32> %sub
   ret <2 x i32> %sel
 }
 
 define i32 @select_icmp_slt0_xor(i32 %x) {
;
   %add = add i8 %x, 1
   %sel = select i1 %cmp, i8 %add, i8 %y
   ret i8 %sel
 }
 
 define <2 x i8> @select_replacement_add_eq_vec(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @select_replacement_add_eq_vec(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 1, i8 1>
-; 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 2>, <2 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[SEL]]
 ;
   %cmp = icmp eq <2 x i8> %x, <i8 1, i8 1>
   %add = add <2 x i8> %x, <i8 1, i8 1>
   %sel = select <2 x i1> %cmp, <2 x i8> %add, <2 x i8> %y
   ret <2 x i8> %sel
 }
 
 define <2 x i8> @select_replacement_add_eq_vec_nonuniform(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @select_replacement_add_eq_vec_nonuniform(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 1, i8 2>
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 3, i8 4>
-; 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 4, i8 6>, <2 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[SEL]]
 ;
   %cmp = icmp eq <2 x i8> %x, <i8 1, i8 2>
   %add = add <2 x i8> %x, <i8 3, i8 4>
   %sel = select <2 x i1> %cmp, <2 x i8> %add, <2 x i8> %y
   ret <2 x i8> %sel
 }
 

llvm/test/Transforms/InstSimplify/select.ll

;
   %cmp = icmp ne i32 %y, 0
   %add = add i32 %y, %x
   %sel = select i1 %cmp, i32 %add, i32 %x
   ret i32 %sel
 }
 
 define <2 x i32> @select_neutral_add_rhs_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @select_neutral_add_rhs_vec(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[Y:%.*]], zeroinitializer
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i32> [[X:%.*]], [[Y]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i32> [[ADD]], <2 x i32> [[X]]
-; CHECK-NEXT:    ret <2 x i32> [[SEL]]
+; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i32> [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[ADD]]
 ;
   %cmp = icmp ne <2 x i32> %y, zeroinitializer
   %add = add <2 x i32> %x, %y
   %sel = select <2 x i1> %cmp, <2 x i32> %add, <2 x i32> %x
   ret <2 x i32> %sel
 }
 
 define <2 x i32> @select_neutral_add_lhs_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @select_neutral_add_lhs_vec(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[Y:%.*]], zeroinitializer
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i32> [[Y]], [[X:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i32> [[ADD]], <2 x i32> [[X]]
-; CHECK-NEXT:    ret <2 x i32> [[SEL]]
+; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i32> [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[ADD]]
 ;
   %cmp = icmp ne <2 x i32> %y, zeroinitializer
   %add = add <2 x i32> %y, %x
   %sel = select <2 x i1> %cmp, <2 x i32> %add, <2 x i32> %x
   ret <2 x i32> %sel
 }
 
 define i32 @select_neutral_sub_rhs(i32 %x, i32 %y) {

 ; CHECK-NEXT:    ret i32 [[T1]]
 ;
   %t0 = icmp eq i32 %x, 0
   %t1 = call i32 @llvm.ctpop.i32(i32 %x)
   %sel = select i1 %t0, i32 0, i32 %t1
   ret i32 %sel
 }
 
+; FIXME: This is safe to fold.
 define <2 x i32> @select_ctpop_zero_vec(<2 x i32> %x) {
 ; CHECK-LABEL: @select_ctpop_zero_vec(
 ; CHECK-NEXT:    [[T0:%.*]] = icmp eq <2 x i32> [[X:%.*]], zeroinitializer
 ; CHECK-NEXT:    [[T1:%.*]] = call <2 x i32> @llvm.ctpop.v2i32(<2 x i32> [[X]])
 ; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[T0]], <2 x i32> zeroinitializer, <2 x i32> [[T1]]
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %t0 = icmp eq <2 x i32> %x, zeroinitializer
   %t1 = call <2 x i32> @llvm.ctpop.v2i32(<2 x i32> %x)
   %sel = select <2 x i1> %t0, <2 x i32> zeroinitializer, <2 x i32> %t1
   ret <2 x i32> %sel
 }
 
+; Negative test: Cannot fold due to cross-lane intrinsic.
 define <2 x i32> @select_vector_reverse(<2 x i32> %x) {
 ; CHECK-LABEL: @select_vector_reverse(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[X:%.*]], zeroinitializer
 ; CHECK-NEXT:    [[REV:%.*]] = call <2 x i32> @llvm.experimental.vector.reverse.v2i32(<2 x i32> [[X]])
 ; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i32> zeroinitializer, <2 x i32> [[REV]]
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %cmp = icmp eq <2 x i32> %x, zeroinitializer