[InstCombine] canonicalize sign bit checks

jacobly0 · jacobly0 · commit 32ab6eef8876 · 2025-01-15T13:04:42.000-05:00
This is a generalization of #122520 to other instructions with constrained result ranges.
diff --git a/llvm/include/llvm/IR/ConstantRange.h b/llvm/include/llvm/IR/ConstantRange.h
@@ -128,6 +128,12 @@ class [[nodiscard]] ConstantRange {
   /// NOTE: false does not mean that inverse predicate holds!
   bool icmp(CmpInst::Predicate Pred, const ConstantRange &Other) const;
 
+  /// Does the predicate \p Pred or its inverse hold between ranges this and \p
+  /// Other? Returns `true` if the predicate always holds, `false` if the
+  /// inverse always holds, or `std::nullopt` otherwise.
+  std::optional<bool> icmpOrInverse(CmpInst::Predicate Pred,
+                                    const ConstantRange &Other) const;
+
   /// Return true iff CR1 ult CR2 is equivalent to CR1 slt CR2.
   /// Does not depend on strictness/direction of the predicate.
   static bool
diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -3783,13 +3783,9 @@ static Value *simplifyICmpInst(CmpPredicate Pred, Value *LHS, Value *RHS,
   // If both operands have range metadata, use the metadata
   // to simplify the comparison.
   if (std::optional<ConstantRange> RhsCr = getRange(RHS, Q.IIQ))
-    if (std::optional<ConstantRange> LhsCr = getRange(LHS, Q.IIQ)) {
-      if (LhsCr->icmp(Pred, *RhsCr))
-        return ConstantInt::getTrue(ITy);
-
-      if (LhsCr->icmp(CmpInst::getInversePredicate(Pred), *RhsCr))
-        return ConstantInt::getFalse(ITy);
-    }
+    if (std::optional<ConstantRange> LhsCr = getRange(LHS, Q.IIQ))
+      if (auto Res = LhsCr->icmpOrInverse(Pred, *RhsCr))
+        return ConstantInt::getBool(ITy, *Res);
 
   // Compare of cast, for example (zext X) != 0 -> X != 0
   if (isa<CastInst>(LHS) && (isa<Constant>(RHS) || isa<CastInst>(RHS))) {
diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
@@ -274,6 +274,16 @@ bool ConstantRange::icmp(CmpInst::Predicate Pred,
   }
 }
 
+std::optional<bool>
+ConstantRange::icmpOrInverse(CmpInst::Predicate Pred,
+                             const ConstantRange &Other) const {
+  if (icmp(Pred, Other))
+    return true;
+  if (icmp(CmpInst::getInversePredicate(Pred), Other))
+    return false;
+  return std::nullopt;
+}
+
 /// Exact mul nuw region for single element RHS.
 static ConstantRange makeExactMulNUWRegion(const APInt &V) {
   unsigned BitWidth = V.getBitWidth();
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2674,41 +2674,10 @@ Instruction *InstCombinerImpl::foldICmpShrConstant(ICmpInst &Cmp,
 Instruction *InstCombinerImpl::foldICmpSRemConstant(ICmpInst &Cmp,
                                                     BinaryOperator *SRem,
                                                     const APInt &C) {
-  const ICmpInst::Predicate Pred = Cmp.getPredicate();
-  if (Pred == ICmpInst::ICMP_UGT || Pred == ICmpInst::ICMP_ULT) {
-    // Canonicalize unsigned predicates to signed:
-    // (X s% DivisorC) u> C -> (X s% DivisorC) s< 0
-    //   iff (C s< 0 ? ~C : C) u>= abs(DivisorC)-1
-    // (X s% DivisorC) u< C+1 -> (X s% DivisorC) s> -1
-    //   iff (C+1 s< 0 ? ~C : C) u>= abs(DivisorC)-1
-
-    const APInt *DivisorC;
-    if (!match(SRem->getOperand(1), m_APInt(DivisorC)))
-      return nullptr;
-
-    APInt NormalizedC = C;
-    if (Pred == ICmpInst::ICMP_ULT) {
-      assert(!NormalizedC.isZero() &&
-             "ult X, 0 should have been simplified already.");
-      --NormalizedC;
-    }
-    if (C.isNegative())
-      NormalizedC.flipAllBits();
-    assert(!DivisorC->isZero() &&
-           "srem X, 0 should have been simplified already.");
-    if (!NormalizedC.uge(DivisorC->abs() - 1))
-      return nullptr;
-
-    Type *Ty = SRem->getType();
-    if (Pred == ICmpInst::ICMP_UGT)
-      return new ICmpInst(ICmpInst::ICMP_SLT, SRem,
-                          ConstantInt::getNullValue(Ty));
-    return new ICmpInst(ICmpInst::ICMP_SGT, SRem,
-                        ConstantInt::getAllOnesValue(Ty));
-  }
   // Match an 'is positive' or 'is negative' comparison of remainder by a
   // constant power-of-2 value:
   // (X % pow2C) sgt/slt 0
+  const ICmpInst::Predicate Pred = Cmp.getPredicate();
   if (Pred != ICmpInst::ICMP_SGT && Pred != ICmpInst::ICMP_SLT &&
       Pred != ICmpInst::ICMP_EQ && Pred != ICmpInst::ICMP_NE)
     return nullptr;
@@ -3164,7 +3133,10 @@ Instruction *InstCombinerImpl::foldICmpAddConstant(ICmpInst &Cmp,
 
   if (ICmpInst::isUnsigned(Pred) && Add->hasNoSignedWrap() &&
       C.isNonNegative() && (C - *C2).isNonNegative() &&
-      computeConstantRange(X, /*ForSigned=*/true).add(*C2).isAllNonNegative())
+      computeConstantRange(X, /*ForSigned=*/true, /*UseInstrInfo=*/true, &AC,
+                           Add, &DT)
+          .add(*C2)
+          .isAllNonNegative())
     return new ICmpInst(ICmpInst::getSignedPredicate(Pred), X,
                         ConstantInt::get(Ty, C - *C2));
 
@@ -7056,6 +7028,66 @@ static Instruction *canonicalizeICmpBool(ICmpInst &I,
   }
 }
 
+// (icmp X, Y) --> (icmp slt/sgt X, 0/-1) iff Y is outside the signed range of X
+static ICmpInst *canonicalizeSignBitCheck(ICmpInst::Predicate Pred, Value *X,
+                                          const ConstantRange &XRange,
+                                          const ConstantRange &YRange) {
+  if (XRange.isSignWrappedSet())
+    return nullptr;
+  unsigned BitWidth = XRange.getBitWidth();
+  APInt SMin = APInt::getSignedMinValue(BitWidth);
+  APInt Zero = APInt::getZero(BitWidth);
+  auto NegResult =
+      XRange.intersectWith(ConstantRange(SMin, Zero), ConstantRange::Signed)
+          .icmpOrInverse(Pred, YRange);
+  if (!NegResult)
+    return nullptr;
+  auto PosResult =
+      XRange.intersectWith(ConstantRange(Zero, SMin), ConstantRange::Signed)
+          .icmpOrInverse(Pred, YRange);
+  if (!PosResult)
+    return nullptr;
+  assert(NegResult != PosResult &&
+         "Known result should been simplified already.");
+  Type *Ty = X->getType();
+  if (*NegResult)
+    return new ICmpInst(ICmpInst::ICMP_SLT, X, ConstantInt::getNullValue(Ty));
+  return new ICmpInst(ICmpInst::ICMP_SGT, X, ConstantInt::getAllOnesValue(Ty));
+}
+
+// Try to fold an icmp using the constant ranges of its operands.
+Instruction *InstCombinerImpl::foldICmpUsingConstantRanges(ICmpInst &Cmp) {
+  Value *X = Cmp.getOperand(0);
+  if (!X->getType()->isIntOrIntVectorTy())
+    return nullptr;
+  Value *Y = Cmp.getOperand(1);
+  ICmpInst::Predicate Pred = Cmp.getPredicate();
+  ConstantRange XRange = computeConstantRange(
+      X, ICmpInst::isSigned(Pred), /*UseInstrInfo=*/true, &AC, &Cmp, &DT);
+  if (XRange.isFullSet())
+    return nullptr; // early out if we don't have any information
+  ConstantRange YRange = computeConstantRange(
+      Y, ICmpInst::isSigned(Pred), /*UseInstrInfo=*/true, &AC, &Cmp, &DT);
+  if (YRange.isFullSet())
+    return nullptr; // early out if we don't have any information
+  if (auto Res = XRange.icmpOrInverse(Pred, YRange))
+    return replaceInstUsesWith(Cmp, ConstantInt::getBool(Cmp.getType(), *Res));
+  if (ICmpInst::isUnsigned(Pred)) {
+    // Check if this icmp is actually a sign bit check.
+    const APInt *C;
+    bool IgnoreTrueIfSigned;
+    if (!match(Y, m_APInt(C)) ||
+        !isSignBitCheck(Pred, *C, IgnoreTrueIfSigned)) {
+      if (ICmpInst *Res = canonicalizeSignBitCheck(Pred, X, XRange, YRange))
+        return Res;
+      if (ICmpInst *Res = canonicalizeSignBitCheck(
+              ICmpInst::getSwappedPredicate(Pred), Y, YRange, XRange))
+        return Res;
+    }
+  }
+  return nullptr;
+}
+
 // Transform pattern like:
 //   (1 << Y) u<= X  or  ~(-1 << Y) u<  X  or  ((1 << Y)+(-1)) u<  X
 //   (1 << Y) u>  X  or  ~(-1 << Y) u>= X  or  ((1 << Y)+(-1)) u>= X
@@ -7428,6 +7460,9 @@ Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
   if (Instruction *Res = canonicalizeICmpPredicate(I))
     return Res;
 
+  if (Instruction *Res = foldICmpUsingConstantRanges(I))
+    return Res;
+
   if (Instruction *Res = foldICmpWithConstant(I))
     return Res;
 
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -668,6 +668,7 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
   Instruction *foldICmpWithCastOp(ICmpInst &ICmp);
   Instruction *foldICmpWithZextOrSext(ICmpInst &ICmp);
 
+  Instruction *foldICmpUsingConstantRanges(ICmpInst &Cmp);
   Instruction *foldICmpUsingKnownBits(ICmpInst &Cmp);
   Instruction *foldICmpWithDominatingICmp(ICmpInst &Cmp);
   Instruction *foldICmpWithConstant(ICmpInst &Cmp);
diff --git a/llvm/test/Transforms/InstCombine/icmp-dom.ll b/llvm/test/Transforms/InstCombine/icmp-dom.ll
@@ -381,11 +381,11 @@ falselabel:
 
 define i8 @PR48900_alt(i8 %i, ptr %p) {
 ; CHECK-LABEL: @PR48900_alt(
-; CHECK-NEXT:    [[SMAX:%.*]] = call i8 @llvm.smax.i8(i8 [[I:%.*]], i8 -127)
-; CHECK-NEXT:    [[I4:%.*]] = icmp ugt i8 [[SMAX]], -128
+; CHECK-NEXT:    [[I4:%.*]] = icmp slt i8 [[I:%.*]], 0
 ; CHECK-NEXT:    br i1 [[I4]], label [[TRUELABEL:%.*]], label [[FALSELABEL:%.*]]
 ; CHECK:       truelabel:
-; CHECK-NEXT:    [[UMIN:%.*]] = call i8 @llvm.smin.i8(i8 [[SMAX]], i8 -126)
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt i8 [[I]], -126
+; CHECK-NEXT:    [[UMIN:%.*]] = select i1 [[TMP1]], i8 -127, i8 -126
 ; CHECK-NEXT:    ret i8 [[UMIN]]
 ; CHECK:       falselabel:
 ; CHECK-NEXT:    ret i8 0
diff --git a/llvm/test/Transforms/InstCombine/icmp-srem.ll b/llvm/test/Transforms/InstCombine/icmp-srem.ll
@@ -398,9 +398,7 @@ define i1 @icmp_slt_srem_pos_range(i32 %x, i32 range(i32 999, -2147483648) %y) {
 define i1 @icmp_sle_srem_pos_range(i32 %x, i32 range(i32 999, -2147483648) %y) {
 ; CHECK-LABEL: define i1 @icmp_sle_srem_pos_range(
 ; CHECK-SAME: i32 [[X:%.*]], i32 range(i32 999, -2147483648) [[Y:%.*]]) {
-; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X]], 1000
-; CHECK-NEXT:    [[C:%.*]] = icmp sle i32 [[R]], [[Y]]
-; CHECK-NEXT:    ret i1 [[C]]
+; CHECK-NEXT:    ret i1 true
 ;
   %r = srem i32 %x, 1000
   %c = icmp sle i32 %r, %y
@@ -410,9 +408,7 @@ define i1 @icmp_sle_srem_pos_range(i32 %x, i32 range(i32 999, -2147483648) %y) {
 define i1 @icmp_sgt_srem_pos_range(i32 %x, i32 range(i32 999, -2147483648) %y) {
 ; CHECK-LABEL: define i1 @icmp_sgt_srem_pos_range(
 ; CHECK-SAME: i32 [[X:%.*]], i32 range(i32 999, -2147483648) [[Y:%.*]]) {
-; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X]], 1000
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i32 [[R]], [[Y]]
-; CHECK-NEXT:    ret i1 [[C]]
+; CHECK-NEXT:    ret i1 false
 ;
   %r = srem i32 %x, 1000
   %c = icmp sgt i32 %r, %y
@@ -434,9 +430,7 @@ define i1 @icmp_sge_srem_pos_range(i32 %x, i32 range(i32 999, -2147483648) %y) {
 define i1 @icmp_slt_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y) {
 ; CHECK-LABEL: define i1 @icmp_slt_srem_neg_range(
 ; CHECK-SAME: i32 [[X:%.*]], i32 range(i32 -2147483648, -999) [[Y:%.*]]) {
-; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X]], 1000
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i32 [[R]], [[Y]]
-; CHECK-NEXT:    ret i1 [[C]]
+; CHECK-NEXT:    ret i1 false
 ;
   %r = srem i32 %x, 1000
   %c = icmp slt i32 %r, %y
@@ -446,9 +440,7 @@ define i1 @icmp_slt_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y)
 define i1 @icmp_sle_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y) {
 ; CHECK-LABEL: define i1 @icmp_sle_srem_neg_range(
 ; CHECK-SAME: i32 [[X:%.*]], i32 range(i32 -2147483648, -999) [[Y:%.*]]) {
-; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X]], 1000
-; CHECK-NEXT:    [[C:%.*]] = icmp sle i32 [[R]], [[Y]]
-; CHECK-NEXT:    ret i1 [[C]]
+; CHECK-NEXT:    ret i1 false
 ;
   %r = srem i32 %x, 1000
   %c = icmp sle i32 %r, %y
@@ -458,9 +450,7 @@ define i1 @icmp_sle_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y)
 define i1 @icmp_sgt_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y) {
 ; CHECK-LABEL: define i1 @icmp_sgt_srem_neg_range(
 ; CHECK-SAME: i32 [[X:%.*]], i32 range(i32 -2147483648, -999) [[Y:%.*]]) {
-; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X]], 1000
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i32 [[R]], [[Y]]
-; CHECK-NEXT:    ret i1 [[C]]
+; CHECK-NEXT:    ret i1 true
 ;
   %r = srem i32 %x, 1000
   %c = icmp sgt i32 %r, %y
@@ -470,9 +460,7 @@ define i1 @icmp_sgt_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y)
 define i1 @icmp_sge_srem_neg_range(i32 %x, i32 range(i32 -2147483648, -999) %y) {
 ; CHECK-LABEL: define i1 @icmp_sge_srem_neg_range(
 ; CHECK-SAME: i32 [[X:%.*]], i32 range(i32 -2147483648, -999) [[Y:%.*]]) {
-; CHECK-NEXT:    [[R:%.*]] = srem i32 [[X]], 1000
-; CHECK-NEXT:    [[C:%.*]] = icmp sge i32 [[R]], [[Y]]
-; CHECK-NEXT:    ret i1 [[C]]
+; CHECK-NEXT:    ret i1 true
 ;
   %r = srem i32 %x, 1000
   %c = icmp sge i32 %r, %y
diff --git a/llvm/test/Transforms/InstCombine/smin-icmp.ll b/llvm/test/Transforms/InstCombine/smin-icmp.ll
@@ -965,9 +965,9 @@ define void @eq_smin_v2i32_constant(<2 x i32> %y) {
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP4]])
 ; CHECK-NEXT:    [[CMP5:%.*]] = icmp ult <2 x i32> [[COND]], splat (i32 10)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP5]])
