attributes: introduce allockind attr for describing allocator fn behavior

I chose to encode the allockind information in a string constant because
otherwise we would get a bit of an explosion of keywords to deal with
the possible permutations of allocation function types.

I'm not sure that CodeGen.h is the correct place for this enum, but it
seemed to kind of match the UWTableKind enum so I put it in the same
place. Constructive suggestions on a better location most certainly
encouraged.

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

llvm/docs/LangRef.rst

@@ -1584,6 +1584,27 @@ example:
     "_ZnwmSt11align_val_t" for aligned ``::operator::new`` and
     ``::operator::delete``. Matching malloc/realloc/free calls within a family
     can be optimized, but mismatched ones will be left alone.
+``allockind("KIND")``
+    Describes the behavior of an allocation function. The KIND string contains comma
+    separated entries from the following options:
+     * "alloc": the function returns a new block of memory or null.
+     * "realloc": the function returns a new block of memory or null. If the
+       result is non-null the memory contents from the start of the block up to
+       the smaller of the original allocation size and the new allocation size
+       will match that of the ``allocptr`` argument and the ``allocptr``
+       argument is invalidated, even if the function returns the same address.
+     * "free": the function frees the block of memory specified by ``allocptr``.
+     * "uninitialized": Any newly-allocated memory (either a new block from 
+       a "alloc" function or the enlarged capacity from a "realloc" function)
+       will be uninitialized.
+     * "zeroed": Any newly-allocated memory (either a new block from a "alloc"
+       function or the enlarged capacity from a "realloc" function) will be
+       zeroed.
+     * "aligned": the function returns memory aligned according to the 
+       ``allocalign`` parameter.
+    The first three options are mutually exclusive, and the remaining options
+    describe more details of how the function behaves. The remaining options
+    are invalid for "free"-type functions.
 ``allocsize(<EltSizeParam>[, <NumEltsParam>])``
     This attribute indicates that the annotated function will always return at
     least a given number of bytes (or null). Its arguments are zero-indexed

llvm/include/llvm/AsmParser/LLParser.h

@@ -274,6 +274,7 @@ namespace llvm {
                                 bool AllowParens = false);
     bool parseOptionalDerefAttrBytes(lltok::Kind AttrKind, uint64_t &Bytes);
     bool parseOptionalUWTableKind(UWTableKind &Kind);
+    bool parseAllocKind(AllocFnKind &Kind);
     bool parseScopeAndOrdering(bool IsAtomic, SyncScope::ID &SSID,
                                AtomicOrdering &Ordering);
     bool parseScope(SyncScope::ID &SSID);

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -684,6 +684,7 @@ enum AttributeKindCodes {
   ATTR_KIND_NO_SANITIZE_BOUNDS = 79,
   ATTR_KIND_ALLOC_ALIGN = 80,
   ATTR_KIND_ALLOCATED_POINTER = 81,
+  ATTR_KIND_ALLOC_KIND = 82,
 };
 
 enum ComdatSelectionKindCodes {

llvm/include/llvm/IR/Attributes.h

@@ -17,6 +17,7 @@
 
 #include "llvm-c/Types.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
@@ -43,6 +44,18 @@ class Function;
 class LLVMContext;
 class Type;
 
+enum class AllocFnKind : uint64_t {
+  Unknown = 0,
+  Alloc = 1 << 0,         // Allocator function returns a new allocation
+  Realloc = 1 << 1,       // Allocator function resizes the `allocptr` argument
+  Free = 1 << 2,          // Allocator function frees the `allocptr` argument
+  Uninitialized = 1 << 3, // Allocator function returns uninitialized memory
+  Zeroed = 1 << 4,        // Allocator function returns zeroed memory
+  Aligned = 1 << 5,       // Allocator function aligns allocations per the
+                          // `allocalign` argument
+  LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Aligned)
+};
+
 //===----------------------------------------------------------------------===//
 /// \class
 /// Functions, function parameters, and return types can have attributes
@@ -228,6 +241,9 @@ public:
   // Returns the unwind table kind.
   UWTableKind getUWTableKind() const;
 
+  // Returns the allocator function kind.
+  AllocFnKind getAllocKind() const;
+
   /// The Attribute is converted to a string of equivalent mnemonic. This
   /// is, presumably, for writing out the mnemonics for the assembly writer.
   std::string getAsString(bool InAttrGrp = false) const;
@@ -359,6 +375,7 @@ public:
   unsigned getVScaleRangeMin() const;
   Optional<unsigned> getVScaleRangeMax() const;
   UWTableKind getUWTableKind() const;
+  AllocFnKind getAllocKind() const;
   std::string getAsString(bool InAttrGrp = false) const;
 
   /// Return true if this attribute set belongs to the LLVMContext.
