Range assume bundles

llvm/include/llvm/Analysis/AssumeBundleQueries.h

@@ -16,11 +16,13 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include <optional>
 
 namespace llvm {
 class AssumptionCache;
 class DominatorTree;
 class Instruction;
+class ConstantRange;
 
 /// Index of elements in the operand bundle.
 /// If the element exist it is guaranteed to be what is specified in this enum
@@ -170,6 +172,9 @@ RetainedKnowledge getKnowledgeValidInContext(
 RetainedKnowledge getKnowledgeFromBundle(AssumeInst &Assume,
                                          const CallBase::BundleOpInfo &BOI);
 
+std::optional<ConstantRange>
+getRangeFromBundle(AssumeInst &Assume, const CallBase::BundleOpInfo &BOI);
+
 } // namespace llvm
 
 #endif

llvm/lib/Analysis/AssumeBundleQueries.cpp

@@ -10,6 +10,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -69,8 +70,14 @@ bool llvm::hasAttributeInAssume(AssumeInst &Assume, Value *IsOn,
 
 void llvm::fillMapFromAssume(AssumeInst &Assume, RetainedKnowledgeMap &Result) {
   for (auto &Bundles : Assume.bundle_op_infos()) {
-    std::pair<Value *, Attribute::AttrKind> Key{
-        nullptr, Attribute::getAttrKindFromName(Bundles.Tag->getKey())};
+    Attribute::AttrKind AttrKind =
+        Attribute::getAttrKindFromName(Bundles.Tag->getKey());
+
+    if (!Attribute::isEnumAttrKind(AttrKind) &&
+        !Attribute::isIntAttrKind(AttrKind))
+      continue;
+
+    std::pair<Value *, Attribute::AttrKind> Key{nullptr, AttrKind};
     if (bundleHasArgument(Bundles, ABA_WasOn))
       Key.first = getValueFromBundleOpInfo(Assume, Bundles, ABA_WasOn);
 
@@ -101,8 +108,14 @@ llvm::getKnowledgeFromBundle(AssumeInst &Assume,
   RetainedKnowledge Result;
   if (!DebugCounter::shouldExecute(AssumeQueryCounter))
     return Result;
+  Attribute::AttrKind AttrKind =
+      Attribute::getAttrKindFromName(BOI.Tag->getKey());
 
-  Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
+  if (!Attribute::isEnumAttrKind(AttrKind) &&
+      !Attribute::isIntAttrKind(AttrKind))
+    return Result;
+
+  Result.AttrKind = AttrKind;
   if (bundleHasArgument(BOI, ABA_WasOn))
     Result.WasOn = getValueFromBundleOpInfo(Assume, BOI, ABA_WasOn);
   auto GetArgOr1 = [&](unsigned Idx) -> uint64_t {
@@ -202,3 +215,21 @@ RetainedKnowledge llvm::getKnowledgeValidInContext(
                                 return isValidAssumeForContext(I, CtxI, DT);
                               });
 }
+
+std::optional<ConstantRange>
+llvm::getRangeFromBundle(AssumeInst &Assume,
+                         const CallBase::BundleOpInfo &BOI) {
+  if (Attribute::getAttrKindFromName(BOI.Tag->getKey()) != Attribute::Range)
+    return std::nullopt;
+
+  assert(BOI.End - BOI.Begin > ABA_Argument + 1 &&
+         "range assumptions should have 3 arguments");
+
+  if (auto *Lower = dyn_cast<ConstantInt>(
+          getValueFromBundleOpInfo(Assume, BOI, ABA_Argument)))
+    if (auto *Upper = dyn_cast<ConstantInt>(
+            getValueFromBundleOpInfo(Assume, BOI, ABA_Argument + 1)))
+      if (Lower->getValue() != Upper->getValue())
+        return ConstantRange(Lower->getValue(), Upper->getValue());
+  return std::nullopt;
+}

llvm/lib/Analysis/ValueTracking.cpp

@@ -595,17 +595,24 @@ static bool isKnownNonZeroFromAssume(const Value *V, const SimplifyQuery &Q) {
            "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 (RK.WasOn == V &&
-            (RK.AttrKind == Attribute::NonNull ||
-             (RK.AttrKind == Attribute::Dereferenceable &&
-              !NullPointerIsDefined(Q.CxtI->getFunction(),
-                                    V->getType()->getPointerAddressSpace()))) &&
-            isValidAssumeForContext(I, Q.CxtI, Q.DT))
-          return true;
+      if (V->getType()->isPointerTy()) {
+        if (RetainedKnowledge RK = getKnowledgeFromBundle(
+                *I, I->bundle_op_info_begin()[Elem.Index])) {
+          if (RK.WasOn == V &&
+              (RK.AttrKind == Attribute::NonNull ||
+               (RK.AttrKind == Attribute::Dereferenceable &&
+                !NullPointerIsDefined(
+                    Q.CxtI->getFunction(),
+                    V->getType()->getPointerAddressSpace()))) &&
+              isValidAssumeForContext(I, Q.CxtI, Q.DT))
+            return true;
+        }
+      } else if (V->getType()->isIntOrIntVectorTy()) {
+        if (std::optional<ConstantRange> Range =
+                getRangeFromBundle(*I, I->bundle_op_info_begin()[Elem.Index]))
+          if (!Range->contains(APInt::getZero(Range->getBitWidth())) &&
+              isValidAssumeForContext(I, Q.CxtI, Q.DT))
+            return true;
       }
       continue;
     }
@@ -824,17 +831,23 @@ 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])) {
-        // 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 (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));
+        }
+      } 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;
     }
