[InstCombine] generalize safe vector constant utility

This is almost NFC, but there could be some case where the original
code had undefs in the constants (rather than just the shuffle mask),
and we'll use safe constants rather than undefs now.

The FIXME noted in foldShuffledBinop() is already visible in existing
tests, so correcting that is the next step.

llvm-svn: 336558

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

@@ -212,19 +212,26 @@ IntrinsicIDToOverflowCheckFlavor(unsigned ID) {
   }
 }
 
-/// Integer division/remainder require special handling to avoid undefined
+/// Some binary operators require special handling to avoid poison and 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");
+/// identity constants because those are always safe to execute. If no identity
+/// constant exists, replace undef with '1' or '1.0'.
+static inline Constant *getSafeVectorConstantForBinop(
+      BinaryOperator::BinaryOps Opcode, Constant *In) {
+  Type *Ty = In->getType();
+  assert(Ty->isVectorTy() && "Not expecting scalars here");
 
-  unsigned NumElts = In->getType()->getVectorNumElements();
+  Type *EltTy = Ty->getVectorElementType();
+  Constant *IdentityC = ConstantExpr::getBinOpIdentity(Opcode, EltTy, true);
+  if (!IdentityC)
+    IdentityC = EltTy->isIntegerTy() ? ConstantInt::get(EltTy, 1):
+                                       ConstantFP::get(EltTy, 1.0);
+
+  unsigned NumElts = Ty->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;
+    Out[i] = isa<UndefValue>(C) ? IdentityC : C;
   }
   return ConstantVector::get(Out);
 }

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -1288,33 +1288,18 @@ static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
       Mask->containsUndefElement() &&
       (Instruction::isIntDivRem(BOpc) || Instruction::isShift(BOpc));
 
-  Constant *NewC;
+  // Select the constant elements needed for the single binop.
+  Constant *NewC = ConstantExpr::getShuffleVector(C0, C1, Mask);
   Value *V;
   if (X == Y) {
-    NewC = ConstantExpr::getShuffleVector(C0, C1, Mask);
-
     // The new binop constant must not have any potential for extra poison/UB.
     if (MightCreatePoisonOrUB) {
       // TODO: Use getBinOpAbsorber for LHS replacement constants?
       if (!ConstantsAreOp1)
         return nullptr;
 
-      Type *EltTy = Shuf.getType()->getVectorElementType();
-      auto *IdC = ConstantExpr::getBinOpIdentity(BOpc, EltTy, true);
-      if (!IdC)
-        return nullptr;
-
-      // Replace undef elements caused by the mask with identity constants.
-      NewC = ConstantExpr::getShuffleVector(C0, C1, Mask);
-      unsigned NumElts = Shuf.getType()->getVectorNumElements();
-      SmallVector<Constant *, 16> VectorOfNewC(NumElts);
-      for (unsigned i = 0; i != NumElts; i++) {
-        if (isa<UndefValue>(Mask->getAggregateElement(i)))
-          VectorOfNewC[i] = IdC;
-        else
-          VectorOfNewC[i] = NewC->getAggregateElement(i);
-      }
-      NewC = ConstantVector::get(VectorOfNewC);
+      // Replace undef elements with identity constants.
+      NewC = getSafeVectorConstantForBinop(BOpc, NewC);
     }
 
     // Remove a binop and the shuffle by rearranging the constant:
@@ -1340,7 +1325,6 @@ static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
     // 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)
-    NewC = ConstantExpr::getShuffleVector(C0, C1, Mask);
     V = Builder.CreateShuffleVector(X, Y, Mask);
   }
 

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -1423,8 +1423,10 @@ Instruction *InstCombiner::foldShuffledBinop(BinaryOperator &Inst) {
       // 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).
+      // FIXME: This transform is also not poison-safe. Eg, shift-by-undef would
+      // create poison that may not exist in the original code.
       if (Inst.isIntDivRem())
-        NewC = getSafeVectorConstantForIntDivRem(NewC);
+        NewC = getSafeVectorConstantForBinop(Inst.getOpcode(), NewC);
       
       // Op(shuffle(V1, Mask), C) -> shuffle(Op(V1, NewC), Mask)
       // Op(C, shuffle(V1, Mask)) -> shuffle(Op(NewC, V1), Mask)