[InstCombine] try to fold icmp with mismatched extended operands

[spatel]

Summary

If a value is known to be non-negative and zexted, that's the same thing as sexted.

So for the purpose of looking past the casts with an icmp, treat it as if it was a sext:
https://alive2.llvm.org/ce/z/_BDsGV

This is necessary, but not enough to solve the motivating problem:
https://github.com/llvm/llvm-project/issues/55013

[nikic]

Relevant: https://discourse.llvm.org/t/aggressive-conversion-of-sext-to-zext-blocks-indvarsimplify/61561/9?u=nikic

[spatel]

In D124419#3473046, @nikic wrote:

Relevant: https://discourse.llvm.org/t/aggressive-conversion-of-sext-to-zext-blocks-indvarsimplify/61561/9?u=nikic

Yes, I saw that. Putting flags on casts would probably change the way this is implemented, but it does not change the need for this transform? It seems similar to the recent bf09a925f21a - we'd end up using the ValueTracking call to annotate the zext somewhere else, then pattern match the 'zext nsw' here.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4740

Isn't the sext handling here redundant due to https://github.com/llvm/llvm-project/blob/76d1f5eaa87b3d1ae5f4c07114313350152c24f2/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp#L1506-L1508, which will canonicalize to zext?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4740

No - we are checking if Op1 (right-side of the cmp) is sext, then is Op0 source (zext on left-side) non-negative. The next clause is checking if Op1 is zext, then is Op1 source non-negative.

I coded it this way because we didn't actually test that Op1 is zext or sext (it could be some other cast). I can change that or add a code comment to better explain.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4740

Okay, so we're using CastOp1 == sext here as a proxy for CastOp0 == zext, because we know that it's not the same opcode, but must be one of zext or sext. This does seem pretty roundabout.

[spatel]

Patch updated:
Improve readability by clearly matching zext/sext and adding code comments.

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

if (Pred == ICmpInst::ICMP_UGT && (~*C).isPowerOf2()) {
     Constant *MaskC = ConstantInt::get(X->getType(), (*C + 1).zext(SrcBits));
     Value *And = Builder.CreateAnd(X, MaskC);
     return new ICmpInst(ICmpInst::ICMP_EQ, And, MaskC);
   }
 
   return nullptr;
 }
 
