[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 383840b8110a2fc014d363fd6038995e901442b4 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Mon, 2 Dec 2024 23:00:12 +0100
Subject: [PATCH 1/2] [ValueTracking] Handle assume(trunc x to i1) in
 knownNonZero

---
 llvm/lib/Analysis/ValueTracking.cpp                        | 7 ++++++-
 llvm/test/Transforms/Attributor/lvi-after-jumpthreading.ll | 5 +----
 llvm/test/Transforms/InstSimplify/assume-non-zero.ll       | 3 +--
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 3b0249f91d6d7..8c2561595db9b 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -643,11 +643,16 @@ static bool isKnownNonZeroFromAssume(const Value *V, const SimplifyQuery &Q) {
     // Warning: This loop can end up being somewhat performance sensitive.
     // We're running this loop for once for each value queried resulting in a
     // runtime of ~O(#assumes * #values).
+    Value *Arg = I->getArgOperand(0);
+    if (match(Arg, m_TruncOrSelf(m_Specific(V))) &&
+        isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+      return true;
+    }
 
     Value *RHS;
     CmpPredicate Pred;
     auto m_V = m_CombineOr(m_Specific(V), m_PtrToInt(m_Specific(V)));
-    if (!match(I->getArgOperand(0), m_c_ICmp(Pred, m_V, m_Value(RHS))))
+    if (!match(Arg, m_c_ICmp(Pred, m_V, m_Value(RHS))))
       continue;
 
     if (cmpExcludesZero(Pred, RHS) && isValidAssumeForContext(I, Q.CxtI, Q.DT))
diff --git a/llvm/test/Transforms/Attributor/lvi-after-jumpthreading.ll b/llvm/test/Transforms/Attributor/lvi-after-jumpthreading.ll
index 2aa95216a6656..fd9d2e4de61ba 100644
--- a/llvm/test/Transforms/Attributor/lvi-after-jumpthreading.ll
+++ b/llvm/test/Transforms/Attributor/lvi-after-jumpthreading.ll
@@ -130,8 +130,6 @@ exit2:
   ret i32 30
 }
 
-; FIXME: We should be able to merge cont into do.
-; FIXME: COND should be replaced with false. This will be fixed by improving LVI.
 define i32 @test4(i32 %i, i1 %f, i32 %n) {
 ; CHECK-LABEL: define {{[^@]+}}@test4
 ; CHECK-SAME: (i32 [[I:%.*]], i1 [[F:%.*]], i32 [[N:%.*]]) {
@@ -148,8 +146,7 @@ define i32 @test4(i32 %i, i1 %f, i32 %n) {
 ; CHECK-NEXT:    call void @dummy(i1 [[F]]) #[[ATTR2]]
 ; CHECK-NEXT:    [[CONSUME:%.*]] = call i32 @exit()
 ; CHECK-NEXT:    call void @llvm.assume(i1 noundef [[F]])
-; CHECK-NEXT:    [[COND:%.*]] = icmp eq i1 [[F]], false
-; CHECK-NEXT:    br i1 [[COND]], label [[EXIT]], label [[CONT:%.*]]
+; CHECK-NEXT:    br label [[CONT:%.*]]
 ; CHECK:       exit2:
 ; CHECK-NEXT:    ret i32 30
 ;
diff --git a/llvm/test/Transforms/InstSimplify/assume-non-zero.ll b/llvm/test/Transforms/InstSimplify/assume-non-zero.ll
index 22bdf374487da..331c9504f2da6 100644
--- a/llvm/test/Transforms/InstSimplify/assume-non-zero.ll
+++ b/llvm/test/Transforms/InstSimplify/assume-non-zero.ll
@@ -236,8 +236,7 @@ define i1 @nonnull_trunc_true(i8 %x) {
 ; CHECK-LABEL: @nonnull_trunc_true(
 ; CHECK-NEXT:    [[A:%.*]] = trunc i8 [[X:%.*]] to i1
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[A]])
-; CHECK-NEXT:    [[Q:%.*]] = icmp ne i8 [[X]], 0
-; CHECK-NEXT:    ret i1 [[Q]]
+; CHECK-NEXT:    ret i1 true
 ;
   %a = trunc i8 %x to i1
   call void @llvm.assume(i1 %a)

>From 570f3c613e1db493070c4d7c62b52621fc0d505e 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 2/2] [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 8c2561595db9b..2606eee35c94d 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -912,6 +912,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)