1 files changed, 38 insertions, 12 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index d4fb0863c33..dd579a66325 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -1140,7 +1140,8 @@ static bool isShuffleExtractingFromLHS(ShuffleVectorInst &SVI,
   return true;
 }
 
-static Instruction *foldSelectShuffles(ShuffleVectorInst &Shuf) {
+static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
+                                      InstCombiner::BuilderTy &Builder) {
   // Folds under here require the equivalent of a vector select.
   if (!Shuf.isSelect())
     return nullptr;
@@ -1150,16 +1151,14 @@ static Instruction *foldSelectShuffles(ShuffleVectorInst &Shuf) {
       !match(Shuf.getOperand(1), m_BinOp(B1)))
     return nullptr;
 
-  // TODO: Fold the case with different variable operands (requires creating a
-  // new shuffle and checking number of uses).
-  Value *X;
+  Value *X, *Y;
   Constant *C0, *C1;
   bool ConstantsAreOp1;
   if (match(B0, m_BinOp(m_Value(X), m_Constant(C0))) &&
-      match(B1, m_BinOp(m_Specific(X), m_Constant(C1))))
+      match(B1, m_BinOp(m_Value(Y), m_Constant(C1))))
     ConstantsAreOp1 = true;
   else if (match(B0, m_BinOp(m_Constant(C0), m_Value(X))) &&
-           match(B1, m_BinOp(m_Constant(C1), m_Specific(X))))
+           match(B1, m_BinOp(m_Constant(C1), m_Value(Y))))
     ConstantsAreOp1 = false;
   else
     return nullptr;
@@ -1191,9 +1190,36 @@ static Instruction *foldSelectShuffles(ShuffleVectorInst &Shuf) {
   // The opcodes must be the same. Use a new name to make that clear.
   BinaryOperator::BinaryOps BOpc = Opc0;
 
-  // Remove a binop and the shuffle by rearranging the constant:
-  // shuffle (op X, C0), (op X, C1), M --> op X, C'
-  // shuffle (op C0, X), (op C1, X), M --> op C', X
+  Value *V;
+  if (X == Y) {
+    // Remove a binop and the shuffle by rearranging the constant:
+    // shuffle (op V, C0), (op V, C1), M --> op V, C'
+    // shuffle (op C0, V), (op C1, V), M --> op C', V
+    V = X;
+  } else if (!Instruction::isIntDivRem(BOpc) &&
+             (B0->hasOneUse() || B1->hasOneUse())) {
+    // If there are 2 different variable operands, we must create a new shuffle
+    // (select) first, so check uses to ensure that we don't end up with more
+    // instructions than we started with.
+    //
+    // Note: In general, we do not create new shuffles in InstCombine because we
+    // do not know if a target can lower an arbitrary shuffle optimally. In this
+    // case, the shuffle uses the existing mask, so there is no additional risk.
+    //
+    // TODO: We are disallowing div/rem because a shuffle with an undef mask
+    // element would propagate an undef value to the div/rem. That's not
+    // safe in general because div/rem allow for undefined behavior. We can
+    // loosen this restriction (eg, check if the mask has no undefs or replace
+    // undef elements).
+
+    // Select the variable vectors first, then perform the binop:
+    // shuffle (op X, C0), (op Y, C1), M --> op (shuffle X, Y, M), C'
+    // shuffle (op C0, X), (op C1, Y), M --> op C', (shuffle X, Y, M)
+    V = Builder.CreateShuffleVector(X, Y, Shuf.getMask());
+  } else {
+    return nullptr;
+  }
+
   Constant *NewC = ConstantExpr::getShuffleVector(C0, C1, Shuf.getMask());
 
   // If the shuffle mask contains undef elements, then the new constant
@@ -1202,8 +1228,8 @@ static Instruction *foldSelectShuffles(ShuffleVectorInst &Shuf) {
   if (Instruction::isIntDivRem(BOpc))
     NewC = getSafeVectorConstantForIntDivRem(NewC);
 
-  Instruction *NewBO = ConstantsAreOp1 ? BinaryOperator::Create(BOpc, X, NewC) :
-                                         BinaryOperator::Create(BOpc, NewC, X);
+  Instruction *NewBO = ConstantsAreOp1 ? BinaryOperator::Create(BOpc, V, NewC) :
+                                         BinaryOperator::Create(BOpc, NewC, V);
 
   // Flags are intersected from the 2 source binops.
   NewBO->copyIRFlags(B0);
@@ -1223,7 +1249,7 @@ Instruction *InstCombiner::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
           LHS, RHS, SVI.getMask(), SVI.getType(), SQ.getWithInstruction(&SVI)))
     return replaceInstUsesWith(SVI, V);
 
-  if (Instruction *I = foldSelectShuffles(SVI))
+  if (Instruction *I = foldSelectShuffle(SVI, Builder))
     return I;
 
   bool MadeChange = false;