[llvm] [ValueTracking] Handle `trunc nuw` in `computeKnownBitsFromICmpCond` (PR #125414)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/125414

>From 4405af50027855bad6518fc96f65d84b7b66c8ff Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sun, 2 Feb 2025 22:10:10 +0800
Subject: [PATCH 1/3] [InstCombine] Add pre-commit tests. NFC.

---
 .../test/Transforms/InstCombine/known-bits.ll | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/known-bits.ll b/llvm/test/Transforms/InstCombine/known-bits.ll
index bd6b2f015145e53..653aadbf08b7591 100644
--- a/llvm/test/Transforms/InstCombine/known-bits.ll
+++ b/llvm/test/Transforms/InstCombine/known-bits.ll
@@ -480,6 +480,56 @@ if.else:
   ret i64 13
 }
 
+define i64 @test_icmp_trunc_nuw(i64 %a) {
+; CHECK-LABEL: @test_icmp_trunc_nuw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CAST:%.*]] = trunc nuw i64 [[A:%.*]] to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[CAST]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[B:%.*]] = and i64 [[A]], 2147483647
+; CHECK-NEXT:    ret i64 [[B]]
+; CHECK:       if.else:
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  %cast = trunc nuw i64 %a to i32
+  %cmp = icmp sgt i32 %cast, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  %b = and i64 %a, 2147483647
+  ret i64 %b
+
+if.else:
+  ret i64 0
+}
+
+define i64 @test_icmp_trunc_no_nuw(i64 %a) {
+; CHECK-LABEL: @test_icmp_trunc_no_nuw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CAST:%.*]] = trunc i64 [[A:%.*]] to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[CAST]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[B:%.*]] = and i64 [[A]], 2147483647
+; CHECK-NEXT:    ret i64 [[B]]
+; CHECK:       if.else:
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  %cast = trunc i64 %a to i32
+  %cmp = icmp sgt i32 %cast, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  %b = and i64 %a, 2147483647
+  ret i64 %b
+
+if.else:
+  ret i64 0
+}
+
 define i1 @test_icmp_or_distjoint(i8 %n, i1 %other) {
 ; CHECK-LABEL: @test_icmp_or_distjoint(
 ; CHECK-NEXT:  entry:

>From a02eadeca7ee45e6d4c73e61eefdc841e1180645 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sun, 2 Feb 2025 22:11:07 +0800
Subject: [PATCH 2/3] [ValueTracking] Handle `trunc nuw` in
 `computeKnownBitsFromICmpCond`

---
 llvm/lib/Analysis/ValueTracking.cpp            | 5 ++++-
 llvm/test/Transforms/InstCombine/known-bits.ll | 3 +--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index b3216f8fb78db98..c1c5cf0e614c956 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -776,7 +776,10 @@ static void computeKnownBitsFromICmpCond(const Value *V, ICmpInst *Cmp,
   if (match(LHS, m_Trunc(m_Specific(V)))) {
     KnownBits DstKnown(LHS->getType()->getScalarSizeInBits());
     computeKnownBitsFromCmp(LHS, Pred, LHS, RHS, DstKnown, SQ);
-    Known = Known.unionWith(DstKnown.anyext(Known.getBitWidth()));
+    if (match(LHS, m_NUWTrunc(m_Value())))
+      Known = Known.unionWith(DstKnown.zext(Known.getBitWidth()));
+    else
+      Known = Known.unionWith(DstKnown.anyext(Known.getBitWidth()));
     return;
   }
 
diff --git a/llvm/test/Transforms/InstCombine/known-bits.ll b/llvm/test/Transforms/InstCombine/known-bits.ll
index 653aadbf08b7591..cbd9737415f1f0e 100644
--- a/llvm/test/Transforms/InstCombine/known-bits.ll
+++ b/llvm/test/Transforms/InstCombine/known-bits.ll
@@ -487,8 +487,7 @@ define i64 @test_icmp_trunc_nuw(i64 %a) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[CAST]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    [[B:%.*]] = and i64 [[A]], 2147483647
-; CHECK-NEXT:    ret i64 [[B]]
+; CHECK-NEXT:    ret i64 [[A]]
 ; CHECK:       if.else:
 ; CHECK-NEXT:    ret i64 0
 ;

>From ae59b5977d2c29b0be3b7f8bc3cc2aca693e4506 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 3 Feb 2025 08:43:02 +0800
Subject: [PATCH 3/3] [ValueTracking] Address review comments. NFC.

---
 llvm/lib/Analysis/ValueTracking.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index c1c5cf0e614c956..6b61a3546e8b7c5 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -776,7 +776,7 @@ static void computeKnownBitsFromICmpCond(const Value *V, ICmpInst *Cmp,
   if (match(LHS, m_Trunc(m_Specific(V)))) {
     KnownBits DstKnown(LHS->getType()->getScalarSizeInBits());
     computeKnownBitsFromCmp(LHS, Pred, LHS, RHS, DstKnown, SQ);
-    if (match(LHS, m_NUWTrunc(m_Value())))
+    if (cast<TruncInst>(LHS)->hasNoUnsignedWrap())
       Known = Known.unionWith(DstKnown.zext(Known.getBitWidth()));
     else
       Known = Known.unionWith(DstKnown.anyext(Known.getBitWidth()));