[InstCombine] Generalize icmp handling in isKnownNonZero()

The dominating condition handling in isKnownNonZero() currently
only takes into account conditions of the form "x != 0" or "x == 0".
However, there are plenty of other conditions that imply non-zero,
a common one being "x s> 0".

Peculiarly, the handling for assumes was already dealing with more
general non-zero-ness conditions, so this just reuses the same
logic for the dominating condition case.

llvm/lib/Analysis/ValueTracking.cpp

@@ -590,41 +590,30 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv,
   return false;
 }
 
+static bool cmpExcludesZero(CmpInst::Predicate Pred, const Value *RHS) {
+  // v u> y implies v != 0.
+  if (Pred == ICmpInst::ICMP_UGT)
+    return true;
+
+  // Special-case v != 0 to also handle v != null.
+  if (Pred == ICmpInst::ICMP_NE)
+    return match(RHS, m_Zero());
+
+  // All other predicates - rely on generic ConstantRange handling.
+  const APInt *C;
+  if (!match(RHS, m_APInt(C)))
+    return false;
+
+  ConstantRange TrueValues = ConstantRange::makeExactICmpRegion(Pred, *C);
+  return !TrueValues.contains(APInt::getNullValue(C->getBitWidth()));
+}
+
 static bool isKnownNonZeroFromAssume(const Value *V, const Query &Q) {
   // Use of assumptions is context-sensitive. If we don't have a context, we
   // cannot use them!
   if (!Q.AC || !Q.CxtI)
     return false;
 
-  // Note that the patterns below need to be kept in sync with the code
-  // in AssumptionCache::updateAffectedValues.
-
-  auto CmpExcludesZero = [V](ICmpInst *Cmp) {
-    auto m_V = m_CombineOr(m_Specific(V), m_PtrToInt(m_Specific(V)));
-
-    Value *RHS;
-    CmpInst::Predicate Pred;
-    if (!match(Cmp, m_c_ICmp(Pred, m_V, m_Value(RHS))))
-      return false;
-    // assume(v u> y) -> assume(v != 0)
-    if (Pred == ICmpInst::ICMP_UGT)
-      return true;
-
-    // assume(v != 0)
-    // We special-case this one to ensure that we handle `assume(v != null)`.
-    if (Pred == ICmpInst::ICMP_NE)
-      return match(RHS, m_Zero());
-
-    // All other predicates - rely on generic ConstantRange handling.
-    ConstantInt *CI;
-    if (!match(RHS, m_ConstantInt(CI)))
-      return false;
-    ConstantRange RHSRange(CI->getValue());
-    ConstantRange TrueValues =
-        ConstantRange::makeAllowedICmpRegion(Pred, RHSRange);
-    return !TrueValues.contains(APInt::getNullValue(CI->getBitWidth()));
-  };
-
   if (Q.CxtI && V->getType()->isPointerTy()) {
     SmallVector<Attribute::AttrKind, 2> AttrKinds{Attribute::NonNull};
     if (!NullPointerIsDefined(Q.CxtI->getFunction(),
@@ -651,12 +640,13 @@ static bool isKnownNonZeroFromAssume(const Value *V, const Query &Q) {
     assert(I->getCalledFunction()->getIntrinsicID() == Intrinsic::assume &&
            "must be an assume intrinsic");
 
-    Value *Arg = I->getArgOperand(0);
+    Value *RHS;
-    ICmpInst *Cmp = dyn_cast<ICmpInst>(Arg);
+    CmpInst::Predicate Pred;
-    if (!Cmp)
+    auto m_V = m_CombineOr(m_Specific(V), m_PtrToInt(m_Specific(V)));
-      continue;
+    if (!match(I->getArgOperand(0), m_c_ICmp(Pred, m_V, m_Value(RHS))))
+      return false;
 
-    if (CmpExcludesZero(Cmp) && isValidAssumeForContext(I, Q.CxtI, Q.DT))
+    if (cmpExcludesZero(Pred, RHS) && isValidAssumeForContext(I, Q.CxtI, Q.DT))
       return true;
   }
 
@@ -2113,10 +2103,17 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
     }
 
     // Consider only compare instructions uniquely controlling a branch
+    Value *RHS;
     CmpInst::Predicate Pred;
-    if (!match(const_cast<User *>(U),
+    if (!match(U, m_c_ICmp(Pred, m_Specific(V), m_Value(RHS))))
-               m_c_ICmp(Pred, m_Specific(V), m_Zero())) ||
+      continue;
-        (Pred != ICmpInst::ICMP_EQ && Pred != ICmpInst::ICMP_NE))
+
+    bool NonNullIfTrue;
+    if (cmpExcludesZero(Pred, RHS))
+      NonNullIfTrue = true;
+    else if (Pred == ICmpInst::ICMP_EQ && match(RHS, m_Zero()))
+      NonNullIfTrue = false;
+    else
       continue;
 
     SmallVector<const User *, 4> WorkList;
@@ -2133,7 +2130,7 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         // propagate "pred != null" condition through AND because it is only
         // correct to assume that all conditions of AND are met in true branch.
         // TODO: Support similar logic of OR and EQ predicate?
-        if (Pred == ICmpInst::ICMP_NE)
+        if (NonNullIfTrue)
           if (auto *BO = dyn_cast<BinaryOperator>(Curr))
             if (BO->getOpcode() == Instruction::And) {
               for (auto *BOU : BO->users())
@@ -2146,11 +2143,11 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
           assert(BI->isConditional() && "uses a comparison!");
 
           BasicBlock *NonNullSuccessor =
-              BI->getSuccessor(Pred == ICmpInst::ICMP_EQ ? 1 : 0);
+              BI->getSuccessor(NonNullIfTrue ? 0 : 1);
           BasicBlockEdge Edge(BI->getParent(), NonNullSuccessor);
           if (Edge.isSingleEdge() && DT->dominates(Edge, CtxI->getParent()))
             return true;
-        } else if (Pred == ICmpInst::ICMP_NE && isGuard(Curr) &&
+        } else if (NonNullIfTrue && isGuard(Curr) &&
                    DT->dominates(cast<Instruction>(Curr), CtxI)) {
           return true;
         }

llvm/test/Transforms/Attributor/nonnull.ll

@@ -1404,24 +1404,40 @@ hd2:
 ; Original from PR43833
 declare void @sink(i32*)
 
-; FIXME: the sink argument should be marked nonnull as in @PR43833_simple.
 define void @PR43833(i32* %0, i32 %1) {
-; CHECK-LABEL: define {{[^@]+}}@PR43833
+; IS________OPM-LABEL: define {{[^@]+}}@PR43833
-; CHECK-SAME: (i32* [[TMP0:%.*]], i32 [[TMP1:%.*]]) {
+; IS________OPM-SAME: (i32* [[TMP0:%.*]], i32 [[TMP1:%.*]]) {
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP1]], 1
+; IS________OPM-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP1]], 1
-; CHECK-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP7:%.*]]
+; IS________OPM-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP7:%.*]]
-; CHECK:       4:
+; IS________OPM:       4:
-; CHECK-NEXT:    [[TMP5:%.*]] = zext i32 [[TMP1]] to i64
+; IS________OPM-NEXT:    [[TMP5:%.*]] = zext i32 [[TMP1]] to i64
-; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[TMP0]], i64 [[TMP5]]
+; IS________OPM-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[TMP0]], i64 [[TMP5]]
-; CHECK-NEXT:    br label [[TMP8:%.*]]
+; IS________OPM-NEXT:    br label [[TMP8:%.*]]
-; CHECK:       7:
+; IS________OPM:       7:
-; CHECK-NEXT:    ret void
+; IS________OPM-NEXT:    ret void
-; CHECK:       8:
+; IS________OPM:       8:
-; CHECK-NEXT:    [[TMP9:%.*]] = phi i32 [ 1, [[TMP4]] ], [ [[TMP10:%.*]], [[TMP8]] ]
+; IS________OPM-NEXT:    [[TMP9:%.*]] = phi i32 [ 1, [[TMP4]] ], [ [[TMP10:%.*]], [[TMP8]] ]
-; CHECK-NEXT:    tail call void @sink(i32* [[TMP6]])
+; IS________OPM-NEXT:    tail call void @sink(i32* [[TMP6]])
-; CHECK-NEXT:    [[TMP10]] = add nuw nsw i32 [[TMP9]], 1
+; IS________OPM-NEXT:    [[TMP10]] = add nuw nsw i32 [[TMP9]], 1
-; CHECK-NEXT:    [[TMP11:%.*]] = icmp eq i32 [[TMP10]], [[TMP1]]
+; IS________OPM-NEXT:    [[TMP11:%.*]] = icmp eq i32 [[TMP10]], [[TMP1]]
-; CHECK-NEXT:    br i1 [[TMP11]], label [[TMP7]], label [[TMP8]]
+; IS________OPM-NEXT:    br i1 [[TMP11]], label [[TMP7]], label [[TMP8]]
+;
+; IS________NPM-LABEL: define {{[^@]+}}@PR43833
+; IS________NPM-SAME: (i32* [[TMP0:%.*]], i32 [[TMP1:%.*]]) {
+; IS________NPM-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP1]], 1
+; IS________NPM-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP7:%.*]]
+; IS________NPM:       4:
+; IS________NPM-NEXT:    [[TMP5:%.*]] = zext i32 [[TMP1]] to i64
+; IS________NPM-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[TMP0]], i64 [[TMP5]]
+; IS________NPM-NEXT:    br label [[TMP8:%.*]]
+; IS________NPM:       7:
+; IS________NPM-NEXT:    ret void
+; IS________NPM:       8:
+; IS________NPM-NEXT:    [[TMP9:%.*]] = phi i32 [ 1, [[TMP4]] ], [ [[TMP10:%.*]], [[TMP8]] ]
+; IS________NPM-NEXT:    tail call void @sink(i32* nonnull [[TMP6]])
+; IS________NPM-NEXT:    [[TMP10]] = add nuw nsw i32 [[TMP9]], 1
+; IS________NPM-NEXT:    [[TMP11:%.*]] = icmp eq i32 [[TMP10]], [[TMP1]]
+; IS________NPM-NEXT:    br i1 [[TMP11]], label [[TMP7]], label [[TMP8]]
 ;
   %3 = icmp sgt i32 %1, 1
   br i1 %3, label %4, label %7

