[llvm] [ValueTracking] Handle assume( trunc x to i1) (PR #118406)

[Andreas Jonson via llvm-commits]

https://github.com/andjo403 updated https://github.com/llvm/llvm-project/pull/118406

>From 2b19e3d5dda3dd7b81ea3f73f706b0c5c6a2cde0 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Tue, 4 Feb 2025 23:59:38 +0100
Subject: [PATCH] [ValueTracking] Handle assume( trunc x to i1) in
 ComputeKnownBits

---
 llvm/lib/Analysis/ValueTracking.cpp          | 10 ++++++++++
 llvm/test/Transforms/InstCombine/assume.ll   |  3 +--
 llvm/test/Transforms/InstSimplify/shr-nop.ll |  3 +--
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index d9c55330f8664..ff8f3b6ae3bc8 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -910,6 +910,16 @@ void llvm::computeKnownBitsFromContext(const Value *V, KnownBits &Known,
       Known.setAllZero();
       return;
     }
+    auto *Trunc = dyn_cast<TruncInst>(Arg);
+    if (Trunc && Trunc->getOperand(0) == V &&
+        isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+      if (Trunc->hasNoUnsignedWrap()) {
+        Known = KnownBits::makeConstant(APInt(BitWidth, 1));
+        return;
+      }
+      Known.One.setBit(0);
+      return;
+    }
 
     // The remaining tests are all recursive, so bail out if we hit the limit.
     if (Depth == MaxAnalysisRecursionDepth)
diff --git a/llvm/test/Transforms/InstCombine/assume.ll b/llvm/test/Transforms/InstCombine/assume.ll
index 1f8fa1716b2aa..e87a61a57ea47 100644
--- a/llvm/test/Transforms/InstCombine/assume.ll
+++ b/llvm/test/Transforms/InstCombine/assume.ll
@@ -1013,8 +1013,7 @@ define i1 @assume_trunc_nuw_eq_one(i8 %x) {
 ; CHECK-LABEL: @assume_trunc_nuw_eq_one(
 ; CHECK-NEXT:    [[A:%.*]] = trunc nuw i8 [[X:%.*]] to i1
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[A]])
-; CHECK-NEXT:    [[Q:%.*]] = icmp eq i8 [[X]], 1
-; CHECK-NEXT:    ret i1 [[Q]]
+; CHECK-NEXT:    ret i1 true
 ;
   %a = trunc nuw i8 %x to i1
   call void @llvm.assume(i1 %a)
diff --git a/llvm/test/Transforms/InstSimplify/shr-nop.ll b/llvm/test/Transforms/InstSimplify/shr-nop.ll
index 4b9292aab4bb3..b18a1c6d30e74 100644
--- a/llvm/test/Transforms/InstSimplify/shr-nop.ll
+++ b/llvm/test/Transforms/InstSimplify/shr-nop.ll
@@ -385,8 +385,7 @@ define i8 @exact_lshr_lowbit_set_assume_trunc(i8 %x) {
 ; CHECK-LABEL: @exact_lshr_lowbit_set_assume_trunc(
 ; CHECK-NEXT:    [[COND:%.*]] = trunc i8 [[X:%.*]] to i1
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[SHR:%.*]] = lshr exact i8 [[X]], 1
-; CHECK-NEXT:    ret i8 [[SHR]]
+; CHECK-NEXT:    ret i8 [[X]]
 ;
   %cond = trunc i8 %x to i1
   call void @llvm.assume(i1 %cond)