InstCombine: extend select-equiv to support vectors

foldSelectEquivalence currently doesn't support comparisons on vector
types due to correctness concerns. Note that the only concern is
lane-crossing; ShuffleVector is the only possible lane-crossing
instruction, and ShuffleVectorInst::{isIdentity,isShuffle} are the exact
properties that should not be broken for valid vector-replacements. Put
in the checks, and lift the limitation.

llvm/lib/Analysis/ValueTracking.cpp

@@ -6950,7 +6950,9 @@ bool llvm::onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V) {
 bool llvm::isNotCrossLaneOperation(const Instruction *I) {
   if (auto *II = dyn_cast<IntrinsicInst>(I))
     return isTriviallyVectorizable(II->getIntrinsicID());
-  return !isa<CallBase, BitCastInst, ShuffleVectorInst, ExtractElementInst>(I);
+  auto *Shuffle = dyn_cast<ShuffleVectorInst>(I);
+  return (!Shuffle || Shuffle->isSelect()) &&
+         !isa<CallBase, BitCastInst, ExtractElementInst>(I);
 }
 
 bool llvm::isSafeToSpeculativelyExecute(const Instruction *Inst,

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -1288,6 +1288,10 @@ bool InstCombinerImpl::replaceInInstruction(Value *V, Value *Old, Value *New,
       !isSafeToSpeculativelyExecuteWithVariableReplaced(I))
     return false;
 
+  // Forbid potentially lane-crossing instructions.
+  if (Old->getType()->isVectorTy() && !isNotCrossLaneOperation(I))
+    return false;
+
   bool Changed = false;
   for (Use &U : I->operands()) {
     if (U == Old) {
@@ -1366,9 +1370,8 @@ Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
     // 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: Support vectors.
     if (OldOp == CmpLHS && match(NewOp, m_ImmConstant()) &&
-        !match(OldOp, m_Constant()) && !Cmp.getType()->isVectorTy() &&
+        !match(OldOp, m_Constant()) &&
         isGuaranteedNotToBeUndef(NewOp, SQ.AC, &Sel, &DT))
       if (replaceInInstruction(TrueVal, OldOp, NewOp))
         return &Sel;

llvm/test/Transforms/InstCombine/and-or-icmps.ll

@@ -983,7 +983,7 @@ define <2 x i1> @substitute_constant_or_ne_slt_swap_vec_poison(<2 x i8> %x, <2 x
 define <2 x i1> @substitute_constant_or_ne_slt_swap_vec_logical(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @substitute_constant_or_ne_slt_swap_vec_logical(
 ; CHECK-NEXT:    [[C1:%.*]] = icmp ne <2 x i8> [[X:%.*]], <i8 42, i8 poison>
-; CHECK-NEXT:    [[C2:%.*]] = icmp slt <2 x i8> [[Y:%.*]], [[X]]
+; CHECK-NEXT:    [[C2:%.*]] = icmp slt <2 x i8> [[Y:%.*]], <i8 42, i8 poison>
 ; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[C1]], <2 x i1> <i1 true, i1 true>, <2 x i1> [[C2]]
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;

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

@@ -552,12 +552,12 @@ define i32 @select_xor_icmp_bad_6(i32 %x, i32 %y, i32 %z) {
   ret i32 %C
 }
 
-; Value equivalence substitution is all-or-nothing, so needs a scalar compare.
+; Value equivalence substitution is valid.
 
-define <2 x i8> @select_xor_icmp_vec_bad(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
-; CHECK-LABEL: @select_xor_icmp_vec_bad(
+define <2 x i8> @select_xor_icmp_vec_equivalence(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
+; CHECK-LABEL: @select_xor_icmp_vec_equivalence(
 ; CHECK-NEXT:    [[A:%.*]] = icmp eq <2 x i8> [[X:%.*]], <i8 5, i8 3>
-; CHECK-NEXT:    [[B:%.*]] = xor <2 x i8> [[X]], [[Z:%.*]]
+; CHECK-NEXT:    [[B:%.*]] = xor <2 x i8> [[Z:%.*]], <i8 5, i8 3>
 ; CHECK-NEXT:    [[C:%.*]] = select <2 x i1> [[A]], <2 x i8> [[B]], <2 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[C]]
 ;
@@ -567,7 +567,7 @@ define <2 x i8> @select_xor_icmp_vec_bad(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z)
   ret <2 x i8>  %C
 }
 
-; Value equivalence substitution is all-or-nothing, so needs a scalar compare.
+; Value equivalence substitution is invalid due to lane-crossing shufflevector.
 
 define <2 x i32> @vec_select_no_equivalence(<2 x i32> %x) {
 ; CHECK-LABEL: @vec_select_no_equivalence(

llvm/test/Transforms/InstCombine/select-value-equivalence.ll

@@ -5,7 +5,7 @@ define <2 x i8> @select_icmp_insertelement_eq(<2 x i8> %x, <2 x i8> %y, i8 %i) {
 ; CHECK-LABEL: define <2 x i8> @select_icmp_insertelement_eq(
 ; CHECK-SAME: <2 x i8> [[X:%.*]], <2 x i8> [[Y:%.*]], i8 [[I:%.*]]) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i8> [[Y]], <i8 2, i8 2>
-; CHECK-NEXT:    [[INSERT:%.*]] = insertelement <2 x i8> [[Y]], i8 0, i8 [[I]]
+; CHECK-NEXT:    [[INSERT:%.*]] = insertelement <2 x i8> <i8 2, i8 2>, i8 0, i8 [[I]]
 ; CHECK-NEXT:    [[RETVAL:%.*]] = select <2 x i1> [[CMP]], <2 x i8> [[INSERT]], <2 x i8> [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[RETVAL]]
 ;
@@ -19,7 +19,7 @@ define <2 x i8> @select_icmp_insertelement_ne(<2 x i8> %x, <2 x i8> %y, i8 %i) {
 ; CHECK-LABEL: define <2 x i8> @select_icmp_insertelement_ne(
 ; CHECK-SAME: <2 x i8> [[X:%.*]], <2 x i8> [[Y:%.*]], i8 [[I:%.*]]) {
 ; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp eq <2 x i8> [[Y]], <i8 2, i8 2>
-; CHECK-NEXT:    [[INSERT:%.*]] = insertelement <2 x i8> [[Y]], i8 0, i8 [[I]]
+; CHECK-NEXT:    [[INSERT:%.*]] = insertelement <2 x i8> <i8 2, i8 2>, i8 0, i8 [[I]]
 ; CHECK-NEXT:    [[RETVAL:%.*]] = select <2 x i1> [[CMP_NOT]], <2 x i8> [[INSERT]], <2 x i8> [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[RETVAL]]
 ;
@@ -46,7 +46,7 @@ define <4 x i8> @select_icmp_shufflevector_select(<4 x i8> %x, <4 x i8> %y, <4 x
 ; CHECK-LABEL: define <4 x i8> @select_icmp_shufflevector_select(
 ; CHECK-SAME: <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i8> [[Z:%.*]]) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <4 x i8> [[Y]], <i8 2, i8 2, i8 2, i8 2>
-; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i8> [[Z]], <4 x i8> [[Y]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i8> [[Z]], <4 x i8> <i8 poison, i8 2, i8 poison, i8 2>, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
 ; CHECK-NEXT:    [[RETVAL:%.*]] = select <4 x i1> [[CMP]], <4 x i8> [[SHUFFLE]], <4 x i8> [[X]]
 ; CHECK-NEXT:    ret <4 x i8> [[RETVAL]]
 ;