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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

<details>
<summary>Changes</summary>

This patch extends https://github.com/llvm/llvm-project/pull/82803 to further infer high bits when `nuw` is set. It will save some and instructions on induction variables.
Alive2: https://alive2.llvm.org/ce/z/j-YFvt


---
Full diff: https://github.com/llvm/llvm-project/pull/125414.diff


2 Files Affected:

- (modified) llvm/lib/Analysis/ValueTracking.cpp (+4-1) 
- (modified) llvm/test/Transforms/InstCombine/known-bits.ll (+49) 


``````````diff
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 bd6b2f015145e53..cbd9737415f1f0e 100644
--- a/llvm/test/Transforms/InstCombine/known-bits.ll
+++ b/llvm/test/Transforms/InstCombine/known-bits.ll
@@ -480,6 +480,55 @@ 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:    ret i64 [[A]]
+; 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:

``````````

</details>


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