handle range attribute in computeKnownBits

Author: andjo403

llvm/lib/Analysis/ValueTracking.cpp

@@ -1467,14 +1467,20 @@ static void computeKnownBitsFromOperator(const Operator *I,
     break;
   }
   case Instruction::Call:
-  case Instruction::Invoke:
+  case Instruction::Invoke: {
     // If range metadata is attached to this call, set known bits from that,
     // and then intersect with known bits based on other properties of the
     // function.
     if (MDNode *MD =
             Q.IIQ.getMetadata(cast<Instruction>(I), LLVMContext::MD_range))
       computeKnownBitsFromRangeMetadata(*MD, Known);
-    if (const Value *RV = cast<CallBase>(I)->getReturnedArgOperand()) {
+
+    const CallBase *CB = cast<CallBase>(I);
+
+    if (std::optional<ConstantRange> Range = CB->getRange())
+      Known = Range->toKnownBits();
+
+    if (const Value *RV = CB->getReturnedArgOperand()) {
       if (RV->getType() == I->getType()) {
         computeKnownBits(RV, Known2, Depth + 1, Q);
         Known = Known.unionWith(Known2);
@@ -1646,6 +1652,7 @@ static void computeKnownBitsFromOperator(const Operator *I,
       }
     }
     break;
+  }
   case Instruction::ShuffleVector: {
     auto *Shuf = dyn_cast<ShuffleVectorInst>(I);
     // FIXME: Do we need to handle ConstantExpr involving shufflevectors?
@@ -1900,6 +1907,10 @@ void computeKnownBits(const Value *V, const APInt &DemandedElts,
   // assumptions.  Confirm that we've handled them all.
   assert(!isa<ConstantData>(V) && "Unhandled constant data!");
 
+  if (const Argument *A = dyn_cast<Argument>(V))
+    if (std::optional<ConstantRange> Range = A->getRange())
+      Known = Range->toKnownBits();
+
   // All recursive calls that increase depth must come after this.
   if (Depth == MaxAnalysisRecursionDepth)
     return;

llvm/test/Transforms/InstSimplify/shift-knownbits.ll

@@ -16,8 +16,7 @@ define i32 @shl_amount_is_known_bogus(i32 %a, i32 %b) {
 
 define i32 @shl_amount_is_known_bogus_range_attr(i32 %a, i32 range(i32 32, 64) %b) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_attr(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    ret i32 [[SHL]]
+; CHECK-NEXT:    ret i32 poison
 ;
   %shl = shl i32 %a, %b
   ret i32 %shl
@@ -38,8 +37,7 @@ declare range(i32 0, 32) i32 @returns_in_range_helper()
 define i32 @shl_amount_is_known_bogus_range_return(i32 %a) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_return(
 ; CHECK-NEXT:    [[B:%.*]] = call i32 @returns_out_of_range_helper()
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret i32 [[SHL]]
+; CHECK-NEXT:    ret i32 poison
 ;
   %b = call i32 @returns_out_of_range_helper()
   %shl = shl i32 %a, %b
@@ -62,8 +60,7 @@ declare i32 @returns_i32_helper()
 define i32 @shl_amount_is_known_bogus_range_call(i32 %a) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_call(
 ; CHECK-NEXT:    [[B:%.*]] = call range(i32 32, 64) i32 @returns_i32_helper()
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret i32 [[SHL]]
+; CHECK-NEXT:    ret i32 poison
 ;
   %b = call range(i32 32, 64) i32 @returns_i32_helper()
   %shl = shl i32 %a, %b
@@ -83,8 +80,7 @@ define i32 @neg_shl_amount_is_known_bogus_range_call(i32 %a) {
 
 define <2 x i32> @shl_amount_is_known_bogus_range_attr_vec(<2 x i32> %a, <2 x i32> range(i32 32, 64) %b) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_attr_vec(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+; CHECK-NEXT:    ret <2 x i32> poison
 ;
   %shl = shl <2 x i32> %a, %b
   ret <2 x i32> %shl
@@ -105,8 +101,7 @@ declare range(i32 0, 32) <2 x i32> @returns_in_range_helper_vec()
 define <2 x i32> @shl_amount_is_known_bogus_range_return_vec(<2 x i32> %a) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_return_vec(
 ; CHECK-NEXT:    [[B:%.*]] = call <2 x i32> @returns_out_of_range_helper_vec()
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+; CHECK-NEXT:    ret <2 x i32> poison
 ;
   %b = call <2 x i32> @returns_out_of_range_helper_vec()
   %shl = shl <2 x i32> %a, %b
@@ -129,8 +124,7 @@ declare <2 x i32> @returns_i32_helper_vec()
 define <2 x i32> @shl_amount_is_known_bogus_range_call_vec(<2 x i32> %a) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_call_vec(
 ; CHECK-NEXT:    [[B:%.*]] = call range(i32 32, 64) <2 x i32> @returns_i32_helper_vec()
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+; CHECK-NEXT:    ret <2 x i32> poison
 ;
   %b = call range(i32 32, 64) <2 x i32> @returns_i32_helper_vec()
   %shl = shl <2 x i32> %a, %b