@@ -9906,11 +9919,17 @@ ConstantRange llvm::computeConstantRange(const Value *V, bool ForSigned,
     for (auto &AssumeVH : AC->assumptionsFor(V)) {
       if (!AssumeVH)
         continue;
-      CallInst *I = cast<CallInst>(AssumeVH);
+      AssumeInst *I = cast<AssumeInst>(AssumeVH);
       assert(I->getParent()->getParent() == CtxI->getParent()->getParent() &&
              "Got assumption for the wrong function!");
-      assert(I->getIntrinsicID() == Intrinsic::assume &&
-             "must be an assume intrinsic");
+
+      if (AssumeVH.Index != AssumptionCache::ExprResultIdx) {
+        if (std::optional<ConstantRange> Range = getRangeFromBundle(
+                *I, I->bundle_op_info_begin()[AssumeVH.Index]))
+          if (isValidAssumeForContext(I, CtxI, DT))
+            CR = CR.intersectWith(*Range);
+        continue;
+      }
 
       if (!isValidAssumeForContext(I, CtxI, DT))
         continue;

llvm/lib/IR/Verifier.cpp

@@ -5416,6 +5416,23 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
                 "third argument should be an integer if present", Call);
         return;
       }
+      if (Kind == Attribute::Range) {
+        Check(ArgCount == 3, "range assumptions should have 3 arguments", Call);
+        Type *FirstType = Call.getOperand(Elem.Begin)->getType();
+        Check(FirstType->isIntOrIntVectorTy(),
+              "first argument should be an integer or vector of integers",
+              Call);
+        Type *ST = FirstType->getScalarType();
+        Check(Call.getOperand(Elem.Begin + 1)->getType() == ST,
+              "second argument should be an integer with same bit width as the "
+              "first argument",
+              Call);
+        Check(Call.getOperand(Elem.Begin + 2)->getType() == ST,
+              "third argument should be an integer with same bit width as the "
+              "first argument",
+              Call);
+        return;
+      }
       Check(ArgCount <= 2, "too many arguments", Call);
       if (Kind == Attribute::None)
         break;

llvm/test/Analysis/BasicAA/range.ll

@@ -271,6 +271,18 @@ entry:
   ret i32 %load_
 }
 
+; CHECK-LABEL: Function: range_assume
+; CHECK: NoAlias: i32* %gep1, i32* %gep2
+define void @range_assume(ptr %s, ptr %q) {
+  %in_array = load i32, ptr %q
+  call void @llvm.assume(i1 true) ["range"(i32 %in_array, i32 0, i32 2)]
+  %gep1 = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1, i32 %in_array
+  %gep2 = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 2
+  load i32, ptr %gep1
+  load i32, ptr %gep2
+  ret void
+}
+
 declare void @llvm.assume(i1)
 
 !0 = !{ i32 0, i32 2 }

llvm/test/Analysis/ValueTracking/known-non-zero.ll

@@ -1546,4 +1546,52 @@ define i1 @vec_reverse_non_zero_demanded_fail(<4 x i8> %xx) {
   ret i1 %r
 }
 
+define i1 @range_assume(i8 %x, i8 %y) {
+; CHECK-LABEL: @range_assume(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i8 [[X:%.*]], i8 1, i8 0) ]
+; CHECK-NEXT:    ret i1 false
+;
+  call void @llvm.assume(i1 true) ["range"(i8 %x, i8 1, i8 0)]
+  %or = or i8 %y, %x
+  %cmp = icmp eq i8 %or, 0
+  ret i1 %cmp
+}
+
+define i1 @neg_range_assum(i8 %x, i8 %y) {
+; CHECK-LABEL: @neg_range_assum(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i8 [[X:%.*]], i8 -1, i8 1) ]
+; CHECK-NEXT:    [[OR:%.*]] = or i8 [[Y:%.*]], [[X]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[OR]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  call void @llvm.assume(i1 true) ["range"(i8 %x, i8 -1, i8 1)]
+  %or = or i8 %y, %x
+  %cmp = icmp eq i8 %or, 0
+  ret i1 %cmp
+}
+
+define <2 x i1> @range_assum_vec(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @range_assum_vec(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i8> [[X:%.*]], i8 1, i8 0) ]
+; CHECK-NEXT:    ret <2 x i1> <i1 true, i1 true>
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i8> %x, i8 1, i8 0)]
+  %or = or <2 x i8> %y, %x
+  %cmp = icmp ne <2 x i8> %or, zeroinitializer
+  ret <2 x i1> %cmp
+}
+
+define <2 x i1> @neg_range_assum_vec(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @neg_range_assum_vec(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i8> [[X:%.*]], i8 -1, i8 1) ]
+; CHECK-NEXT:    [[OR:%.*]] = or <2 x i8> [[Y:%.*]], [[X]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i8> [[OR]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i8> %x, i8 -1, i8 1)]
+  %or = or <2 x i8> %y, %x
+  %cmp = icmp ne <2 x i8> %or, zeroinitializer
+  ret <2 x i1> %cmp
+}
+
 declare i32 @llvm.experimental.get.vector.length.i32(i32, i32, i1)

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

