[InstCombine] make icmp vector canonicalization safe for constant with undef elements

[spatel]

Summary

This is a fix for:
https://bugs.llvm.org/show_bug.cgi?id=43730
...and as shown there, we have an existing test case that shows a potential miscompile.

We could just bail out for vector constants that contain any undef elements, or we can do as shown here:
allow the transform, but replace the undefs with a safe value.

For the tests shown, this results in a full splat constant (no undefs) which is probably a win for further IR analysis because we conservatively don't match undefs in most cases. Codegen can probably recover these kinds of undef lanes via demanded elements analysis if that's profitable.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5197–5208

Time to hoist replaceUndefsWith() from llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp ?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5197–5208

Yes - I forgot where that was hiding. :)

[spatel]

Patch updated:
Hoist the existing function for replacing undefs and use it. I can make that part an NFC pre-commit, but I kept it all together here in the review for the context.

[lebedev.ri]

I didn't verify that alive no longer barfs on that transform,
but i do believe the change makes sense in the context of
the example provided in that bugreport.

Should there be a icmp sge test?

Other than that, thanks, LG.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5190

This will result in the last one being used

5198

with last safe constant we found

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5190

Good catch. Yes, the test coverage isn't sufficient. I'll add tests before splitting/committing/fixing.

llvm/include/llvm/IR/Constant.h

public:
   const Constant *stripPointerCasts() const {
     return cast<Constant>(Value::stripPointerCasts());
   }
 
   Constant *stripPointerCasts() {
     return const_cast<Constant*>(
                       static_cast<const Constant *>(this)->stripPointerCasts());
   }
+
+  /// Try to replace undefined constant C or undefined elements in C with
+  /// Replacement. If no changes are made, the constant C is returned.
+  static Constant *replaceUndefsWith(Constant *C, Constant *Replacement);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_IR_CONSTANT_H

llvm/lib/IR/Constants.cpp

 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 
 using namespace llvm;
