[ValueTracking] Handle range assume bundles in computeKnownBits

Author: andjo403

llvm/lib/Analysis/ValueTracking.cpp

@@ -831,17 +831,22 @@ void llvm::computeKnownBitsFromContext(const Value *V, KnownBits &Known,
            "Got assumption for the wrong function!");
 
     if (Elem.Index != AssumptionCache::ExprResultIdx) {
-      if (!V->getType()->isPointerTy())
-        continue;
-      if (RetainedKnowledge RK = getKnowledgeFromBundle(
-              *I, I->bundle_op_info_begin()[Elem.Index])) {
+      if (V->getType()->isPointerTy()) {
+        if (RetainedKnowledge RK = getKnowledgeFromBundle(
+                *I, I->bundle_op_info_begin()[Elem.Index])) {
         // Allow AllowEphemerals in isValidAssumeForContext, as the CxtI might
         // be the producer of the pointer in the bundle. At the moment, align
         // assumptions aren't optimized away.
-        if (RK.WasOn == V && RK.AttrKind == Attribute::Alignment &&
-            isPowerOf2_64(RK.ArgValue) &&
-            isValidAssumeForContext(I, Q.CxtI, Q.DT, /*AllowEphemerals*/ true))
-          Known.Zero.setLowBits(Log2_64(RK.ArgValue));
+          if (RK.WasOn == V && RK.AttrKind == Attribute::Alignment &&
+              isPowerOf2_64(RK.ArgValue) &&
+              isValidAssumeForContext(I, Q.CxtI, Q.DT, /*AllowEphemerals*/ true))
+            Known.Zero.setLowBits(Log2_64(RK.ArgValue));
+        }
+      } else if (V->getType()->isIntOrIntVectorTy()) {
+        if (std::optional<ConstantRange> Range =
+                getRangeFromBundle(*I, I->bundle_op_info_begin()[Elem.Index]))
+          if (isValidAssumeForContext(I, Q.CxtI, Q.DT))
+            Known = Known.unionWith(Range->toKnownBits());
       }
       continue;
     }

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

@@ -503,8 +503,7 @@ define <1 x i64> @bitcast_noshift_vector_wrong_type(<2 x float> %v1, <1 x i64> %
 define i32 @shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_assum(
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i32 [[B:%.*]], i32 32, i32 64) ]
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret i32 [[SHL]]
+; CHECK-NEXT:    ret i32 poison
 ;
   call void @llvm.assume(i1 true) ["range"(i32 %b, i32 32, i32 64)]
   %shl = shl i32 %a, %b
@@ -526,8 +525,7 @@ define i32 @neg_shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
 define <2 x i32> @shl_amount_is_known_bogus_range_assum_vec(<2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: @shl_amount_is_known_bogus_range_assum_vec(
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i32> [[B:%.*]], i32 32, i32 64) ]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+; CHECK-NEXT:    ret <2 x i32> poison
 ;
   call void @llvm.assume(i1 true) ["range"(<2 x i32> %b, i32 32, i32 64)]
   %shl = shl <2 x i32> %a, %b