2 files changed, 24 insertions, 13 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 23888526aae..75fec3e79d3 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -212,6 +212,23 @@ IntrinsicIDToOverflowCheckFlavor(unsigned ID) {
   }
 }
 
+/// Integer division/remainder require special handling to avoid undefined
+/// behavior. If a constant vector has undef elements, replace those undefs with
+/// '1' because that's always safe to execute.
+static inline Constant *getSafeVectorConstantForIntDivRem(Constant *In) {
+  assert(In->getType()->isVectorTy() && "Not expecting scalars here");
+  assert(In->getType()->getVectorElementType()->isIntegerTy() &&
+         "Not expecting FP opcodes/operands/constants here");
+
+  unsigned NumElts = In->getType()->getVectorNumElements();
+  SmallVector<Constant *, 16> Out(NumElts);
+  for (unsigned i = 0; i != NumElts; ++i) {
+    Constant *C = In->getAggregateElement(i);
+    Out[i] = isa<UndefValue>(C) ? ConstantInt::get(C->getType(), 1) : C;
+  }
+  return ConstantVector::get(Out);
+}
+
 /// The core instruction combiner logic.
 ///
 /// This class provides both the logic to recursively visit instructions and
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 1281c99f8dd..84535f971b5 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -1417,23 +1417,17 @@ Instruction *InstCombiner::foldShuffledBinop(BinaryOperator &Inst) {
       }
     }
     if (MayChange) {
+      Constant *NewC = ConstantVector::get(NewVecC);
       // With integer div/rem instructions, it is not safe to use a vector with
       // undef elements because the entire instruction can be folded to undef.
-      // So replace undef elements with '1' because that can never induce
-      // undefined behavior. All other binop opcodes are always safe to
-      // speculate, and therefore, it is fine to include undef elements for
-      // unused lanes (and using undefs may help optimization).
-      if (Inst.isIntDivRem()) {
-        assert(C->getType()->getScalarType()->isIntegerTy() &&
-               "Not expecting FP opcodes/operands/constants here");
-        for (unsigned i = 0; i < VWidth; ++i)
-          if (isa<UndefValue>(NewVecC[i]))
-            NewVecC[i] = ConstantInt::get(NewVecC[i]->getType(), 1);
-      }
-
+      // All other binop opcodes are always safe to speculate, and therefore, it
+      // is fine to include undef elements for unused lanes (and using undefs
+      // may help optimization).
+      if (Inst.isIntDivRem())
+        NewC = getSafeVectorConstantForIntDivRem(NewC);
+      
       // Op(shuffle(V1, Mask), C) -> shuffle(Op(V1, NewC), Mask)
       // Op(C, shuffle(V1, Mask)) -> shuffle(Op(NewC, V1), Mask)
-      Constant *NewC = ConstantVector::get(NewVecC);
       Value *NewLHS = isa<Constant>(LHS) ? NewC : V1;
       Value *NewRHS = isa<Constant>(LHS) ? V1 : NewC;
       return createBinOpShuffle(NewLHS, NewRHS, Mask);