[InstCombine] drop poison flags for shuffle transforms with undefs

[spatel]

Summary

As discussed in D48893, we have an existing bug in the recent shuffle transforms when we combine poison flags with undefs in the shuffle mask.

There are 2 ways to solve that bug: (1) change the undef constants to safe constants (which we must do for div/rem because they have immediate UB potential), or (2) drop the poison flags.

I suggested using safe constants in D48893, but I changed my mind for 2 reasons: (1) we can't always be sure what qualifies as a safe constant (operand 0 of a binop?) and (2) undef vector lanes are likely more valuable for further analysis and vector transforms than any poison decorations.

[lebedev.ri]

So my concerns weren't unfounded :/
Do we want to also do such kind of fixup in VisitShl() in instcombine, for every shl we visit?

[spatel]

In D48987#1153565, @lebedev.ri wrote:

So my concerns weren't unfounded :/
Do we want to also do such kind of fixup in VisitShl() in instcombine, for every shl we visit?

Do you have an example of the problem you're thinking of? AFAIK, we won't hit this problem with scalars or general transforms because they don't modify constants like we're doing here. We do have a related foldShuffledBinop() function that I will review/fix if this patch looks right.

[spatel]

In D48987#1153586, @spatel wrote:

In D48987#1153565, @lebedev.ri wrote:

So my concerns weren't unfounded :/
Do we want to also do such kind of fixup in VisitShl() in instcombine, for every shl we visit?

Do you have an example of the problem you're thinking of? AFAIK, we won't hit this problem with scalars or general transforms because they don't modify constants like we're doing here. We do have a related foldShuffledBinop() function that I will review/fix if this patch looks right.

Note also that most of instcombine drops flags by default when it does a transform, and we generally favor dropping flags if that results in a better instruction canonicalization.

My main concern with these recent shuffle patches was not dropping fast-math-flags because that can inhibit later optimizations. I was too aggressive by trying to preserve poison flags along with those.

[lebedev.ri]

test/Transforms/InstCombine/shuffle_select.ll

30–42

In D48987#1153586, @spatel wrote:

In D48987#1153565, @lebedev.ri wrote:

So my concerns weren't unfounded :/
Do we want to also do such kind of fixup in VisitShl() in instcombine, for every shl we visit?

Do you have an example of the problem you're thinking of? AFAIK, we won't hit this problem with scalars or general transforms because they don't modify constants like we're doing here. We do have a related foldShuffledBinop() function that I will review/fix if this patch looks right.

https://godbolt.org/g/y7hCki

Should instcombine be dropping these flags just because?
Or we only care about not creating more of these situations than we already had?

[spatel]

test/Transforms/InstCombine/shuffle_select.ll

30–42

No, we don't want to drop flags on existing instructions. Presumably, whoever created that instruction did so correctly, so there actually is poison in the vector lane with undef.

We can't create poison where there wasn't any to start with, but we also don't want to remove poison when it exists without good reason. At least that's my understanding...happy to be corrected by people closer to poison/undef than me. :)

[lebedev.ri]

LG, but maybe @nlopes wants to comment.

test/Transforms/InstCombine/shuffle_select.ll

30–42

Ok, makes sense :)

[nlopes]

test/Transforms/InstCombine/shuffle_select.ll

521

Everything looks good, except this one.
As discussed in the other patch, a << undef is poison, so you can't introduce it. (since it's not correct to replace undef with poison).

[spatel]

test/Transforms/InstCombine/shuffle_select.ll

521

Hmm...how is this case different than the 'shl_shl' above?

Are there 4 cases?

1. Instructions with immediate UB (srem, urem).
2. Instructions with immediate UB and poison with flags (sdiv, udiv).
3. Instructions with poison regardless of flags (shl, ashr, lshr).
4. Instructions with poison only with flags (mul, add, sub).

[nlopes]

test/Transforms/InstCombine/shuffle_select.ll

521

This case is the same as shl_shl, yes.

Your 4 cases are correct. I would only add that case 3 is poison regardless of flags + poison with flags.

[spatel]

Abandoning - see D49047 instead.

include/llvm/IR/Constant.h

public:
   /// Return true if this scalar has an exact multiplicative inverse or this
   /// vector has an exact multiplicative inverse for each element in the vector.
   bool hasExactInverseFP() const;
 
   /// Return true if this is a floating-point NaN constant or a vector
   /// floating-point constant with all NaN elements.
   bool isNaN() const;
 