llvm/test/Transforms/InstCombine/known-non-zero.ll

@@ -140,7 +140,7 @@ define i64 @test_sgt_zero(i64 %x) {
 ; CHECK-NEXT:    [[C:%.*]] = icmp sgt i64 [[X:%.*]], 0
 ; CHECK-NEXT:    br i1 [[C]], label [[NON_ZERO:%.*]], label [[EXIT:%.*]]
 ; CHECK:       non_zero:
-; CHECK-NEXT:    [[CTZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 false), [[RNG0]]
+; CHECK-NEXT:    [[CTZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 true), [[RNG0]]
 ; CHECK-NEXT:    ret i64 [[CTZ]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i64 -1
@@ -163,7 +163,7 @@ define i64 @test_slt_neg_ten(i64 %x) {
 ; CHECK-NEXT:    [[C:%.*]] = icmp slt i64 [[X:%.*]], -10
 ; CHECK-NEXT:    br i1 [[C]], label [[NON_ZERO:%.*]], label [[EXIT:%.*]]
 ; CHECK:       non_zero:
-; CHECK-NEXT:    [[CTZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 false), [[RNG0]]
+; CHECK-NEXT:    [[CTZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 true), [[RNG0]]
 ; CHECK-NEXT:    ret i64 [[CTZ]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i64 -1
@@ -209,7 +209,7 @@ define i64 @test_ugt_unknown(i64 %x, i64 %y) {
 ; CHECK-NEXT:    [[C:%.*]] = icmp ugt i64 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    br i1 [[C]], label [[NON_ZERO:%.*]], label [[EXIT:%.*]]
 ; CHECK:       non_zero:
-; CHECK-NEXT:    [[CTZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 false), [[RNG0]]
+; CHECK-NEXT:    [[CTZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[X]], i1 true), [[RNG0]]
 ; CHECK-NEXT:    ret i64 [[CTZ]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i64 -1