[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

5199–5210

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

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5199–5210

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

5200

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/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;
+
+      if (!SafeReplacementConstant)
+        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/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 <3 x i1> @PR27756_2(<3 x i8> %a) {
 ; CHECK-LABEL: @PR27756_2(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt <3 x i8> [[A:%.*]], <i8 43, i8 undef, i8 1>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt <3 x i8> [[A:%.*]], <i8 43, i8 43, i8 1>
 ; CHECK-NEXT:    ret <3 x i1> [[CMP]]
 ;
   %cmp = icmp sle <3 x i8> %a, <i8 42, i8 undef, i8 0>
   ret <3 x i1> %cmp
 }
 
 define <3 x i1> @PR27756_3(<3 x i8> %a) {
 ; CHECK-LABEL: @PR27756_3(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt <3 x i8> [[A:%.*]], <i8 undef, i8 0, i8 41>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt <3 x i8> [[A:%.*]], <i8 0, i8 0, i8 41>
 ; CHECK-NEXT:    ret <3 x i1> [[CMP]]
 ;
   %cmp = icmp sge <3 x i8> %a, <i8 undef, i8 1, i8 42>
   ret <3 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
 }