@@ -850,6 +867,8 @@ public:
   /// Get the unwind table kind requested for the function.
   UWTableKind getUWTableKind() const;
 
+  AllocFnKind getAllocKind() const;
+
   /// Return the attributes at the index as a string.
   std::string getAsString(unsigned Index, bool InAttrGrp = false) const;
 
@@ -1203,6 +1222,9 @@ public:
   /// Attribute.
   AttrBuilder &addUWTableAttr(UWTableKind Kind);
 
+  // This turns the allocator kind into the form used internally in Attribute.
+  AttrBuilder &addAllocKindAttr(AllocFnKind Kind);
+
   ArrayRef<Attribute> attrs() const { return Attrs; }
 
   bool operator==(const AttrBuilder &B) const;

llvm/include/llvm/IR/Attributes.td

@@ -51,6 +51,9 @@ def Alignment : IntAttr<"align", [ParamAttr, RetAttr]>;
 /// aligned_alloc and aligned ::operator::new.
 def AllocAlign: EnumAttr<"allocalign", [ParamAttr]>;
 
+/// Describes behavior of an allocator function in terms of known properties.
+def AllocKind: IntAttr<"allockind", [FnAttr]>;
+
 /// Parameter is the pointer to be manipulated by the allocator function.
 def AllocatedPointer : EnumAttr<"allocptr", [ParamAttr]>;
 

llvm/lib/AsmParser/LLParser.cpp

@@ -1362,6 +1362,13 @@ bool LLParser::parseEnumAttribute(Attribute::AttrKind Attr, AttrBuilder &B,
     B.addUWTableAttr(Kind);
     return false;
   }
+  case Attribute::AllocKind: {
+    AllocFnKind Kind = AllocFnKind::Unknown;
+    if (parseAllocKind(Kind))
+      return true;
+    B.addAllocKindAttr(Kind);
+    return false;
+  }
   default:
     B.addAttribute(Attr);
     Lex.Lex();
@@ -2041,6 +2048,40 @@ bool LLParser::parseOptionalUWTableKind(UWTableKind &Kind) {
   return parseToken(lltok::rparen, "expected ')'");
 }
 
+bool LLParser::parseAllocKind(AllocFnKind &Kind) {
+  Lex.Lex();
+  LocTy ParenLoc = Lex.getLoc();
+  if (!EatIfPresent(lltok::lparen))
+    return error(ParenLoc, "expected '('");
+  LocTy KindLoc = Lex.getLoc();
+  std::string Arg;
+  if (parseStringConstant(Arg))
+    return error(KindLoc, "expected allockind value");
+  for (StringRef A : llvm::split(Arg, ",")) {
+    if (A == "alloc") {
+      Kind |= AllocFnKind::Alloc;
+    } else if (A == "realloc") {
+      Kind |= AllocFnKind::Realloc;
+    } else if (A == "free") {
+      Kind |= AllocFnKind::Free;
+    } else if (A == "uninitialized") {
+      Kind |= AllocFnKind::Uninitialized;
+    } else if (A == "zeroed") {
+      Kind |= AllocFnKind::Zeroed;
+    } else if (A == "aligned") {
+      Kind |= AllocFnKind::Aligned;
+    } else {
+      return error(KindLoc, Twine("unknown allockind ") + A);
+    }
+  }
+  ParenLoc = Lex.getLoc();
+  if (!EatIfPresent(lltok::rparen))
+    return error(ParenLoc, "expected ')'");
+  if (Kind == AllocFnKind::Unknown)
+    return error(KindLoc, "expected allockind value");
+  return false;
+}
+
 /// parseOptionalCommaAlign
 ///   ::=
 ///   ::= ',' align 4

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -1536,6 +1536,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::DereferenceableOrNull;
   case bitc::ATTR_KIND_ALLOC_ALIGN:
     return Attribute::AllocAlign;
+  case bitc::ATTR_KIND_ALLOC_KIND:
+    return Attribute::AllocKind;
   case bitc::ATTR_KIND_ALLOC_SIZE:
     return Attribute::AllocSize;
   case bitc::ATTR_KIND_ALLOCATED_POINTER:
@@ -1736,6 +1738,8 @@ Error BitcodeReader::parseAttributeGroupBlock() {
             B.addVScaleRangeAttrFromRawRepr(Record[++i]);
           else if (Kind == Attribute::UWTable)
             B.addUWTableAttr(UWTableKind(Record[++i]));
+          else if (Kind == Attribute::AllocKind)
+            B.addAllocKindAttr(static_cast<AllocFnKind>(Record[++i]));
         } else if (Record[i] == 3 || Record[i] == 4) { // String attribute
           bool HasValue = (Record[i++] == 4);
           SmallString<64> KindStr;

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -650,6 +650,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_MIN_SIZE;
   case Attribute::AllocatedPointer:
     return bitc::ATTR_KIND_ALLOCATED_POINTER;
+  case Attribute::AllocKind:
+    return bitc::ATTR_KIND_ALLOC_KIND;
   case Attribute::Naked:
     return bitc::ATTR_KIND_NAKED;
   case Attribute::Nest:

llvm/lib/IR/AttributeImpl.h

@@ -256,6 +256,7 @@ public:
   unsigned getVScaleRangeMin() const;
   Optional<unsigned> getVScaleRangeMax() const;
   UWTableKind getUWTableKind() const;
+  AllocFnKind getAllocKind() const;
   std::string getAsString(bool InAttrGrp) const;
   Type *getAttributeType(Attribute::AttrKind Kind) const;
 

llvm/lib/IR/Attributes.cpp

@@ -376,6 +376,12 @@ UWTableKind Attribute::getUWTableKind() const {
   return UWTableKind(pImpl->getValueAsInt());
 }
 
+AllocFnKind Attribute::getAllocKind() const {
+  assert(hasAttribute(Attribute::AllocKind) &&
+         "Trying to get allockind value from non-allockind attribute");
+  return AllocFnKind(pImpl->getValueAsInt());
+}
+
 std::string Attribute::getAsString(bool InAttrGrp) const {
   if (!pImpl) return {};
 
@@ -447,6 +453,26 @@ std::string Attribute::getAsString(bool InAttrGrp) const {
     }
   }
 
+  if (hasAttribute(Attribute::AllocKind)) {
+    AllocFnKind Kind = getAllocKind();
+    SmallVector<StringRef> parts;
+    if ((Kind & AllocFnKind::Alloc) != AllocFnKind::Unknown)
+      parts.push_back("alloc");
+    if ((Kind & AllocFnKind::Realloc) != AllocFnKind::Unknown)
+      parts.push_back("realloc");
+    if ((Kind & AllocFnKind::Free) != AllocFnKind::Unknown)
+      parts.push_back("free");
+    if ((Kind & AllocFnKind::Uninitialized) != AllocFnKind::Unknown)
+      parts.push_back("uninitialized");
+    if ((Kind & AllocFnKind::Zeroed) != AllocFnKind::Unknown)
+      parts.push_back("zeroed");
+    if ((Kind & AllocFnKind::Aligned) != AllocFnKind::Unknown)
+      parts.push_back("aligned");
+    return ("allockind(\"" +
+            Twine(llvm::join(parts.begin(), parts.end(), ",")) + "\")")
+        .str();
+  }
+
   // Convert target-dependent attributes to strings of the form:
   //
   //   "kind"
@@ -735,6 +761,10 @@ UWTableKind AttributeSet::getUWTableKind() const {
   return SetNode ? SetNode->getUWTableKind() : UWTableKind::None;
 }
 
+AllocFnKind AttributeSet::getAllocKind() const {
+  return SetNode ? SetNode->getAllocKind() : AllocFnKind::Unknown;
+}
+
 std::string AttributeSet::getAsString(bool InAttrGrp) const {
   return SetNode ? SetNode->getAsString(InAttrGrp) : "";
 }
@@ -907,6 +937,12 @@ UWTableKind AttributeSetNode::getUWTableKind() const {
   return UWTableKind::None;
 }
 
+AllocFnKind AttributeSetNode::getAllocKind() const {
+  if (auto A = findEnumAttribute(Attribute::AllocKind))
+    return A->getAllocKind();
+  return AllocFnKind::Unknown;
+}
+
 std::string AttributeSetNode::getAsString(bool InAttrGrp) const {
   std::string Str;
   for (iterator I = begin(), E = end(); I != E; ++I) {
@@ -1463,6 +1499,10 @@ UWTableKind AttributeList::getUWTableKind() const {
   return getFnAttrs().getUWTableKind();
 }
 
+AllocFnKind AttributeList::getAllocKind() const {
+  return getFnAttrs().getAllocKind();
+}
+
 std::string AttributeList::getAsString(unsigned Index, bool InAttrGrp) const {
   return getAttributes(Index).getAsString(InAttrGrp);
 }
@@ -1690,6 +1730,10 @@ AttrBuilder &AttrBuilder::addUWTableAttr(UWTableKind Kind) {
   return addRawIntAttr(Attribute::UWTable, uint64_t(Kind));
 }
 
+AttrBuilder &AttrBuilder::addAllocKindAttr(AllocFnKind Kind) {
+  return addRawIntAttr(Attribute::AllocKind, static_cast<uint64_t>(Kind));
+}
+
 Type *AttrBuilder::getTypeAttr(Attribute::AttrKind Kind) const {
   assert(Attribute::isTypeAttrKind(Kind) && "Not a type attribute");
   Attribute A = getAttribute(Kind);

llvm/lib/IR/Verifier.cpp

@@ -2080,6 +2080,25 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
       return;
   }
 
+  if (Attrs.hasFnAttr(Attribute::AllocKind)) {
+    AllocFnKind K = Attrs.getAllocKind();
+    AllocFnKind Type =
+        K & (AllocFnKind::Alloc | AllocFnKind::Realloc | AllocFnKind::Free);
+    if (!is_contained(
+            {AllocFnKind::Alloc, AllocFnKind::Realloc, AllocFnKind::Free},
+            Type))
+      CheckFailed(
+          "'allockind()' requires exactly one of alloc, realloc, and free");
+    if ((Type == AllocFnKind::Free) &&
+        ((K & (AllocFnKind::Uninitialized | AllocFnKind::Zeroed |
+               AllocFnKind::Aligned)) != AllocFnKind::Unknown))
+      CheckFailed("'allockind(\"free\")' doesn't allow uninitialized, zeroed, "
+                  "or aligned modifiers.");
+    AllocFnKind ZeroedUninit = AllocFnKind::Uninitialized | AllocFnKind::Zeroed;
+    if ((K & ZeroedUninit) == ZeroedUninit)
+      CheckFailed("'allockind()' can't be both zeroed and uninitialized");
+  }
+
   if (Attrs.hasFnAttr(Attribute::VScaleRange)) {
     unsigned VScaleMin = Attrs.getFnAttrs().getVScaleRangeMin();
     if (VScaleMin == 0)

llvm/lib/Transforms/Utils/CodeExtractor.cpp

@@ -920,6 +920,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
       case Attribute::StackAlignment:
       case Attribute::WillReturn:
       case Attribute::WriteOnly:
+      case Attribute::AllocKind:
         continue;
       // Those attributes should be safe to propagate to the extracted function.
       case Attribute::AlwaysInline:

llvm/test/Assembler/allockind-missing.ll

@@ -0,0 +1,4 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+declare void @f0() allockind()
+; CHECK: :[[#@LINE-1]]:30: error: expected allockind value

llvm/test/Assembler/allockind.ll

@@ -0,0 +1,7 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+declare void @f0() allockind("free")
+declare void @f1() allockind("alloc,aligned,uninitialized")
+declare void @f2() allockind("realloc,zeroed,aligned")
+declare void @f3() allockind("fjord")
+; CHECK: :[[#@LINE-1]]:30: error: unknown allockind fjord

llvm/test/Bitcode/compatibility.ll

@@ -1520,7 +1520,7 @@ exit:
   ; CHECK: select <2 x i1> <i1 true, i1 false>, <2 x i8> <i8 2, i8 3>, <2 x i8> <i8 3, i8 2>
 
   call void @f.nobuiltin() builtin
-  ; CHECK: call void @f.nobuiltin() #49
+  ; CHECK: call void @f.nobuiltin() #50
 
   call fastcc noalias i32* @f.noalias() noinline
   ; CHECK: call fastcc noalias i32* @f.noalias() #12
@@ -1943,6 +1943,9 @@ declare void @f.allocsize_two(i32, i32) allocsize(1, 0)
 declare void @f.nosanitize_bounds() nosanitize_bounds
 ; CHECK: declare void @f.nosanitize_bounds() #48
 
+declare void @f.allockind() allockind("alloc,uninitialized")
+; CHECK: declare void @f.allockind() #49
+
 ; CHECK: attributes #0 = { alignstack=4 }
 ; CHECK: attributes #1 = { alignstack=8 }
 ; CHECK: attributes #2 = { alwaysinline }
@@ -1992,7 +1995,8 @@ declare void @f.nosanitize_bounds() nosanitize_bounds
 ; CHECK: attributes #46 = { allocsize(0) }
 ; CHECK: attributes #47 = { allocsize(1,0) }
 ; CHECK: attributes #48 = { nosanitize_bounds }
-; CHECK: attributes #49 = { builtin }
+; CHECK: attributes #49 = { allockind("alloc,uninitialized") }
+; CHECK: attributes #50 = { builtin }
 
 ;; Metadata
 

llvm/test/Verifier/allockind.ll

@@ -0,0 +1,16 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: 'allockind()' requires exactly one of alloc, realloc, and free
+declare i8* @a(i32) allockind("aligned")
+
+; CHECK: 'allockind()' requires exactly one of alloc, realloc, and free
+declare i8* @b(i32*) allockind("free,realloc")
+
+; CHECK: 'allockind("free")' doesn't allow uninitialized, zeroed, or aligned modifiers.
+declare i8* @c(i32) allockind("free,zeroed")
+
+; CHECK: 'allockind()' can't be both zeroed and uninitialized
+declare i8* @d(i32, i32*) allockind("realloc,uninitialized,zeroed")
+
+; CHECK: 'allockind()' requires exactly one of alloc, realloc, and free
+declare i8* @e(i32, i32) allockind("alloc,free")