[IR] Allow range assume bundles

Author: andjo403

llvm/lib/Analysis/AssumeBundleQueries.cpp

@@ -69,8 +69,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 +107,14 @@ llvm::getKnowledgeFromBundle(AssumeInst &Assume,
   RetainedKnowledge Result;
   if (!DebugCounter::shouldExecute(AssumeQueryCounter))
     return Result;
+  Attribute::AttrKind AttrKind =
+      Attribute::getAttrKindFromName(BOI.Tag->getKey());
+
+  if (!Attribute::isEnumAttrKind(AttrKind) &&
+      !Attribute::isIntAttrKind(AttrKind))
+    return Result;
 
-  Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
+  Result.AttrKind = AttrKind;
   if (bundleHasArgument(BOI, ABA_WasOn))
     Result.WasOn = getValueFromBundleOpInfo(Assume, BOI, ABA_WasOn);
   auto GetArgOr1 = [&](unsigned Idx) -> uint64_t {

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/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(