-static Instruction *foldICmpWithZextOrSext(ICmpInst &ICmp,
-                                           InstCombiner::BuilderTy &Builder) {
+Instruction *InstCombinerImpl::foldICmpWithZextOrSext(ICmpInst &ICmp) {
   assert(isa<CastInst>(ICmp.getOperand(0)) && "Expected cast for operand 0");
   auto *CastOp0 = cast<CastInst>(ICmp.getOperand(0));
   Value *X;
   if (!match(CastOp0, m_ZExtOrSExt(m_Value(X))))
     return nullptr;
 
   bool IsSignedExt = CastOp0->getOpcode() == Instruction::SExt;
   bool IsSignedCmp = ICmp.isSigned();
-  if (auto *CastOp1 = dyn_cast<CastInst>(ICmp.getOperand(1))) {
-    // If the signedness of the two casts doesn't agree (i.e. one is a sext
-    // and the other is a zext), then we can't handle this.
-    // TODO: This is too strict. We can handle some predicates (equality?).
-    if (CastOp0->getOpcode() != CastOp1->getOpcode())
+
+  // icmp Pred (ext X), (ext Y)
+  Value *Y;
+  if (match(ICmp.getOperand(1), m_ZExtOrSExt(m_Value(Y)))) {
+    bool IsZext0 = isa<ZExtOperator>(ICmp.getOperand(0));
+    bool IsZext1 = isa<ZExtOperator>(ICmp.getOperand(1));
+
+    // If we have mismatched casts, treat the zext of a non-negative source as
+    // a sext to simulate matching casts. Otherwise, we are done.
+    // TODO: Can we handle some predicates (equality) without non-negative?
+    if (IsZext0 != IsZext1) {
+      if ((IsZext0 && isKnownNonNegative(X, DL, 0, &AC, &ICmp, &DT)) ||
+          (IsZext1 && isKnownNonNegative(Y, DL, 0, &AC, &ICmp, &DT)))
+        IsSignedExt = true;
+      else
         return nullptr;
+    }
 
     // Not an extension from the same type?
-    Value *Y = CastOp1->getOperand(0);
     Type *XTy = X->getType(), *YTy = Y->getType();
     if (XTy != YTy) {
       // One of the casts must have one use because we are creating a new cast.
-      if (!CastOp0->hasOneUse() && !CastOp1->hasOneUse())
+      if (!ICmp.getOperand(0)->hasOneUse() && !ICmp.getOperand(1)->hasOneUse())
         return nullptr;
       // Extend the narrower operand to the type of the wider operand.
+      CastInst::CastOps CastOpcode =
+          IsSignedExt ? Instruction::SExt : Instruction::ZExt;
       if (XTy->getScalarSizeInBits() < YTy->getScalarSizeInBits())
-        X = Builder.CreateCast(CastOp0->getOpcode(), X, YTy);
+        X = Builder.CreateCast(CastOpcode, X, YTy);
       else if (YTy->getScalarSizeInBits() < XTy->getScalarSizeInBits())
-        Y = Builder.CreateCast(CastOp0->getOpcode(), Y, XTy);
+        Y = Builder.CreateCast(CastOpcode, Y, XTy);
       else
         return nullptr;
     }
 
     // (zext X) == (zext Y) --> X == Y
     // (sext X) == (sext Y) --> X == Y
     if (ICmp.isEquality())
       return new ICmpInst(ICmp.getPredicate(), X, Y);
if (CastOp0->getOpcode() == Instruction::PtrToInt &&
 
     if (NewOp1)
       return new ICmpInst(ICmp.getPredicate(), Op0Src, NewOp1);
   }
 
   if (Instruction *R = foldICmpWithTrunc(ICmp, Builder))
     return R;
 
-  return foldICmpWithZextOrSext(ICmp, Builder);
+  return foldICmpWithZextOrSext(ICmp);
 }
 
 static bool isNeutralValue(Instruction::BinaryOps BinaryOp, Value *RHS) {
   switch (BinaryOp) {
     default:
       llvm_unreachable("Unsupported binary op");
     case Instruction::Add:
     case Instruction::Sub:

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *foldCmpLoadFromIndexedGlobal(LoadInst *LI,
                                             GetElementPtrInst *GEP,
                                             GlobalVariable *GV, CmpInst &ICI,
                                             ConstantInt *AndCst = nullptr);
   Instruction *foldFCmpIntToFPConst(FCmpInst &I, Instruction *LHSI,
                                     Constant *RHSC);
   Instruction *foldICmpAddOpConst(Value *X, const APInt &C,
                                   ICmpInst::Predicate Pred);
-  Instruction *foldICmpWithCastOp(ICmpInst &ICI);
+  Instruction *foldICmpWithCastOp(ICmpInst &ICmp);
+  Instruction *foldICmpWithZextOrSext(ICmpInst &ICmp);
 
   Instruction *foldICmpUsingKnownBits(ICmpInst &Cmp);
   Instruction *foldICmpWithDominatingICmp(ICmpInst &Cmp);
   Instruction *foldICmpWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstantNotInt(ICmpInst &Cmp);
   Instruction *foldICmpBinOp(ICmpInst &Cmp, const SimplifyQuery &SQ);
   Instruction *foldICmpEquality(ICmpInst &Cmp);

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

;
   %b = zext i8 %y to i32
   %c = icmp uge i32 %a, %b
   ret i1 %c
 }
 
 define i1 @zext_sext_sgt_known_nonneg(i8 %x, i8 %y) {
 ; CHECK-LABEL: @zext_sext_sgt_known_nonneg(
 ; CHECK-NEXT:    [[N:%.*]] = udiv i8 127, [[X:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[B:%.*]] = sext i8 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i32 [[A]], [[B]]
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i8 [[N]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %n = udiv i8 127, %x
   %a = zext i8 %n to i32
   %b = sext i8 %y to i32
   %c = icmp sgt i32 %a, %b
   ret i1 %c
 }
 
 define i1 @zext_sext_ugt_known_nonneg(i8 %x, i8 %y) {
 ; CHECK-LABEL: @zext_sext_ugt_known_nonneg(
 ; CHECK-NEXT:    [[N:%.*]] = and i8 [[X:%.*]], 127
-; CHECK-NEXT:    [[A:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[B:%.*]] = sext i8 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp ugt i32 [[A]], [[B]]
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i8 [[N]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %n = and i8 %x, 127
   %a = zext i8 %n to i32
   %b = sext i8 %y to i32
   %c = icmp ugt i32 %a, %b
   ret i1 %c
 }
 
 define i1 @zext_sext_eq_known_nonneg(i8 %x, i8 %y) {
 ; CHECK-LABEL: @zext_sext_eq_known_nonneg(
 ; CHECK-NEXT:    [[N:%.*]] = lshr i8 [[X:%.*]], 1
-; CHECK-NEXT:    [[A:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[B:%.*]] = sext i8 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[A]], [[B]]
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[N]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %n = lshr i8 %x, 1
   %a = zext i8 %n to i32
   %b = sext i8 %y to i32
   %c = icmp eq i32 %a, %b
   ret i1 %c
 }
 
 define i1 @zext_sext_sle_known_nonneg_op0_narrow(i8 %x, i16 %y) {
 ; CHECK-LABEL: @zext_sext_sle_known_nonneg_op0_narrow(
 ; CHECK-NEXT:    [[N:%.*]] = and i8 [[X:%.*]], 12
-; CHECK-NEXT:    [[A:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[B:%.*]] = sext i16 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp sle i32 [[A]], [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[N]] to i16
+; CHECK-NEXT:    [[C:%.*]] = icmp sle i16 [[TMP1]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %n = and i8 %x, 12
   %a = zext i8 %n to i32
   %b = sext i16 %y to i32
   %c = icmp sle i32 %a, %b
   ret i1 %c
 }
 
 define i1 @zext_sext_ule_known_nonneg_op0_wide(i9 %x, i8 %y) {
 ; CHECK-LABEL: @zext_sext_ule_known_nonneg_op0_wide(
 ; CHECK-NEXT:    [[N:%.*]] = urem i9 [[X:%.*]], 254
-; CHECK-NEXT:    [[A:%.*]] = zext i9 [[N]] to i32
-; CHECK-NEXT:    [[B:%.*]] = sext i8 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp ule i32 [[A]], [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[Y:%.*]] to i9
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i9 [[N]], [[TMP1]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %n = urem i9 %x, 254
   %a = zext i9 %n to i32
   %b = sext i8 %y to i32
   %c = icmp ule i32 %a, %b
   ret i1 %c
 }
 
 define i1 @sext_zext_slt_known_nonneg(i8 %x, i8 %y) {
 ; CHECK-LABEL: @sext_zext_slt_known_nonneg(
-; CHECK-NEXT:    [[A:%.*]] = sext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[N:%.*]] = and i8 [[Y:%.*]], 126
-; CHECK-NEXT:    [[B:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i32 [[A]], [[B]]
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i8 [[N]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %a = sext i8 %x to i32
   %n = and i8 %y, 126
   %b = zext i8 %n to i32
   %c = icmp slt i32 %a, %b
   ret i1 %c
 }
 
 define i1 @sext_zext_ult_known_nonneg(i8 %x, i8 %y) {
 ; CHECK-LABEL: @sext_zext_ult_known_nonneg(
-; CHECK-NEXT:    [[A:%.*]] = sext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[N:%.*]] = lshr i8 [[Y:%.*]], 6
-; CHECK-NEXT:    [[B:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[A]], [[B]]
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i8 [[N]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %a = sext i8 %x to i32
   %n = lshr i8 %y, 6
   %b = zext i8 %n to i32
   %c = icmp ult i32 %a, %b
   ret i1 %c
 }
 
 define i1 @sext_zext_ne_known_nonneg(i8 %x, i8 %y) {
 ; CHECK-LABEL: @sext_zext_ne_known_nonneg(
-; CHECK-NEXT:    [[A:%.*]] = sext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[N:%.*]] = udiv i8 [[Y:%.*]], 6
-; CHECK-NEXT:    [[B:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp ne i32 [[A]], [[B]]
+; CHECK-NEXT:    [[C:%.*]] = icmp ne i8 [[N]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %a = sext i8 %x to i32
   %n = udiv i8 %y, 6
   %b = zext i8 %n to i32
   %c = icmp ne i32 %a, %b
   ret i1 %c
 }
 
 define <2 x i1> @sext_zext_sge_known_nonneg_op0_narrow(<2 x i5> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @sext_zext_sge_known_nonneg_op0_narrow(
-; CHECK-NEXT:    [[A:%.*]] = sext <2 x i5> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[N:%.*]] = mul nsw <2 x i8> [[Y:%.*]], [[Y]]
-; CHECK-NEXT:    [[B:%.*]] = zext <2 x i8> [[N]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = icmp sge <2 x i32> [[A]], [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sext <2 x i5> [[X:%.*]] to <2 x i8>
+; CHECK-NEXT:    [[C:%.*]] = icmp sle <2 x i8> [[N]], [[TMP1]]
 ; CHECK-NEXT:    ret <2 x i1> [[C]]
 ;
   %a = sext <2 x i5> %x to <2 x i32>
   %n = mul nsw <2 x i8> %y, %y
   %b = zext <2 x i8> %n to <2 x i32>
   %c = icmp sge <2 x i32> %a, %b
   ret <2 x i1> %c
 }
 
 define i1 @sext_zext_uge_known_nonneg_op0_wide(i16 %x, i8 %y) {
 ; CHECK-LABEL: @sext_zext_uge_known_nonneg_op0_wide(
-; CHECK-NEXT:    [[A:%.*]] = sext i16 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[N:%.*]] = and i8 [[Y:%.*]], 12
-; CHECK-NEXT:    [[B:%.*]] = zext i8 [[N]] to i32
-; CHECK-NEXT:    [[C:%.*]] = icmp uge i32 [[A]], [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[N]] to i16
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i16 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %a = sext i16 %x to i32
   %n = and i8 %y, 12
   %b = zext i8 %n to i32
   %c = icmp uge i32 %a, %b
   ret i1 %c
 }