+using namespace PatternMatch;
 
 //===----------------------------------------------------------------------===//
 //                              Constant Class
 //===----------------------------------------------------------------------===//
 
 bool Constant::isNegativeZeroValue() const {
   // Floating point values have an explicit -0.0 value.
   if (const ConstantFP *CFP = dyn_cast<ConstantFP>(this))

 bool Constant::isElementWiseEqual(Value *Y) const {
   // Are they fully identical?
   if (this == Y)
     return true;
   // They may still be identical element-wise (if they have `undef`s).
   auto *Cy = dyn_cast<Constant>(Y);
   if (!Cy)
     return false;
-  return PatternMatch::match(ConstantExpr::getICmp(ICmpInst::Predicate::ICMP_EQ,
-                                                   const_cast<Constant *>(this),
-                                                   Cy),
-                             PatternMatch::m_One());
+  return match(ConstantExpr::getICmp(ICmpInst::Predicate::ICMP_EQ,
+                                     const_cast<Constant *>(this), Cy),
+               m_One());
 }
 
 bool Constant::containsUndefElement() const {
   if (!getType()->isVectorTy())
     return false;
   for (unsigned i = 0, e = getType()->getVectorNumElements(); i != e; ++i)
     if (isa<UndefValue>(getAggregateElement(i)))
       return true;
while (I != E) {
     // If the constant was dead, then the iterator is invalidated.
     if (LastNonDeadUser == E)
       I = user_begin();
     else
       I = std::next(LastNonDeadUser);
   }
 }
 
+Constant *Constant::replaceUndefsWith(Constant *C, Constant *Replacement) {
+  Type *Ty = C->getType();
+  if (C && match(C, m_Undef())) {
+    assert(Ty == Replacement->getType() && "Expected matching types");
+    return Replacement;
+  }
+
+  // Don't know how to deal with this constant.
+  if (!Ty->isVectorTy())
+    return C;
+
+  unsigned NumElts = Ty->getVectorNumElements();
+  SmallVector<Constant *, 32> NewC(NumElts);
+  for (unsigned i = 0; i != NumElts; ++i) {
+    Constant *EltC = C->getAggregateElement(i);
+    assert(EltC->getType() == Replacement->getType() &&
+           "Expected matching types");
+    NewC[i] = EltC && match(EltC, m_Undef()) ? Replacement : EltC;
+  }
+  return ConstantVector::get(NewC);
+}
 
 
 //===----------------------------------------------------------------------===//
 //                                ConstantInt
 //===----------------------------------------------------------------------===//
 
 ConstantInt::ConstantInt(IntegerType *Ty, const APInt &V)
     : ConstantData(Ty, ConstantIntVal), Val(V) {

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

bool WillIncrement =
       UnsignedPred == ICmpInst::ICMP_ULE || UnsignedPred == ICmpInst::ICMP_UGT;
 
   // Check if the constant operand can be safely incremented/decremented
   // without overflowing/underflowing.
   auto ConstantIsOk = [WillIncrement, IsSigned](ConstantInt *C) {
     return WillIncrement ? !C->isMaxValue(IsSigned) : !C->isMinValue(IsSigned);
   };
 
+  Constant *SafeReplacementConstant = nullptr;
   if (auto *CI = dyn_cast<ConstantInt>(C)) {
     // Bail out if the constant can't be safely incremented/decremented.
     if (!ConstantIsOk(CI))
       return llvm::None;
   } else if (Type->isVectorTy()) {
     unsigned NumElts = Type->getVectorNumElements();
     for (unsigned i = 0; i != NumElts; ++i) {
       Constant *Elt = C->getAggregateElement(i);
       if (!Elt)
         return llvm::None;
 
       if (isa<UndefValue>(Elt))
         continue;
 
       // Bail out if we can't determine if this constant is min/max or if we
       // know that this constant is min/max.
       auto *CI = dyn_cast<ConstantInt>(Elt);
       if (!CI || !ConstantIsOk(CI))
         return llvm::None;
+      SafeReplacementConstant = CI;
     }
   } else {
     // ConstantExpr?
     return llvm::None;
   }
 
+  // It may not be safe to change a compare predicate in the presence of
+  // undefined elements, so replace those elements with the first safe constant
+  // that we found.
+  if (C->containsUndefElement()) {
+    assert(SafeReplacementConstant && "Replacement constant not set");
+    C = Constant::replaceUndefsWith(C, SafeReplacementConstant);
+  }
+
   CmpInst::Predicate NewPred = CmpInst::getFlippedStrictnessPredicate(Pred);
 
   // Increment or decrement the constant.
   Constant *OneOrNegOne = ConstantInt::get(Type, WillIncrement ? 1 : -1, true);
   Constant *NewC = ConstantExpr::getAdd(C, OneOrNegOne);
 
   return std::make_pair(NewPred, NewC);
 }
 
 /// If we have an icmp le or icmp ge instruction with a constant operand, turn
 /// it into the appropriate icmp lt or icmp gt instruction. This transform
 /// allows them to be folded in visitICmpInst.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

Value *InstCombiner::reassociateShiftAmtsOfTwoSameDirectionShifts(
   if (Trunc) {
     Builder.Insert(NewShift);
     Ret = CastInst::Create(Instruction::Trunc, NewShift, Sh0->getType());
   }
 
   return Ret;
 }
 
-// Try to replace `undef` constants in C with Replacement.
-static Constant *replaceUndefsWith(Constant *C, Constant *Replacement) {
-  if (C && match(C, m_Undef()))
-    return Replacement;
-
-  if (auto *CV = dyn_cast<ConstantVector>(C)) {
-    llvm::SmallVector<Constant *, 32> NewOps(CV->getNumOperands());
-    for (unsigned i = 0, NumElts = NewOps.size(); i != NumElts; ++i) {
-      Constant *EltC = CV->getOperand(i);
-      NewOps[i] = EltC && match(EltC, m_Undef()) ? Replacement : EltC;
-    }
-    return ConstantVector::get(NewOps);
-  }
-
-  // Don't know how to deal with this constant.
-  return C;
-}
-
 // If we have some pattern that leaves only some low bits set, and then performs
 // left-shift of those bits, if none of the bits that are left after the final
 // shift are modified by the mask, we can omit the mask.
 //
 // There are many variants to this pattern:
 //   a)  (x & ((1 << MaskShAmt) - 1)) << ShiftShAmt
 //   b)  (x & (~(-1 << MaskShAmt))) << ShiftShAmt
 //   c)  (x & (-1 >> MaskShAmt)) << ShiftShAmt
if (!SumOfShAmts)
       return nullptr; // Did not simplify.
     // In this pattern SumOfShAmts correlates with the number of low bits
     // that shall remain in the root value (OuterShift).
 
     // An extend of an undef value becomes zero because the high bits are never
     // completely unknown. Replace the the `undef` shift amounts with final
     // shift bitwidth to ensure that the value remains undef when creating the
     // subsequent shift op.
-    SumOfShAmts = replaceUndefsWith(
+    SumOfShAmts = Constant::replaceUndefsWith(
         SumOfShAmts, ConstantInt::get(SumOfShAmts->getType()->getScalarType(),
                                       ExtendedTy->getScalarSizeInBits()));
     auto *ExtendedSumOfShAmts = ConstantExpr::getZExt(SumOfShAmts, ExtendedTy);
     // And compute the mask as usual: ~(-1 << (SumOfShAmts))
     auto *ExtendedAllOnes = ConstantExpr::getAllOnesValue(ExtendedTy);
     auto *ExtendedInvertedMask =
         ConstantExpr::getShl(ExtendedAllOnes, ExtendedSumOfShAmts);
     NewMask = ConstantExpr::getNot(ExtendedInvertedMask);
   } else if (match(Masked, m_c_And(m_CombineOr(MaskC, MaskD), m_Value(X))) ||
              match(Masked, m_Shr(m_Shl(m_Value(X), m_Value(MaskShAmt)),
                                  m_Deferred(MaskShAmt)))) {
     // Can we simplify (ShiftShAmt-MaskShAmt) ?
     auto *ShAmtsDiff = dyn_cast_or_null<Constant>(SimplifySubInst(
         ShiftShAmt, MaskShAmt, /*IsNSW=*/false, /*IsNUW=*/false, Q));
     if (!ShAmtsDiff)
       return nullptr; // Did not simplify.
     // In this pattern ShAmtsDiff correlates with the number of high bits that
     // shall be unset in the root value (OuterShift).
 
     // An extend of an undef value becomes zero because the high bits are never
     // completely unknown. Replace the the `undef` shift amounts with negated
     // bitwidth of innermost shift to ensure that the value remains undef when
     // creating the subsequent shift op.
     unsigned WidestTyBitWidth = WidestTy->getScalarSizeInBits();
-    ShAmtsDiff = replaceUndefsWith(
+    ShAmtsDiff = Constant::replaceUndefsWith(
         ShAmtsDiff, ConstantInt::get(ShAmtsDiff->getType()->getScalarType(),
                                      -WidestTyBitWidth));
     auto *ExtendedNumHighBitsToClear = ConstantExpr::getZExt(
         ConstantExpr::getSub(ConstantInt::get(ShAmtsDiff->getType(),
                                               WidestTyBitWidth,
                                               /*isSigned=*/false),
                              ShAmtsDiff),
         ExtendedTy);

llvm/test/Transforms/InstCombine/canonicalize-constant-low-bit-mask-and-icmp-eq-to-icmp-ule.ll

 ;
   %tmp0 = and <2 x i8> %x, <i8 3, i8 -1>
   %ret = icmp eq <2 x i8> %tmp0, %x
   ret <2 x i1> %ret
 }
 
 define <3 x i1> @p3_vec_splat_undef(<3 x i8> %x) {
 ; CHECK-LABEL: @p3_vec_splat_undef(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <3 x i8> [[X:%.*]], <i8 4, i8 undef, i8 4>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <3 x i8> [[X:%.*]], <i8 4, i8 4, i8 4>
 ; CHECK-NEXT:    ret <3 x i1> [[TMP1]]
 ;
   %tmp0 = and <3 x i8> %x, <i8 3, i8 undef, i8 3>
   %ret = icmp eq <3 x i8> %tmp0, %x
   ret <3 x i1> %ret
 }
 
 ; ============================================================================ ;

llvm/test/Transforms/InstCombine/canonicalize-constant-low-bit-mask-and-icmp-sge-to-icmp-sle.ll

 ;
   %tmp0 = and <2 x i8> %x, <i8 3, i8 0>
   %ret = icmp sge <2 x i8> %tmp0, %x
   ret <2 x i1> %ret
 }
 
 define <3 x i1> @p3_vec_splat_undef(<3 x i8> %x) {
 ; CHECK-LABEL: @p3_vec_splat_undef(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt <3 x i8> [[X:%.*]], <i8 4, i8 undef, i8 4>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt <3 x i8> [[X:%.*]], <i8 4, i8 4, i8 4>
 ; CHECK-NEXT:    ret <3 x i1> [[TMP1]]
 ;
   %tmp0 = and <3 x i8> %x, <i8 3, i8 undef, i8 3>
   %ret = icmp sge <3 x i8> %tmp0, %x
   ret <3 x i1> %ret
 }
 
 ; ============================================================================ ;

llvm/test/Transforms/InstCombine/canonicalize-constant-low-bit-mask-and-icmp-sle-to-icmp-sle.ll

;
   %tmp0 = and <2 x i8> %x, <i8 3, i8 0>
   %ret = icmp sle <2 x i8> %x, %tmp0
   ret <2 x i1> %ret
 }
 
 define <3 x i1> @p3_vec_splat_undef() {
 ; CHECK-LABEL: @p3_vec_splat_undef(
 ; CHECK-NEXT:    [[X:%.*]] = call <3 x i8> @gen3x8()
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt <3 x i8> [[X]], <i8 4, i8 undef, i8 4>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt <3 x i8> [[X]], <i8 4, i8 4, i8 4>
 ; CHECK-NEXT:    ret <3 x i1> [[TMP1]]
 ;
   %x = call <3 x i8> @gen3x8()
   %tmp0 = and <3 x i8> %x, <i8 3, i8 undef, i8 3>
   %ret = icmp sle <3 x i8> %x, %tmp0
   ret <3 x i1> %ret
 }
 

llvm/test/Transforms/InstCombine/canonicalize-constant-low-bit-mask-and-icmp-uge-to-icmp-ule.ll

 ;
   %tmp0 = and <2 x i8> %x, <i8 3, i8 -1>
   %ret = icmp uge <2 x i8> %tmp0, %x
   ret <2 x i1> %ret
 }
 
 define <3 x i1> @p3_vec_splat_undef(<3 x i8> %x) {
 ; CHECK-LABEL: @p3_vec_splat_undef(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <3 x i8> [[X:%.*]], <i8 4, i8 undef, i8 4>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <3 x i8> [[X:%.*]], <i8 4, i8 4, i8 4>
 ; CHECK-NEXT:    ret <3 x i1> [[TMP1]]
 ;
   %tmp0 = and <3 x i8> %x, <i8 3, i8 undef, i8 3>
   %ret = icmp uge <3 x i8> %tmp0, %x
   ret <3 x i1> %ret
 }
 
 ; ============================================================================ ;

llvm/test/Transforms/InstCombine/canonicalize-constant-low-bit-mask-and-icmp-ule-to-icmp-ule.ll

;
   %tmp0 = and <2 x i8> %x, <i8 3, i8 -1>
   %ret = icmp ule <2 x i8> %x, %tmp0
   ret <2 x i1> %ret
 }
 
 define <3 x i1> @p3_vec_splat_undef() {
 ; CHECK-LABEL: @p3_vec_splat_undef(
 ; CHECK-NEXT:    [[X:%.*]] = call <3 x i8> @gen3x8()
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <3 x i8> [[X]], <i8 4, i8 undef, i8 4>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult <3 x i8> [[X]], <i8 4, i8 4, i8 4>
 ; CHECK-NEXT:    ret <3 x i1> [[TMP1]]
 ;
   %x = call <3 x i8> @gen3x8()
   %tmp0 = and <3 x i8> %x, <i8 3, i8 undef, i8 3>
   %ret = icmp ule <3 x i8> %x, %tmp0
   ret <3 x i1> %ret
 }
 

llvm/test/Transforms/InstCombine/icmp-vec.ll

;
   %cmp = icmp sle <2 x i8> %a, <i8 bitcast (<2 x i4> <i4 1, i4 2> to i8), i8 0>
   ret <2 x i1> %cmp
 }
 
 ; Undef elements don't prevent the transform of the comparison.
 
 define <2 x i1> @PR27756_2(<2 x i8> %a) {
 ; CHECK-LABEL: @PR27756_2(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt <2 x i8> [[A:%.*]], <i8 undef, i8 1>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt <2 x i8> [[A:%.*]], <i8 1, i8 1>
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %cmp = icmp sle <2 x i8> %a, <i8 undef, i8 0>
   ret <2 x i1> %cmp
 }
 
 @someglobal = global i32 0
 

llvm/test/Transforms/InstCombine/reuse-constant-from-select-in-icmp.ll

;
   %r = select <2 x i1> %t, <2 x i32> %y, <2 x i32> <i32 65535, i32 65535>
   ret <2 x i32> %r
 }
 
 ; Vectors with undef
 
 define <2 x i32> @p8_vec_nonsplat_undef0(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @p8_vec_nonsplat_undef0(
-; CHECK-NEXT:    [[T_INV:%.*]] = icmp ugt <2 x i32> [[X:%.*]], <i32 65535, i32 undef>
+; CHECK-NEXT:    [[T_INV:%.*]] = icmp ugt <2 x i32> [[X:%.*]], <i32 65535, i32 65535>
 ; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[T_INV]], <2 x i32> <i32 65535, i32 65535>, <2 x i32> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %t = icmp ult <2 x i32> %x, <i32 65536, i32 undef>
   %r = select <2 x i1> %t, <2 x i32> %y, <2 x i32> <i32 65535, i32 65535>
   ret <2 x i32> %r
 }
 define <2 x i32> @p9_vec_nonsplat_undef1(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @p9_vec_nonsplat_undef1(
 ; CHECK-NEXT:    [[T_INV:%.*]] = icmp ugt <2 x i32> [[X:%.*]], <i32 65535, i32 65535>
 ; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[T_INV]], <2 x i32> <i32 65535, i32 undef>, <2 x i32> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %t = icmp ult <2 x i32> %x, <i32 65536, i32 65536>
   %r = select <2 x i1> %t, <2 x i32> %y, <2 x i32> <i32 65535, i32 undef>
   ret <2 x i32> %r
 }
 define <2 x i32> @p10_vec_nonsplat_undef2(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @p10_vec_nonsplat_undef2(
-; CHECK-NEXT:    [[T_INV:%.*]] = icmp ugt <2 x i32> [[X:%.*]], <i32 65535, i32 undef>
+; CHECK-NEXT:    [[T_INV:%.*]] = icmp ugt <2 x i32> [[X:%.*]], <i32 65535, i32 65535>
 ; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[T_INV]], <2 x i32> <i32 65535, i32 undef>, <2 x i32> [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %t = icmp ult <2 x i32> %x, <i32 65536, i32 undef>
   %r = select <2 x i1> %t, <2 x i32> %y, <2 x i32> <i32 65535, i32 undef>
   ret <2 x i32> %r
 }