Fix llvm::ComputeNumSignBits with some operations and llvm.assume

[rampitec]

Summary

Currently ComputeNumSignBits does early exit while processing some
of the operations (add, sub, mul, and select). This prevents the
function from using AssumptionCacheTracker if passed.

[artem.tamazov]

Good job, thanks!

llvm/trunk/lib/Analysis/ValueTracking.cpp

if (Tmp != 1) {
       // We computed what we know about the sign bits as our first
       // answer. Now proceed to the generic code that uses
       // computeKnownBits, and pick whichever answer is better.
     }
     break;
 
   case Instruction::Select:
     Tmp = ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
-    if (Tmp == 1) return 1;  // Early out.
+    if (Tmp == 1) break;
     Tmp2 = ComputeNumSignBits(U->getOperand(2), Depth + 1, Q);
     return std::min(Tmp, Tmp2);
 
   case Instruction::Add:
     // Add can have at most one carry bit.  Thus we know that the output
     // is, at worst, one more bit than the inputs.
     Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
-    if (Tmp == 1) return 1;  // Early out.
+    if (Tmp == 1) break;
 
     // Special case decrementing a value (ADD X, -1):
     if (const auto *CRHS = dyn_cast<Constant>(U->getOperand(1)))
       if (CRHS->isAllOnesValue()) {
         KnownBits Known(TyBits);
         computeKnownBits(U->getOperand(0), Known, Depth + 1, Q);
 
         // If the input is known to be 0 or 1, the output is 0/-1, which is all
         // sign bits set.
         if ((Known.Zero | 1).isAllOnesValue())
           return TyBits;
 
         // If we are subtracting one from a positive number, there is no carry
         // out of the result.
         if (Known.isNonNegative())
           return Tmp;
       }
 
     Tmp2 = ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
-    if (Tmp2 == 1) return 1;
+    if (Tmp2 == 1) break;
     return std::min(Tmp, Tmp2)-1;
 
   case Instruction::Sub:
     Tmp2 = ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
-    if (Tmp2 == 1) return 1;
+    if (Tmp2 == 1) break;
 
     // Handle NEG.
     if (const auto *CLHS = dyn_cast<Constant>(U->getOperand(0)))
       if (CLHS->isNullValue()) {
         KnownBits Known(TyBits);
         computeKnownBits(U->getOperand(1), Known, Depth + 1, Q);
         // If the input is known to be 0 or 1, the output is 0/-1, which is all
         // sign bits set.
         if ((Known.Zero | 1).isAllOnesValue())
           return TyBits;
 
         // If the input is known to be positive (the sign bit is known clear),
         // the output of the NEG has the same number of sign bits as the input.
         if (Known.isNonNegative())
           return Tmp2;
 
         // Otherwise, we treat this like a SUB.
       }
 
     // Sub can have at most one carry bit.  Thus we know that the output
     // is, at worst, one more bit than the inputs.
     Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
-    if (Tmp == 1) return 1;  // Early out.
+    if (Tmp == 1) break;
     return std::min(Tmp, Tmp2)-1;
 
   case Instruction::Mul: {
     // The output of the Mul can be at most twice the valid bits in the inputs.
     unsigned SignBitsOp0 = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
-    if (SignBitsOp0 == 1) return 1;  // Early out.
+    if (SignBitsOp0 == 1) break;
     unsigned SignBitsOp1 = ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
-    if (SignBitsOp1 == 1) return 1;
+    if (SignBitsOp1 == 1) break;
     unsigned OutValidBits =
         (TyBits - SignBitsOp0 + 1) + (TyBits - SignBitsOp1 + 1);
     return OutValidBits > TyBits ? 1 : TyBits - OutValidBits + 1;
   }
 
   case Instruction::PHI: {
     const PHINode *PN = cast<PHINode>(U);
     unsigned NumIncomingValues = PN->getNumIncomingValues();

llvm/trunk/test/Analysis/ValueTracking/numsignbits-from-assume.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define i32 @computeNumSignBits_add1(i32 %in) {
+; CHECK-LABEL: @computeNumSignBits_add1(
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[IN:%.*]], 1
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[ADD]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[ADD]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %add = add i32 %in, 1
+  %cond = icmp ule i32 %add, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %add, 3
+  ret i32 %sh
+}
+
+define i32 @computeNumSignBits_add2(i32 %in1, i32 %in2) {
+; CHECK-LABEL: @computeNumSignBits_add2(
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[IN1:%.*]], [[IN2:%.*]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[ADD]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[ADD]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %add = add i32 %in1, %in2
+  %cond = icmp ule i32 %add, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %add, 3
+  ret i32 %sh
+}
+
+define i32 @computeNumSignBits_sub1(i32 %in) {
+; CHECK-LABEL: @computeNumSignBits_sub1(
+; CHECK-NEXT:    [[SUB:%.*]] = sub i32 1, [[IN:%.*]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[SUB]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[SUB]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %sub = sub i32 1, %in
+  %cond = icmp ule i32 %sub, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %sub, 3
+  ret i32 %sh
+}
+
+define i32 @computeNumSignBits_sub2(i32 %in) {
+; CHECK-LABEL: @computeNumSignBits_sub2(
+; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[IN:%.*]], -1
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[SUB]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[SUB]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %sub = sub i32 %in, 1
+  %cond = icmp ule i32 %sub, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %sub, 3
+  ret i32 %sh
+}
+
+define i32 @computeNumSignBits_sub3(i32 %in1, i32 %in2) {
+; CHECK-LABEL: @computeNumSignBits_sub3(
+; CHECK-NEXT:    [[SUB:%.*]] = sub i32 [[IN1:%.*]], [[IN2:%.*]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[SUB]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[SUB]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %sub = sub i32 %in1, %in2
+  %cond = icmp ule i32 %sub, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %sub, 3
+  ret i32 %sh
+}
+
+define i32 @computeNumSignBits_mul(i32 %in1, i32 %in2) {
+; CHECK-LABEL: @computeNumSignBits_mul(
+; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[IN1:%.*]], [[IN2:%.*]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[MUL]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[MUL]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %mul = mul i32 %in1, %in2
+  %cond = icmp ule i32 %mul, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %mul, 3
+  ret i32 %sh
+}
+
+define i32 @computeNumSignBits_select(i32 %in, i1 %s) {
+; CHECK-LABEL: @computeNumSignBits_select(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[S:%.*]], i32 [[IN:%.*]], i32 1
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[SEL]], 43
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[SH:%.*]] = shl nuw nsw i32 [[SEL]], 3
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %sel = select i1 %s, i32 %in, i32 1
+  %cond = icmp ule i32 %sel, 42
+  call void @llvm.assume(i1 %cond)
+  %sh = shl i32 %sel, 3
+  ret i32 %sh
+}
+
+declare void @llvm.assume(i1)