+  /// Return true if this is a vector constant that includes undefined elements.
+  bool containsUndefElements() const;
+
   /// Return true if evaluation of this constant could trap. This is true for
   /// things like constant expressions that could divide by zero.
   bool canTrap() const;
 
   /// Return true if the value can vary between threads.
   bool isThreadDependent() const;
 
   /// Return true if the value is dependent on a dllimport variable.

lib/IR/Constants.cpp

bool Constant::isNaN() const {
   for (unsigned i = 0, e = getType()->getVectorNumElements(); i != e; ++i) {
     auto *CFP = dyn_cast_or_null<ConstantFP>(this->getAggregateElement(i));
     if (!CFP || !CFP->isNaN())
       return false;
   }
   return true;
 }
 
+bool Constant::containsUndefElements() const {
+  if (!getType()->isVectorTy())
+    return false;
+  for (unsigned i = 0, e = getType()->getVectorNumElements(); i != e; ++i)
+    if (isa<UndefValue>(getAggregateElement(i)))
+      return true;
+
+  return false;
+}
+
 /// Constructor to create a '0' constant of arbitrary type.
 Constant *Constant::getNullValue(Type *Ty) {
   switch (Ty->getTypeID()) {
   case Type::IntegerTyID:
     return ConstantInt::get(Ty, 0);
   case Type::HalfTyID:
     return ConstantFP::get(Ty->getContext(),
                            APFloat::getZero(APFloat::IEEEhalf()));

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

case Instruction::Or: {
       break;
     }
     default:
       break;
   }
   return {};
 }
 
+/// Helper function for the shuffle transforms below. If a constant operand
+/// contains undefs, we can not safely propagate poison-generating flags to the
+/// instruction. The flags would allow other transforms to assume this binop
+/// produces poison when there might have been no poison potential in the
+/// original code. We choose to drop the flags rather than replace the undef
+/// values with safe constants because that is more likely to be better for
+/// further analysis and codegen.
+static BinaryOperator *getPoisonSafeVectorInst(BinaryOperator *BO) {
+  assert(BO->getType()->isVectorTy() &&
+         "Undef with poison special-case only applies to vector types");
+  Constant *C;
+  if (match(BO->getOperand(0), m_Constant(C)) ||
+      match(BO->getOperand(1), m_Constant(C)))
+    if (C->containsUndefElements())
+      BO->dropPoisonGeneratingFlags();
+  return BO;
+}
+
 static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
   assert(Shuf.isSelect() && "Must have select-equivalent shuffle");
 
   // Are we shuffling together some value and that same value after it has been
   // modified by a binop with a constant?
   Value *Op0 = Shuf.getOperand(0), *Op1 = Shuf.getOperand(1);
   Constant *C;
   bool Op0IsBinop;
static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
   // Example: shuf X, (add X, {-1,-2,-3,-4}), {0,1,6,7} --> add X, {0,0,-3,-4}
   // The existing binop constant vector remains in the same operand position.
   Constant *Mask = Shuf.getMask();
   Constant *NewC = Op0IsBinop ? ConstantExpr::getShuffleVector(C, IdC, Mask) :
                                 ConstantExpr::getShuffleVector(IdC, C, Mask);
 
   // shuf (bop X, C), X, M --> bop X, C'
   // shuf X, (bop X, C), M --> bop X, C'
-  Instruction *NewBO = BinaryOperator::Create(BOpcode, X, NewC);
+  auto *NewBO = BinaryOperator::Create(BOpcode, X, NewC);
   NewBO->copyIRFlags(BO);
-  return NewBO;
+  return getPoisonSafeVectorInst(NewBO);
 }
 
 /// Try to fold shuffles that are the equivalent of a vector select.
 static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
                                       InstCombiner::BuilderTy &Builder,
                                       const DataLayout &DL) {
   if (!Shuf.isSelect())
     return nullptr;
static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
   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);
+  auto *NewBO = ConstantsAreOp1 ? BinaryOperator::Create(BOpc, V, NewC) :
 
   // Flags are intersected from the 2 source binops.
   NewBO->copyIRFlags(B0);
   NewBO->andIRFlags(B1);
   if (DropNSW)
     NewBO->setHasNoSignedWrap(false);
-  return NewBO;
+
+  // We handled div/rem and a special-case for mul above, but there's one more
+  // problem with undef/poison.
+  return getPoisonSafeVectorInst(NewBO);
 }
 
 Instruction *InstCombiner::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   Value *LHS = SVI.getOperand(0);
   Value *RHS = SVI.getOperand(1);
   SmallVector<int, 16> Mask = SVI.getShuffleMask();
   Type *Int32Ty = Type::getInt32Ty(SVI.getContext());
 

