[AssumeBundles] Enforce constraints on the operand bundle of llvm.assume

Summary: Add verification that operand bundles on an llvm.assume are well formed to the verify pass.

Reviewers: jdoerfert

Reviewed By: jdoerfert

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D75269

llvm/include/llvm/IR/Attributes.h

@@ -116,6 +116,10 @@ class Attribute {
   /// Return true if and only if the attribute has an Argument.
   static bool doesAttrKindHaveArgument(Attribute::AttrKind AttrKind);
 
+  /// Return true if the provided string matches the IR name of an attribute.
+  /// example: "noalias" return true but not "NoAlias"
+  static bool isExistingAttribute(StringRef Name);
+
   //===--------------------------------------------------------------------===//
   // Attribute Accessors
   //===--------------------------------------------------------------------===//

llvm/lib/IR/Attributes.cpp

@@ -208,6 +208,14 @@ bool Attribute::doesAttrKindHaveArgument(Attribute::AttrKind AttrKind) {
          AttrKind == Attribute::DereferenceableOrNull;
 }
 
+bool Attribute::isExistingAttribute(StringRef Name) {
+  return StringSwitch<bool>(Name)
+#define GET_ATTR_NAMES
+#define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) .Case(#DISPLAY_NAME, true)
+#include "llvm/IR/Attributes.inc"
+      .Default(false);
+}
+
 //===----------------------------------------------------------------------===//
 // Attribute Accessor Methods
 //===----------------------------------------------------------------------===//

llvm/lib/IR/KnowledgeRetention.cpp

@@ -183,25 +183,14 @@ static Value *getValueFromBundleOpInfo(IntrinsicInst &Assume,
   return (Assume.op_begin() + BOI.Begin + Idx)->get();
 }
 
-#ifndef NDEBUG
-
-static bool isExistingAttribute(StringRef Name) {
-  return StringSwitch<bool>(Name)
-#define GET_ATTR_NAMES
-#define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) .Case(#DISPLAY_NAME, true)
-#include "llvm/IR/Attributes.inc"
-      .Default(false);
-}
-
-#endif
-
 bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
                                 StringRef AttrName, uint64_t *ArgVal,
                                 AssumeQuery AQR) {
   IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
   assert(Assume.getIntrinsicID() == Intrinsic::assume &&
          "this function is intended to be used on llvm.assume");
-  assert(isExistingAttribute(AttrName) && "this attribute doesn't exist");
+  assert(Attribute::isExistingAttribute(AttrName) &&
+         "this attribute doesn't exist");
   assert((ArgVal == nullptr || Attribute::doesAttrKindHaveArgument(
                                    Attribute::getAttrKindFromName(AttrName))) &&
          "requested value for an attribute that has no argument");
@@ -218,16 +207,12 @@ bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
     Lookup =
         llvm::lower_bound(Assume.bundle_op_infos(), AttrName,
                           [](const CallBase::BundleOpInfo &BOI, StringRef RHS) {
-                            assert(isExistingAttribute(BOI.Tag->getKey()) &&
-                                   "this attribute doesn't exist");
                             return BOI.Tag->getKey() < RHS;
                           });
   else
     Lookup = std::prev(
         llvm::upper_bound(Assume.bundle_op_infos(), AttrName,
                           [](StringRef LHS, const CallBase::BundleOpInfo &BOI) {
-                            assert(isExistingAttribute(BOI.Tag->getKey()) &&
-                                   "this attribute doesn't exist");
                             return LHS < BOI.Tag->getKey();
                           }));
 

llvm/lib/IR/Verifier.cpp

@@ -4320,6 +4320,29 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
   switch (ID) {
   default:
     break;
+  case Intrinsic::assume: {
+    for (auto &Elem : Call.bundle_op_infos()) {
+      Assert(Attribute::isExistingAttribute(Elem.Tag->getKey()),
+             "tags must be valid attribute names");
+      Assert(Elem.End - Elem.Begin <= 2, "to many arguments");
+      Attribute::AttrKind Kind =
+          Attribute::getAttrKindFromName(Elem.Tag->getKey());
+      if (Kind == Attribute::None)
+        break;
+      if (Attribute::doesAttrKindHaveArgument(Kind)) {
+        Assert(Elem.End - Elem.Begin == 2,
+               "this attribute should have 2 arguments");
+        Assert(isa<ConstantInt>(Call.getOperand(Elem.Begin + 1)),
+               "the second argument should be a constant integral value");
+      } else if (isFuncOnlyAttr(Kind)) {
+        Assert((Elem.End - Elem.Begin) == 0, "this attribute has no argument");
+      } else if (!isFuncOrArgAttr(Kind)) {
+        Assert((Elem.End - Elem.Begin) == 1,
+               "this attribute should have one argument");
+      }
+    }
+    break;
+  }
   case Intrinsic::coro_id: {
     auto *InfoArg = Call.getArgOperand(3)->stripPointerCasts();
     if (isa<ConstantPointerNull>(InfoArg))

llvm/test/IR/assume-builder.ll

@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -passes='assume-builder' -S %s | FileCheck %s --check-prefixes=BASIC
-; RUN: opt -passes='assume-builder' --assume-preserve-all -S %s | FileCheck %s --check-prefixes=ALL
+; RUN: opt -passes='assume-builder,verify' -S %s | FileCheck %s --check-prefixes=BASIC
+; RUN: opt -passes='assume-builder,verify' --assume-preserve-all -S %s | FileCheck %s --check-prefixes=ALL
 
 declare void @func(i32*, i32*)
 declare void @func_cold(i32*) cold

llvm/test/Verifier/assume-bundles.ll

@@ -0,0 +1,19 @@
+; RUN: not opt -verify < %s 2>&1 | FileCheck %s
+
+declare void @llvm.assume(i1)
+
+define void @func(i32* %P, i32 %P1, i32* %P2, i32* %P3) {
+; CHECK: tags must be valid attribute names
+  call void @llvm.assume(i1 true) ["adazdazd"()]
+; CHECK: the second argument should be a constant integral value
+  call void @llvm.assume(i1 true) ["align"(i32* %P, i32 %P1)]
+; CHECK: to many arguments
+  call void @llvm.assume(i1 true) ["align"(i32* %P, i32 8, i32 8)]
+; CHECK: this attribute should have 2 arguments
+  call void @llvm.assume(i1 true) ["align"(i32* %P)]
+; CHECK: this attribute has no argument
+  call void @llvm.assume(i1 true) ["align"(i32* %P, i32 4), "cold"(i32* %P)]
+; CHECK: this attribute should have one argument
+  call void @llvm.assume(i1 true) ["noalias"()]
+  ret void
+}