-; CHECK-NEXT:    [[CMP6:%.*]] = icmp ult <2 x i32> [[COND]], splat (i32 11)
+; CHECK-NEXT:    [[CMP6:%.*]] = icmp sgt <2 x i32> [[Y]], splat (i32 -1)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP6]])
-; CHECK-NEXT:    [[CMP7:%.*]] = icmp ugt <2 x i32> [[COND]], splat (i32 10)
+; CHECK-NEXT:    [[CMP7:%.*]] = icmp slt <2 x i32> [[Y]], zeroinitializer
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP7]])
 ; CHECK-NEXT:    [[CMP8:%.*]] = icmp ugt <2 x i32> [[COND]], splat (i32 9)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP8]])
@@ -1004,18 +1004,17 @@ define void @eq_smin_v2i32_constant(<2 x i32> %y) {
 ; icmp pred smin(C1, Y), C2 where C1 < C2
 define void @slt_smin_v2i32_constant(<2 x i32> %y) {
 ; CHECK-LABEL: @slt_smin_v2i32_constant(
-; CHECK-NEXT:    [[COND:%.*]] = call <2 x i32> @llvm.smin.v2i32(<2 x i32> [[Y:%.*]], <2 x i32> splat (i32 5))
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> splat (i1 true))
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> splat (i1 true))
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> zeroinitializer)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> zeroinitializer)
-; CHECK-NEXT:    [[CMP5:%.*]] = icmp ult <2 x i32> [[COND]], splat (i32 10)
+; CHECK-NEXT:    [[CMP5:%.*]] = icmp sgt <2 x i32> [[Y:%.*]], splat (i32 -1)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP5]])
-; CHECK-NEXT:    [[CMP6:%.*]] = icmp ult <2 x i32> [[COND]], splat (i32 11)
+; CHECK-NEXT:    [[CMP6:%.*]] = icmp sgt <2 x i32> [[Y]], splat (i32 -1)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP6]])
-; CHECK-NEXT:    [[CMP7:%.*]] = icmp ugt <2 x i32> [[COND]], splat (i32 10)
+; CHECK-NEXT:    [[CMP7:%.*]] = icmp slt <2 x i32> [[Y]], zeroinitializer
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP7]])
-; CHECK-NEXT:    [[CMP8:%.*]] = icmp ugt <2 x i32> [[COND]], splat (i32 9)
+; CHECK-NEXT:    [[CMP8:%.*]] = icmp slt <2 x i32> [[Y]], zeroinitializer
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> [[CMP8]])
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> zeroinitializer)
 ; CHECK-NEXT:    call void @use_v2i1(<2 x i1> splat (i1 true))
