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

Author: andjo403

llvm/lib/Analysis/ValueTracking.cpp

@@ -873,10 +873,9 @@ 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))) &&

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)