[ValueTracking] Handle assume( trunc x to i1) in ComputeKnownBits

Author: andjo403

llvm/lib/Analysis/ValueTracking.cpp

@@ -814,8 +814,6 @@ void llvm::computeKnownBitsFromContext(const Value *V, KnownBits &Known,
   if (!Q.AC)
     return;
 
-  unsigned BitWidth = Known.getBitWidth();
-
   // Note that the patterns below need to be kept in sync with the code
   // in AssumptionCache::updateAffectedValues.
 
@@ -849,17 +847,14 @@ void llvm::computeKnownBitsFromContext(const Value *V, KnownBits &Known,
 
     Value *Arg = I->getArgOperand(0);
 
-    if (Arg == V && isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
-      assert(BitWidth == 1 && "assume operand is not i1?");
-      (void)BitWidth;
-      Known.setAllOnes();
+    if (match(Arg, m_TruncOrSelf(m_Specific(V))) &&
+        isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+      Known.One.setBit(0);
       return;
     }
-    if (match(Arg, m_Not(m_Specific(V))) &&
+    if (match(Arg, m_Not(m_TruncOrSelf(m_Specific(V)))) &&
         isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
-      assert(BitWidth == 1 && "assume operand is not i1?");
-      (void)BitWidth;
-      Known.setAllZero();
+      Known.Zero.setBit(0);
       return;
     }
 

llvm/test/Transforms/InstSimplify/compare.ll

@@ -3429,8 +3429,7 @@ define i1 @icmp_eq_false_by_trunc(i8 %x) {
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[X:%.*]] to i1
 ; CHECK-NEXT:    [[NOT:%.*]] = xor i1 [[TRUNC]], true
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[NOT]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[X]], 1
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %trunc = trunc i8 %x to i1
   %not = xor i1 %trunc, true

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)