test/Transforms/InstCombine/shuffle_select.ll

 ; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[V]], <4 x i32> [[B]], <4 x i32> <i32 0, i32 1, i32 2, i32 7>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = sub <4 x i32> <i32 11, i32 12, i32 13, i32 14>, %v
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 7>
   ret <4 x i32> %s
 }
 
 ; If any element of the shuffle mask operand is undef, that element of the result is undef.
 ; The shuffle is eliminated in this transform, but we can replace a constant element with undef.
-; Preserve flags when possible.
+; Preserve flags when possible. It's not safe to propagate poison-generating flags with undef constants.
 
 define <4 x i32> @mul(<4 x i32> %v) {
 ; CHECK-LABEL: @mul(
-; CHECK-NEXT:    [[S:%.*]] = mul nuw nsw <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 1, i32 14>
+; CHECK-NEXT:    [[S:%.*]] = mul <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 1, i32 14>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = mul nsw nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
   ret <4 x i32> %s
 }
 
 define <4 x i32> @shl(<4 x i32> %v) {
 ; CHECK-LABEL: @shl(
 ; CHECK-NEXT:    [[B:%.*]] = shl nuw <4 x i32> [[V:%.*]], <i32 11, i32 12, i32 13, i32 14>
 ; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[B]], <4 x i32> [[V]], <4 x i32> <i32 undef, i32 5, i32 2, i32 undef>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = shl nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>

 ; CHECK-NEXT:    ret <4 x i32> [[T3]]
 ;
   %t1 = mul <4 x i32> %v0, <i32 1, i32 2, i32 3, i32 4>
   %t2 = mul <4 x i32> %v0, <i32 5, i32 6, i32 7, i32 8>
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
   ret <4 x i32> %t3
 }
 
-; Preserve flags when possible.
+; Preserve flags when possible. It's not safe to propagate poison-generating flags with undef constants.
 
 define <4 x i32> @shl_shl(<4 x i32> %v0) {
 ; CHECK-LABEL: @shl_shl(
-; CHECK-NEXT:    [[T3:%.*]] = shl nuw <4 x i32> [[V0:%.*]], <i32 undef, i32 6, i32 3, i32 undef>
+; CHECK-NEXT:    [[T3:%.*]] = shl <4 x i32> [[V0:%.*]], <i32 undef, i32 6, i32 3, i32 undef>
 ; CHECK-NEXT:    ret <4 x i32> [[T3]]
 ;
   %t1 = shl nuw <4 x i32> %v0, <i32 1, i32 2, i32 3, i32 4>
   %t2 = shl nuw <4 x i32> %v0, <i32 5, i32 6, i32 7, i32 8>
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 undef, i32 5, i32 2, i32 undef>
   ret <4 x i32> %t3
 }
 

 ; CHECK-NEXT:    ret <4 x i32> [[T3]]
 ;
   %t1 = mul <4 x i32> %v0, <i32 1, i32 2, i32 3, i32 4>
   %t2 = mul <4 x i32> %v1, <i32 5, i32 6, i32 7, i32 8>
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
   ret <4 x i32> %t3
 }
 
-; Preserve flags when possible.
+; Preserve flags when possible. It's not safe to propagate poison-generating flags with undef constants.
 
 define <4 x i32> @shl_2_vars(<4 x i32> %v0, <4 x i32> %v1) {
 ; CHECK-LABEL: @shl_2_vars(
 ; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[V0:%.*]], <4 x i32> [[V1:%.*]], <4 x i32> <i32 undef, i32 5, i32 2, i32 undef>
-; CHECK-NEXT:    [[T3:%.*]] = shl nsw <4 x i32> [[TMP1]], <i32 undef, i32 6, i32 3, i32 undef>
+; CHECK-NEXT:    [[T3:%.*]] = shl <4 x i32> [[TMP1]], <i32 undef, i32 6, i32 3, i32 undef>
 ; CHECK-NEXT:    ret <4 x i32> [[T3]]
 ;
   %t1 = shl nsw <4 x i32> %v0, <i32 1, i32 2, i32 3, i32 4>
   %t2 = shl nsw <4 x i32> %v1, <i32 5, i32 6, i32 7, i32 8>
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 undef, i32 5, i32 2, i32 undef>
   ret <4 x i32> %t3
 }