diff --git a/llvm/unittests/IR/ConstantRangeTest.cpp b/llvm/unittests/IR/ConstantRangeTest.cpp
@@ -1694,12 +1694,23 @@ void ICmpTestImpl(CmpInst::Predicate Pred) {
   EnumerateTwoInterestingConstantRanges(
       [&](const ConstantRange &CR1, const ConstantRange &CR2) {
         bool Exhaustive = true;
+        bool ExhaustiveInverse = true;
         ForeachNumInConstantRange(CR1, [&](const APInt &N1) {
           ForeachNumInConstantRange(CR2, [&](const APInt &N2) {
-            Exhaustive &= ICmpInst::compare(N1, N2, Pred);
+            bool Res = ICmpInst::compare(N1, N2, Pred);
+            Exhaustive &= Res;
+            ExhaustiveInverse &= !Res;
           });
         });
+
+        std::optional<bool> ExhaustiveOrInverse;
+        if (Exhaustive) // Expect true if Exhaustive && ExhaustiveInverse.
+          ExhaustiveOrInverse = true;
+        else if (ExhaustiveInverse)
+          ExhaustiveOrInverse = false;
+
         EXPECT_EQ(CR1.icmp(Pred, CR2), Exhaustive);
+        EXPECT_EQ(CR1.icmpOrInverse(Pred, CR2), ExhaustiveOrInverse);
       });
 }
 