@@ -499,3 +499,46 @@ define <1 x i64> @bitcast_noshift_vector_wrong_type(<2 x float> %v1, <1 x i64> %
   %r = shl <1 x i64> %v2, %b
   ret <1 x i64> %r
 }
+
+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:    ret i32 poison
+;
+  call void @llvm.assume(i1 true) ["range"(i32 %b, i32 32, i32 64)]
+  %shl = shl i32 %a, %b
+  ret i32 %shl
+}
+
+define i32 @neg_shl_amount_is_known_bogus_range_assum(i32 %a, i32 %b) {
+; CHECK-LABEL: @neg_shl_amount_is_known_bogus_range_assum(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(i32 [[B:%.*]], i32 0, i32 32) ]
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret i32 [[SHL]]
+;
+  call void @llvm.assume(i1 true) ["range"(i32 %b, i32 0, i32 32)]
+  %shl = shl i32 %a, %b
+  ret i32 %shl
+}
+
+
+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:    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
+  ret <2 x i32> %shl
+}
+
+define <2 x i32> @neg_shl_amount_is_known_bogus_range_assum_vec(<2 x i32> %a, <2 x i32> %b) {
+; CHECK-LABEL: @neg_shl_amount_is_known_bogus_range_assum_vec(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "range"(<2 x i32> [[B:%.*]], i32 0, i32 32) ]
+; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret <2 x i32> [[SHL]]
+;
+  call void @llvm.assume(i1 true) ["range"(<2 x i32> %b, i32 0, i32 32)]
+  %shl = shl <2 x i32> %a, %b
+  ret <2 x i32> %shl
+}

llvm/test/Verifier/assume-bundles.ll

@@ -28,5 +28,13 @@ define void @func(ptr %P, i32 %P1, ptr %P2, ptr %P3) {
   call void @llvm.assume(i1 true) ["separate_storage"(ptr %P)]
 ; CHECK: arguments to separate_storage assumptions should be pointers
   call void @llvm.assume(i1 true) ["separate_storage"(ptr %P, i32 123)]
+; CHECK: range assumptions should have 3 arguments
+  call void @llvm.assume(i1 true) ["range"(i32 %P1, i32 0)]
+; CHECK: first argument should be an integer or vector of integers
+  call void @llvm.assume(i1 true) ["range"(ptr %P, i32 %P1, i32 4)]
+; CHECK: second argument should be an integer with same bit width as the first argument
+  call void @llvm.assume(i1 true) ["range"(i32 %P1, ptr %P, i32 4)]
+; CHECK: third argument should be an integer with same bit width as the first argument
+  call void @llvm.assume(i1 true) ["range"(i32 %P1, i32 4, i8 10)]
   ret void
 }

llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp

@@ -411,8 +411,10 @@ static void RunRandTest(uint64_t Seed, int Size, int MinCount, int MaxCount,
   std::mt19937 Rng(Seed);
   std::uniform_int_distribution<int> DistCount(MinCount, MaxCount);
   std::uniform_int_distribution<unsigned> DistValue(0, MaxValue);
-  std::uniform_int_distribution<unsigned> DistAttr(0,
-                                                   Attribute::EndAttrKinds - 1);
+  std::uniform_int_distribution<unsigned> DistEnumAttr(Attribute::FirstEnumAttr,
+                                                       Attribute::LastEnumAttr);
+  std::uniform_int_distribution<unsigned> DistIntAttr(Attribute::FirstIntAttr,
+                                                      Attribute::LastIntAttr);
 
   std::unique_ptr<Module> Mod = std::make_unique<Module>("AssumeQueryAPI", C);
   if (!Mod)
@@ -449,7 +451,7 @@ static void RunRandTest(uint64_t Seed, int Size, int MinCount, int MaxCount,
   for (int i = 0; i < Size; i++) {
     int count = DistCount(Rng);
     int value = DistValue(Rng);
-    int attr = DistAttr(Rng);
+    int attr = count > 1 ? DistIntAttr(Rng) : DistEnumAttr(Rng);
     std::string str;
     raw_string_ostream ss(str);
     ss << Attribute::getNameFromAttrKind(