[InstCombine] KnownBits::isNonNegative should recognize b - a after a <= b

llvm/lib/Analysis/ValueTracking.cpp

@@ -278,16 +278,7 @@ static bool isKnownNonZero(const Value *V, const APInt &DemandedElts,
 
 bool llvm::isKnownNonNegative(const Value *V, const SimplifyQuery &SQ,
                               unsigned Depth) {
-  if (computeKnownBits(V, SQ, Depth).isNonNegative())
-    return true;
-
-  Value *X, *Y;
-  if (match(V, m_NSWSub(m_Value(X), m_Value(Y))))
-    if (std::optional<bool> result =
-            isImpliedByDomCondition(ICmpInst::ICMP_SLE, Y, X, SQ.CxtI, SQ.DL))
-      return *result;
-
-  return false;
+  return computeKnownBits(V, SQ, Depth).isNonNegative();
 }
 
 bool llvm::isKnownPositive(const Value *V, const SimplifyQuery &SQ,
@@ -371,6 +362,12 @@ static void computeKnownBitsAddSub(bool Add, const Value *Op0, const Value *Op1,
 
   computeKnownBits(Op0, DemandedElts, Known2, Q, Depth + 1);
   KnownOut = KnownBits::computeForAddSub(Add, NSW, NUW, Known2, KnownOut);
+
+  if (!Add && NSW)
+    if (std::optional<bool> result =
+            isImpliedByDomCondition(ICmpInst::ICMP_SLE, Op1, Op0, Q.CxtI, Q.DL);
+        *result)
+      KnownOut.makeNonNegative();
 }
 
 static void computeKnownBitsMul(const Value *Op0, const Value *Op1, bool NSW,

llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll

@@ -1,5 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+declare void @subroutine(i16)
+
 define void @test_as_arg(i8 %a, i8 %b) {
 ; CHECK-LABEL: define void @test_as_arg(
 ; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
@@ -26,8 +28,6 @@ cond.end:
   ret void
 }
 
-declare void @subroutine(i16)
-
 define i16 @test_as_retval(i8 %a, i8 %b) {
 ; CHECK-LABEL: define i16 @test_as_retval(
 ; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
@@ -51,3 +51,53 @@ cond.false:
   %conv = sext i8 %sub to i16
   ret i16 %conv
 }
+
+define void @test_as_arg_neg(i8 %a, i8 %b) {
+; CHECK-LABEL: define void @test_as_arg_neg(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i8 [[A]], [[B]]
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COND_END:.*]], label %[[COND_FALSE:.*]]
+; CHECK:       [[COND_FALSE]]:
+; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i8 [[B]], [[A]]
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[SUB]] to i16
+; CHECK-NEXT:    call void @subroutine(i16 [[CONV]])
+; CHECK-NEXT:    br label %[[COND_END]]
+; CHECK:       [[COND_END]]:
+; CHECK-NEXT:    ret void
+;
+  %cmp = icmp slt i8 %a, %b
+  br i1 %cmp, label %cond.end, label %cond.false
+
+cond.false:
+  %sub = sub nsw i8 %b, %a
+  %conv = sext i8 %sub to i16
+  call void @subroutine(i16 %conv)
+  br label %cond.end
+
+cond.end:
+  ret void
+}
+
+define i16 @test_as_retval_neg(i8 %a, i8 %b) {
+; CHECK-LABEL: define i16 @test_as_retval_neg(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i8 [[A]], [[B]]
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COND_TRUE:.*]], label %[[COND_FALSE:.*]]
+; CHECK:       [[COND_TRUE]]:
+; CHECK-NEXT:    ret i16 0
+; CHECK:       [[COND_FALSE]]:
+; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i8 [[B]], [[A]]
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[SUB]] to i16
+; CHECK-NEXT:    ret i16 [[CONV]]
+;
+  %cmp = icmp slt i8 %a, %b
+  br i1 %cmp, label %cond.true, label %cond.false
+
+cond.true:
+  ret i16 0
+
+cond.false:
+  %sub = sub nsw i8 %b, %a
+  %conv = sext i8 %sub to i16
+  ret i16 %conv
+}