[llvm] [InstCombine] isKnownNonNegative should recognize `b - a` after `a <= b` (PR #145105)

[Ross Kirsling via llvm-commits]

https://github.com/rkirsling updated https://github.com/llvm/llvm-project/pull/145105

>From e6612fe4290bd154e981569b92565f551ce21a21 Mon Sep 17 00:00:00 2001
From: Ross Kirsling <ross.kirsling at sony.com>
Date: Fri, 20 Jun 2025 11:46:11 -0700
Subject: [PATCH 1/2] Add test case.

---
 .../sub-after-sle-is-non-negative.ll          | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll

diff --git a/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll b/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll
new file mode 100644
index 0000000000000..b4427b72ff8a5
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+define void @test_as_arg(i8 %a, i8 %b) {
+; CHECK-LABEL: define void @test_as_arg(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt 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 sgt 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
+}
+
+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:%.*]]) {
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt 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 sgt 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
+}

>From d80031d1c511b91b86f9333a6a32c24c52b7ad32 Mon Sep 17 00:00:00 2001
From: Ross Kirsling <ross.kirsling at sony.com>
Date: Fri, 20 Jun 2025 11:46:58 -0700
Subject: [PATCH 2/2] Address issue.

---
 llvm/lib/Analysis/ValueTracking.cpp                   | 11 ++++++++++-
 .../InstCombine/sub-after-sle-is-non-negative.ll      |  4 ++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 73320b556f825..be9843503298f 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -278,7 +278,16 @@ static bool isKnownNonZero(const Value *V, const APInt &DemandedElts,
 
 bool llvm::isKnownNonNegative(const Value *V, const SimplifyQuery &SQ,
                               unsigned Depth) {
-  return computeKnownBits(V, SQ, Depth).isNonNegative();
+  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;
 }
 
 bool llvm::isKnownPositive(const Value *V, const SimplifyQuery &SQ,
diff --git a/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll b/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll
index b4427b72ff8a5..279a6de4125ed 100644
--- a/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll
+++ b/llvm/test/Transforms/InstCombine/sub-after-sle-is-non-negative.ll
@@ -7,7 +7,7 @@ define void @test_as_arg(i8 %a, i8 %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:    [[CONV:%.*]] = zext nneg i8 [[SUB]] to i16
 ; CHECK-NEXT:    call void @subroutine(i16 [[CONV]])
 ; CHECK-NEXT:    br label %[[COND_END]]
 ; CHECK:       [[COND_END]]:
@@ -37,7 +37,7 @@ define i16 @test_as_retval(i8 %a, i8 %b) {
 ; 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:    [[CONV:%.*]] = zext nneg i8 [[SUB]] to i16
 ; CHECK-NEXT:    ret i16 [[CONV]]
 ;
   %cmp = icmp sgt i8 %a, %b