[InstCombine] fix shuffle-of-binops transform to avoid poison/undef

As noted in D48987, there are many different ways for this transform to go wrong. 
In particular, the poison potential for shifts means we have to more careful with those ops. 
I added tests to make that behavior visible for all of the different cases that I could find.

This is a partial fix. To make this review easier, I did not make changes for the single binop 
pattern (handled in foldSelectShuffleWith1Binop()). I also left out some potential optimizations 
noted with TODO comments. I'll follow-up once we're confident that things are correct here.

The goal is to correct all marked FIXME tests to either avoid the shuffle transform or do it safely.

Note that distinguishing when the shuffle mask contains undefs and using getBinOpIdentity() allows 
for some improvements to div/rem patterns, so there are wins along with the missed opportunities 
and fixes.

Differential Revision: https://reviews.llvm.org/D49047


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@336546 91177308-0d34-0410-b5e6-96231b3b80d8

Author: rotateright

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -1280,52 +1280,83 @@ static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
   // The opcodes must be the same. Use a new name to make that clear.
   BinaryOperator::BinaryOps BOpc = Opc0;
 
+  // We are moving a binop after a shuffle. When a shuffle has an undefined
+  // mask element, the result is undefined, but it is not poison or undefined
+  // behavior. That is not necessarily true for div/rem/shift.
+  Constant *Mask = Shuf.getMask();
+  bool MightCreatePoisonOrUB =
+      Mask->containsUndefElement() &&
+      (Instruction::isIntDivRem(BOpc) || Instruction::isShift(BOpc));
+
+  Constant *NewC;
   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);
+    }
+
     // 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())) {
+  } else {
     // 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.
-    //
+    if (!B0->hasOneUse() && !B1->hasOneUse())
+      return nullptr;
+
+    // Bail out if we can not guarantee safety of the transform.
+    // TODO: We could try harder by replacing the undef mask elements.
+    if (MightCreatePoisonOrUB)
+      return nullptr;
+
     // 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;
+    NewC = ConstantExpr::getShuffleVector(C0, C1, Mask);
+    V = Builder.CreateShuffleVector(X, Y, Mask);
   }
 
-  Constant *NewC = ConstantExpr::getShuffleVector(C0, C1, Shuf.getMask());
-
-  // If the shuffle mask contains undef elements, then the new constant
-  // vector will have undefs in those lanes. This could cause the entire
-  // binop to be undef.
-  if (Instruction::isIntDivRem(BOpc))
-    NewC = getSafeVectorConstantForIntDivRem(NewC);
-
   Instruction *NewBO = ConstantsAreOp1 ? BinaryOperator::Create(BOpc, V, NewC) :
                                          BinaryOperator::Create(BOpc, NewC, V);
 
-  // Flags are intersected from the 2 source binops.
+  // Flags are intersected from the 2 source binops. But there are 2 exceptions:
+  // 1. If we changed an opcode, poison conditions might have changed.
+  // 2. If the shuffle had undef mask elements, the new binop might have undefs
+  //    where the original code did not. Drop all poison potential to be safe.
   NewBO->copyIRFlags(B0);
   NewBO->andIRFlags(B1);
   if (DropNSW)
     NewBO->setHasNoSignedWrap(false);
+  if (Mask->containsUndefElement())
+    NewBO->dropPoisonGeneratingFlags();
   return NewBO;
 }