Add the allocsize attribute to LLVM

[george.burgess.iv]

Summary

Add the allocsize function-level attribute to LLVM.

This attribute allows users to tag arbitrary functions as "malloc-like", so intrinsics like @llvm.objectsize can give higher quality results without requiring modifications to LLVM itself.

[george.burgess.iv]

Ping :)

[nlopes]

I was wondering whether we could go a bit more generic and support arbitrary expressions (like llvm.assume). This design only supports allocation functions that return a pointer to the beginning of a buffer with size p1 * p2 * ... * pn, where pi are the parameters.

It would be nice to support arbitrary expressions for size and offset. There are many cases in practice where these would be helpful. Of course then we would need to expose a generalized version of gcc's attribute to applications.

[george.burgess.iv]

support arbitrary expressions (like llvm.assume).

I'm fine allowing this attribute to store a Function* or Value* that LLVM tries to constant fold at each callsite, but I can't see how llvm.assume fits into the picture at all. Can you give a concrete example of what you're trying to accomplish with the addition you proposed?

[nlopes]

In D14933#311655, @george.burgess.iv wrote:

support arbitrary expressions (like llvm.assume).

I'm fine allowing this attribute to store a Function* or Value* that LLVM tries to constant fold at each callsite, but I can't see how llvm.assume fits into the picture at all. Can you give a concrete example of what you're trying to accomplish with the addition you proposed?

I mentioned assume as a precedent for something that takes an arbitrary expression which is not codegen'ed.
See e.g.: http://lists.llvm.org/pipermail/cfe-dev/2012-June/022272.html

[george.burgess.iv]

My stab at allowing LLVM to support arbitrary expressions in allocsize is available at D16390. It needs some refactoring/cleanup, but it gives you an idea of what I think we need to do to support allocsize with arbitrary expressions well in LLVM. If we end up _really_ wanting to support that, I'll fix the other patch up a bit and submit it for actual-review. Personally, I'm sketched out by the things that we need to do in that patch, so I think I'd rather continue with a review on this one.

If the consensus is that the other patch is worth the modifications that it brings, I'll clean it up + submit it for actual-review.

[george.burgess.iv]

• Rebased
• Added tests for overflow
• Added support in the bitcode reader/writer + tests for it.

[george.burgess.iv]

Ping :)

[george.burgess.iv]

Ping :)

[ab]

D16390 is indeed quite repulsive; what do you think of annotating the callsites instead?

It does increase bitcode size, but perhaps something like this would feel more natural:

define i8* @checked_malloc(i32 %N) {
	%1 = call i8* @object_malloc(i32 %N)
	%2 = call i32 @llvm.objectsize.i32.p0i8(i8* %1)
	%max = add i32 %N, 8
	%c = icmp ule %max, %2
	call void @llvm.assume(i1 %c)
	%3 = icmp lt %2, %N
	br i1 %3, label %Trap, label %Return
	...
}

You could also introduce dedicated intrinsics (and maybe even wrap the calls with them, to make the lookup trivial).

[george.burgess.iv]

Let me modify your code sample a bit, so it represents what the transformation would really look like if we wanted to support arbitrary expressions.

 define i8* @checked_malloc(i32 %N) {
 	%1 = call i8* @object_malloc(i32 %N)
 	%2 = call {i32, i32} @llvm.objectsize.and.object.offset(i8* %1)
	%3 = call {i32, i32} @__autogenerated_get_object_size_of_object_malloc(i32 %N)
	%c = icmp eq %2, %3
	call void @llvm.assume(i1 %c)
	%4 = extractelement {i32, i32} %2, 0
	%5 = icmp lt %4, %N
	br i1 %3, label %Trap, label %Return
	...
}

So let's walk through what needs to happen here:

• If we don't run a special piece of code (ObjectSizeOffsetEvaluator), then the call to @__autogenerated_get_object_size_of_object_malloc must go away without a trace.
• If we do run ObjectSizeOffsetEvaluator on %1, @__autogenerated_get_object_size_of_object_malloc must stay, as it will be used

Because of this, we either have to get rid of @__autogenerated_get_object_size_of_object_malloc before inlining, or we need to mark it noinline so it's not impossible to remove from functions. The former case means we lose our allocsize *really* early, and the latter means we probably get slower code (especially if the function amounts to return {N - 8, 8};)

Also, we'd need a highly restricted user function constant folder to make this work, just like in D16390.

Plus, @llvm.assume says on the tin that it makes optimizing things harder, so this makes optimizing any function that calls an allocation function harder. Ouch. Okay, let's try something else.

define i8* @checked_malloc(i32 %N) {
  %1 = call i8* @object_malloc(i32 %N)
  call void @llvm.assert_object_size_is(i8* %1, {i32, i32}(i32)* @__autogenerated_get_object_size_of_object_malloc, i32 %N)
  ...
}

This seems more reasonable. Here are the problems it poses, as I see it:

• It still makes computing @llvm.objectsize more complex (instead of chasing a pointer through bitcasts/GEPs/... to its definition, you need to check at every GEP/bitcast/... for uses of the pointer that are calls to this intrinsic.
• We still need a special user function constant folder.
• We still need to rely on optimizations happening to @__autogenerated_get_object_size_of_object_malloc in order to have a chance of constant folding it.
• We waste time optimizing @__autogenerated_get_object_size_of_object_malloc, a function that may never even hit an object file.
• We would need to magically turn the intrinsic into an actual function call in some cases. If this happens late, the function won't get inlined. Like said above, if the function amounts to return {N - 8, 8};, that's not good.
• @__autogenerated_get_object_size_of_object_malloc isn't allowed to have any side-effects (because the intrinsic would presumably be treated as though it has no side-effects, much unlike a call instruction).
• We need to keep around anything passed to @__autogenerated_get_object_size_of_object_malloc until the intrinsic is deleted, so it may suffer from similar optimization problems as @llvm.assume.
• We still need to emit this at every callsite to any allocation function.

...And this is the best alternate design I can think of. Does this match what you were thinking? If so, I'm not at all convinced that it's strictly better than what I'm presenting here, which prevents no optimizations, is very lightweight, doesn't require special handling everywhere, doesn't need a custom constant folder, etc.

---

Honestly, and this is me speaking from a position of complete ignorance as to what the actual real-life use case of this would be, in what case are tiny wrapper functions unacceptable? The only things I have to pull from are here: http://lists.llvm.org/pipermail/cfe-dev/2012-June/022272.html -- can you tell me why the below alternatives don't work just as well? Or give me an example where it wouldn't be possible to write a similar wrapper? I'm unable to think of one.

// Original
char *my_strdup(char *str) __attribute__((alloc_size_ex(strlen(str)+1)));

// Modified
extern char *_my_strdup_impl(const char *str, int N) __attribute__((alloc_size(1)));
inline char *my_strdup(const char *str) {
  size_t total_len = strlen(str)+1;
  return _my_strdup_impl(str, total_len);
}

// ------------------

// Original
void *my_complex_alloc(int n, int size, int add) __attribute__((alloc_size_ex(n * size + add)));

// Modified -- _impl returns the start of the buffer
extern void *_my_complex_alloc_impl(int total_bytes, int n, int size, int add) __attribute__((alloc_size(0)));
inline void *my_complex_alloc(int n, int size, int add) {
  return _my_complex_alloc_impl(n * size + add, n, size, add);
}

// ------------------
// Original
char *middle(int size) __attribute__((alloc_size_ex(size, size/2)));

// Modified - Again, _middle_impl returns the start of the buffer
extern char *_middle_impl(int size) __attribute__((alloc_size(0)));
inline char *middle(int size) { return _middle_impl(size) + size/2; }

[sanjoy]

Hi George,

Have you considered using operand bundles to express arbitrary
allocation size information? With operand bundles, the representation
would be something like

define i8* @checked_malloc(i32 %N) {
      %alloc_size = add i32 %N, 8
      %1 = call i8* @object_malloc(i32 %N) [ "alloc-size"(i32 %alloc_size) ]
      %2 = call i32 @llvm.objectsize.i32.p0i8(i8* %1)
      %3 = icmp lt %2, %N
      br i1 %3, label %Trap, label %Return
      ...
}

and you'd fold (objectsize (call fn(...) [ "alloc-size"(N) ])) to (N).
Once you've decided that there are likely no more opportunities to
simplify llvm.objectsize calls, you can drop the "alloc-size" operand
bundles from all calls (this can possibly be folded into CGP or some
other pass, to avoid re-traversing the whole module).

The caveat, of course, is that in this scheme you won't be able to
fold (objectsize (call allocation_fn)) where the call to allocation_fn
was initially an indirect call that got devirtualized, since you need
to know the call sites that will call allocation functions when
emitting the calls.

• Sanjoy

[ab]

Before going further:

Honestly, and this is me speaking from a position of complete ignorance as to what the actual real-life use case of this would be, in what case are tiny wrapper functions unacceptable? The only things I have to pull from are here: http://lists.llvm.org/pipermail/cfe-dev/2012-June/022272.html -- can you tell me why the below alternatives don't work just as well? Or give me an example where it wouldn't be possible to write a similar wrapper? I'm unable to think of one.

FWIW, I agree. One could justify the expression support as simplifying the transition for users (the alternative you show does require manual work), but I think this attribute is enough.

That said, the expression support (D16390) is an interesting thought experiment (Disclaimer: just a curious bystander). So, I was thinking you'd IRGen the whole expression at every callsite; and this is what I meant with "wrapping" the calls:

define i8* @checked_malloc(i32 %N) {
	%p = call i8* @object_malloc(i32 %N)
	%size = add i32 %N, 8
	%1 = call i8* @llvm.assume_alloc_size_and_offset.p0i8(i8* %p, i32 %size, i32 -8)
	%2 = call i32 @llvm.objectsize.i32.p0i8(i8* %1)
	%3 = icmp lt %2, %N
	br i1 %3, label %Trap, label %Return
	...
}

I think this avoids all problems except:

• @__autogenerated_get_object_size_of_object_malloc isn't allowed to have any side-effects (because the intrinsic would presumably be treated as though it has no side-effects, much unlike a call instruction).
• We need to keep around anything passed to @__autogenerated_get_object_size_of_object_malloc until the intrinsic is deleted, so it may suffer from similar optimization problems as @llvm.assume.
• We still need to emit this at every callsite to any allocation function.

Also, +1 to Sanjoy's operand bundles approach: I think it's equivalent but cleaner. (Same caveat applies.)

[george.burgess.iv]

Thank you both for the feedback/thoughts.

@sanjoy -

Have you considered using operand bundles to express arbitrary allocation size information? With operand bundles, the representation would be something like

TIL those exist -- neat! I think that would be a better approach than requiring custom intrinsics in the example I posted above.

@ab -

One could justify the expression support as simplifying the transition for users (the alternative you show does require manual work),

I agree, but I'm not sure how strong of an argument that would be, given that in most cases, the difference seems to be writing a few expressions in an attribute vs writing those same expressions in a small wrapper.

but I think this attribute is enough

Yay :)

So, I was thinking you'd IRGen the whole expression at every callsite; and this is what I meant with "wrapping" the calls

Ah, okay. I ruled that out initially because of the "side-effects must be allowed" restriction, but didn't consider it when we lifted that. Apologies.

I think this avoids all problems except [...]

Agreed

---

Thought experiments aside, ISTM that no one's unhappy with the design of this attribute as-is. Is this correct?

[mehdi_amini]

Just some minor comment I wrote a few weeks ago, but didn't get to the end of the review (and won't have time in the short term).

include/llvm/Analysis/MemoryBuiltins.h

170 ↗

(On Diff #45604)

It seems unrelated to the rest of this patch, can you commit now?

include/llvm/IR/Attributes.h

154 ↗

(On Diff #45604)

s/form/for/

Also no \brief anywhere unless you want multiple sentences as part of the brief (the code that has \brief is legacy).

lib/Analysis/MemoryBuiltins.cpp

487 ↗

(On Diff #45604)

This is not clear to me, is this if we overflow the 32 bits that the pointer can represent?
Is it consider UB or...?

[george.burgess.iv]

Addressed all feedback -- thank you!

include/llvm/Analysis/MemoryBuiltins.h

170 ↗

(On Diff #45604)

r261144

include/llvm/IR/Attributes.h

154 ↗

(On Diff #45604)

Is there something that I should be using instead of \brief, or should I just remove it entirely?

lib/Analysis/MemoryBuiltins.cpp

487 ↗

(On Diff #45604)

This is a remnant of two things:

• Trying to handle overflow sanely for llvm.objectsize (the user can call 32 and 64 bit variants of objectsize)
• Trying to handle sizes larger than IntTyBits properly

...The first part of that has been moved to InstCombineCalls, so this can mostly go away. The interesting bit is when the user gives us e.g. void *special_malloc(uint64_t N) on 32-bit platforms. Other places in the code expect the APInts you hand back to be IntTyBits bits wide, so we need to shrink the uint64_t to be 32 bits.

An argument can be made for doing a checked truncation of N or for doing an unchecked truncation. Given that checking if N requires more than 32 bits is cheap, I think having well-defined behavior (read: failing) in this case is harmless.

I rewrote this bit of code so it's hopefully more clear (and added a test to be sure we get this right). :)

[george.burgess.iv]

Rebased and split.

I split this patch in hopes of making it a bit easier to review. This review now just adds the allocsize attribute. With this patch applied alone, allocsize can sit on functions and look pretty, but is useless; the review for the split out bits (coming in ~5 minutes ;) ) makes llvm.objectsize actually use the allocsize attribute.

I intend to commit both patches together.

[george.burgess.iv]

test/Verifier/allocsize.ll

16 ↗

(On Diff #53285)

I know the tests look tiny; please note that the other patch has quite a few tests that come with it. :)

[mehdi_amini]

LGTM, but please see inline comment.

docs/LangRef.rst

1253 ↗

(On Diff #53285)

The two parameters thing is not "super clean", but I can see how it models the reality (i.e. malloc vs calloc).

Maybe specify it as: allocsize(<EltSize>[, <NumElts = 1>])

1255 ↗

(On Diff #53285)

Is "base-0" the right term? (I read it like "base-2" for binary, or "base-16" for hex).
If there is better it would be nice.

include/llvm/IR/Attributes.h

154 ↗

(On Diff #53285)

You should use brief only if you want multiple sentences to be part of the brief, otherwise we have "autobrief" turned on: http://lists.llvm.org/pipermail/llvm-dev/2015-May/085973.html

279 ↗

(On Diff #53285)

Is brief intended?

337 ↗

(On Diff #53285)

No brief

504 ↗

(On Diff #53285)

brief, and below.

[george.burgess.iv]

Addressed all feedback.

docs/LangRef.rst

1253 ↗

(On Diff #53285)

The two parameters thing is not "super clean"

Agreed.

Maybe specify it as: allocsize(<EltSize>[, <NumElts = 1>])

I like the idea of making the names more descriptive, but I'm not sure that saying = 1 is the right thing here. Both of these are parameter numbers, so I feel that putting NumElts = 1 may be a bit misleading. :)

1255 ↗

(On Diff #53285)

Good catch; I'll try "zero-indexed" instead.

include/llvm/IR/Attributes.h

154 ↗

(On Diff #53285)

Good to know -- thanks!

[mehdi_amini]

docs/LangRef.rst

1253 ↗

(On Diff #53338)

I was trying to express the default value with the "= 1", but I'm fine without.

llvm/trunk/docs/LangRef.rst

.. code-block:: llvm
     define void @f() alwaysinline optsize { ... }
     define void @f() optsize { ... }
 
 ``alignstack(<n>)``
     This attribute indicates that, when emitting the prologue and
     epilogue, the backend should forcibly align the stack pointer.
     Specify the desired alignment, which must be a power of two, in
     parentheses.
+``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
+    parameter numbers; if one argument is provided, then it's assumed that at
+    least ``CallSite.Args[EltSizeParam]`` bytes will be available at the
+    returned pointer. If two are provided, then it's assumed that
+    ``CallSite.Args[EltSizeParam] * CallSite.Args[NumEltsParam]`` bytes are
+    available. The referenced parameters must be integer types. No assumptions
+    are made about the contents of the returned block of memory.
 ``alwaysinline``
     This attribute indicates that the inliner should attempt to inline
     this function into callers whenever possible, ignoring any active
     inlining size threshold for this caller.
 ``builtin``
     This indicates that the callee function at a call site should be
     recognized as a built-in function, even though the function's declaration
     uses the ``nobuiltin`` attribute. This is only valid at call sites for

llvm/trunk/include/llvm/Bitcode/LLVMBitCodes.h

enum AttributeKindCodes {
   ATTR_KIND_DEREFERENCEABLE_OR_NULL = 42,
   ATTR_KIND_CONVERGENT = 43,
   ATTR_KIND_SAFESTACK = 44,
   ATTR_KIND_ARGMEMONLY = 45,
   ATTR_KIND_SWIFT_SELF = 46,
   ATTR_KIND_SWIFT_ERROR = 47,
   ATTR_KIND_NO_RECURSE = 48,
   ATTR_KIND_INACCESSIBLEMEM_ONLY = 49,
-  ATTR_KIND_INACCESSIBLEMEM_OR_ARGMEMONLY = 50
+  ATTR_KIND_INACCESSIBLEMEM_OR_ARGMEMONLY = 50,
+  ATTR_KIND_ALLOC_SIZE = 51
 };
 
 enum ComdatSelectionKindCodes {
   COMDAT_SELECTION_KIND_ANY = 1,
   COMDAT_SELECTION_KIND_EXACT_MATCH = 2,
   COMDAT_SELECTION_KIND_LARGEST = 3,
   COMDAT_SELECTION_KIND_NO_DUPLICATES = 4,
   COMDAT_SELECTION_KIND_SAME_SIZE = 5,
 };
 
 } // End bitc namespace
 } // End llvm namespace
 
 #endif

llvm/trunk/include/llvm/IR/Attributes.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_IR_ATTRIBUTES_H
 #define LLVM_IR_ATTRIBUTES_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include <bitset>
 #include <cassert>
 #include <map>
 #include <string>
 
 namespace llvm {
public:
   /// \brief Return a uniquified Attribute object that has the specific
   /// alignment set.
   static Attribute getWithAlignment(LLVMContext &Context, uint64_t Align);
   static Attribute getWithStackAlignment(LLVMContext &Context, uint64_t Align);
   static Attribute getWithDereferenceableBytes(LLVMContext &Context,
                                               uint64_t Bytes);
   static Attribute getWithDereferenceableOrNullBytes(LLVMContext &Context,
                                                      uint64_t Bytes);
+  static Attribute getWithAllocSizeArgs(LLVMContext &Context,
+                                        unsigned ElemSizeArg,
+                                        const Optional<unsigned> &NumElemsArg);
 
   //===--------------------------------------------------------------------===//
   // Attribute Accessors
   //===--------------------------------------------------------------------===//
 
   /// \brief Return true if the attribute is an Attribute::AttrKind type.
   bool isEnumAttribute() const;
 
public:
   /// \brief Returns the number of dereferenceable bytes from the
   /// dereferenceable attribute.
   uint64_t getDereferenceableBytes() const;
 
   /// \brief Returns the number of dereferenceable_or_null bytes from the
   /// dereferenceable_or_null attribute.
   uint64_t getDereferenceableOrNullBytes() const;
 
+  /// Returns the argument numbers for the allocsize attribute (or pair(0, 0)
+  /// if not known).
+  std::pair<unsigned, Optional<unsigned>> getAllocSizeArgs() const;
+
   /// \brief 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;
 
   /// \brief Equality and non-equality operators.
   bool operator==(Attribute A) const { return pImpl == A.pImpl; }
   bool operator!=(Attribute A) const { return pImpl != A.pImpl; }
 
AttributeSet addDereferenceableAttr(LLVMContext &C, unsigned Index,
                                       uint64_t Bytes) const;
 
   /// \brief Add the dereferenceable_or_null attribute to the attribute set at
   /// the given index. Because attribute sets are immutable, this returns a new
   /// set.
   AttributeSet addDereferenceableOrNullAttr(LLVMContext &C, unsigned Index,
                                             uint64_t Bytes) const;
 
+  /// Add the allocsize attribute to the attribute set at the given index.
+  /// Because attribute sets are immutable, this returns a new set.
+  AttributeSet addAllocSizeAttr(LLVMContext &C, unsigned Index,
+                                unsigned ElemSizeArg,
+                                const Optional<unsigned> &NumElemsArg);
+
   //===--------------------------------------------------------------------===//
   // AttributeSet Accessors
   //===--------------------------------------------------------------------===//
 
   /// \brief Retrieve the LLVM context.
   LLVMContext &getContext() const;
 
   /// \brief The attributes for the specified index are returned.
public:
 
   /// \brief Get the number of dereferenceable bytes (or zero if unknown).
   uint64_t getDereferenceableBytes(unsigned Index) const;
 
   /// \brief Get the number of dereferenceable_or_null bytes (or zero if
   /// unknown).
   uint64_t getDereferenceableOrNullBytes(unsigned Index) const;
 
+  /// Get the allocsize argument numbers (or pair(0, 0) if unknown).
+  std::pair<unsigned, Optional<unsigned>>
+  getAllocSizeArgs(unsigned Index) const;
+
   /// \brief Return the attributes at the index as a string.
   std::string getAsString(unsigned Index, bool InAttrGrp = false) const;
 
   typedef ArrayRef<Attribute>::iterator iterator;
 
   iterator begin(unsigned Slot) const;
   iterator end(unsigned Slot) const;
 

 /// equality, presence of attributes, etc.
 class AttrBuilder {
   std::bitset<Attribute::EndAttrKinds> Attrs;
   std::map<std::string, std::string> TargetDepAttrs;
   uint64_t Alignment;
   uint64_t StackAlignment;
   uint64_t DerefBytes;
   uint64_t DerefOrNullBytes;
+  uint64_t AllocSizeArgs;
 
 public:
   AttrBuilder()
       : Attrs(0), Alignment(0), StackAlignment(0), DerefBytes(0),
-        DerefOrNullBytes(0) {}
+        DerefOrNullBytes(0), AllocSizeArgs(0) {}
   explicit AttrBuilder(uint64_t Val)
       : Attrs(0), Alignment(0), StackAlignment(0), DerefBytes(0),
-        DerefOrNullBytes(0) {
+        DerefOrNullBytes(0), AllocSizeArgs(0) {
     addRawValue(Val);
   }
   AttrBuilder(const Attribute &A)
       : Attrs(0), Alignment(0), StackAlignment(0), DerefBytes(0),
-        DerefOrNullBytes(0) {
+        DerefOrNullBytes(0), AllocSizeArgs(0) {
     addAttribute(A);
   }
   AttrBuilder(AttributeSet AS, unsigned Idx);
 
   void clear();
 
   /// \brief Add an attribute to the builder.
   AttrBuilder &addAttribute(Attribute::AttrKind Val);
public:
   /// \brief Retrieve the number of dereferenceable bytes, if the
   /// dereferenceable attribute exists (zero is returned otherwise).
   uint64_t getDereferenceableBytes() const { return DerefBytes; }
 
   /// \brief Retrieve the number of dereferenceable_or_null bytes, if the
   /// dereferenceable_or_null attribute exists (zero is returned otherwise).
   uint64_t getDereferenceableOrNullBytes() const { return DerefOrNullBytes; }
 
+  /// Retrieve the allocsize args, if the allocsize attribute exists.  If it
+  /// doesn't exist, pair(0, 0) is returned.
+  std::pair<unsigned, Optional<unsigned>> getAllocSizeArgs() const;
+
   /// \brief This turns an int alignment (which must be a power of 2) into the
   /// form used internally in Attribute.
   AttrBuilder &addAlignmentAttr(unsigned Align);
 
   /// \brief This turns an int stack alignment (which must be a power of 2) into
   /// the form used internally in Attribute.
   AttrBuilder &addStackAlignmentAttr(unsigned Align);
 
   /// \brief This turns the number of dereferenceable bytes into the form used
   /// internally in Attribute.
   AttrBuilder &addDereferenceableAttr(uint64_t Bytes);
 
   /// \brief This turns the number of dereferenceable_or_null bytes into the
   /// form used internally in Attribute.
   AttrBuilder &addDereferenceableOrNullAttr(uint64_t Bytes);
 
+  /// This turns one (or two) ints into the form used internally in Attribute.
+  AttrBuilder &addAllocSizeAttr(unsigned ElemSizeArg,
+                                const Optional<unsigned> &NumElemsArg);
+
+  /// Add an allocsize attribute, using the representation returned by
+  /// Attribute.getIntValue().
+  AttrBuilder &addAllocSizeAttrFromRawRepr(uint64_t RawAllocSizeRepr);
+
   /// \brief Return true if the builder contains no target-independent
   /// attributes.
   bool empty() const { return Attrs.none(); }
 
   // Iterators for target-dependent attributes.
   typedef std::pair<std::string, std::string>                td_type;
   typedef std::map<std::string, std::string>::iterator       td_iterator;
   typedef std::map<std::string, std::string>::const_iterator td_const_iterator;

llvm/trunk/include/llvm/IR/Attributes.td

 class StrBoolAttr<string S> : Attr<S>;
 
 /// Target-independent enum attributes.
 
 /// Alignment of parameter (5 bits) stored as log2 of alignment with +1 bias.
 /// 0 means unaligned (different from align(1)).
 def Alignment : EnumAttr<"align">;
 
+/// The result of the function is guaranteed to point to a number of bytes that
+/// we can determine if we know the value of the function's arguments.
+def AllocSize : EnumAttr<"allocsize">;
+
 /// inline=always.
 def AlwaysInline : EnumAttr<"alwaysinline">;
 
 /// Function can access memory only using pointers based on its arguments.
 def ArgMemOnly : EnumAttr<"argmemonly">;
 
 /// Callee is recognized as a builtin, despite nobuiltin attribute on its
 /// declaration.

llvm/trunk/lib/Analysis/MemoryBuiltins.cpp

enum AllocType : uint8_t {
   CallocLike         = 1<<2, // allocates + bzero
   ReallocLike        = 1<<3, // reallocates
   StrDupLike         = 1<<4,
   AllocLike          = MallocLike | CallocLike | StrDupLike,
   AnyAlloc           = AllocLike | ReallocLike
 };
 
 struct AllocFnsTy {
-  LibFunc::Func Func;
   AllocType AllocTy;
-  unsigned char NumParams;
+  unsigned NumParams;
   // First and Second size parameters (or -1 if unused)
-  signed char FstParam, SndParam;
+  int FstParam, SndParam;
 };
 
 // FIXME: certain users need more information. E.g., SimplifyLibCalls needs to
 // know which functions are nounwind, noalias, nocapture parameters, etc.
-static const AllocFnsTy AllocationFnData[] = {
-  {LibFunc::malloc,              MallocLike,  1, 0,  -1},
-  {LibFunc::valloc,              MallocLike,  1, 0,  -1},
-  {LibFunc::Znwj,                OpNewLike,   1, 0,  -1}, // new(unsigned int)
-  {LibFunc::ZnwjRKSt9nothrow_t,  MallocLike,  2, 0,  -1}, // new(unsigned int, nothrow)
-  {LibFunc::Znwm,                OpNewLike,   1, 0,  -1}, // new(unsigned long)
-  {LibFunc::ZnwmRKSt9nothrow_t,  MallocLike,  2, 0,  -1}, // new(unsigned long, nothrow)
-  {LibFunc::Znaj,                OpNewLike,   1, 0,  -1}, // new[](unsigned int)
-  {LibFunc::ZnajRKSt9nothrow_t,  MallocLike,  2, 0,  -1}, // new[](unsigned int, nothrow)
-  {LibFunc::Znam,                OpNewLike,   1, 0,  -1}, // new[](unsigned long)
-  {LibFunc::ZnamRKSt9nothrow_t,  MallocLike,  2, 0,  -1}, // new[](unsigned long, nothrow)
-  {LibFunc::msvc_new_int,         OpNewLike,   1, 0,  -1}, // new(unsigned int)
-  {LibFunc::msvc_new_int_nothrow, MallocLike,  2, 0,  -1}, // new(unsigned int, nothrow)
-  {LibFunc::msvc_new_longlong,         OpNewLike,   1, 0,  -1}, // new(unsigned long long)
-  {LibFunc::msvc_new_longlong_nothrow, MallocLike,  2, 0,  -1}, // new(unsigned long long, nothrow)
-  {LibFunc::msvc_new_array_int,         OpNewLike,   1, 0,  -1}, // new[](unsigned int)
-  {LibFunc::msvc_new_array_int_nothrow, MallocLike,  2, 0,  -1}, // new[](unsigned int, nothrow)
-  {LibFunc::msvc_new_array_longlong,         OpNewLike,   1, 0,  -1}, // new[](unsigned long long)
-  {LibFunc::msvc_new_array_longlong_nothrow, MallocLike,  2, 0,  -1}, // new[](unsigned long long, nothrow)
-  {LibFunc::calloc,              CallocLike,  2, 0,   1},
-  {LibFunc::realloc,             ReallocLike, 2, 1,  -1},
-  {LibFunc::reallocf,            ReallocLike, 2, 1,  -1},
-  {LibFunc::strdup,              StrDupLike,  1, -1, -1},
-  {LibFunc::strndup,             StrDupLike,  2, 1,  -1}
+static const std::pair<LibFunc::Func, AllocFnsTy> AllocationFnData[] = {
+  {LibFunc::malloc,              {MallocLike,  1, 0,  -1}},
+  {LibFunc::valloc,              {MallocLike,  1, 0,  -1}},
+  {LibFunc::Znwj,                {OpNewLike,   1, 0,  -1}}, // new(unsigned int)
+  {LibFunc::ZnwjRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new(unsigned int, nothrow)
+  {LibFunc::Znwm,                {OpNewLike,   1, 0,  -1}}, // new(unsigned long)
+  {LibFunc::ZnwmRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new(unsigned long, nothrow)
+  {LibFunc::Znaj,                {OpNewLike,   1, 0,  -1}}, // new[](unsigned int)
+  {LibFunc::ZnajRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new[](unsigned int, nothrow)
+  {LibFunc::Znam,                {OpNewLike,   1, 0,  -1}}, // new[](unsigned long)
+  {LibFunc::ZnamRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new[](unsigned long, nothrow)
+  {LibFunc::msvc_new_int,         {OpNewLike,   1, 0,  -1}}, // new(unsigned int)
+  {LibFunc::msvc_new_int_nothrow, {MallocLike,  2, 0,  -1}}, // new(unsigned int, nothrow)
+  {LibFunc::msvc_new_longlong,         {OpNewLike,   1, 0,  -1}}, // new(unsigned long long)
+  {LibFunc::msvc_new_longlong_nothrow, {MallocLike,  2, 0,  -1}}, // new(unsigned long long, nothrow)
+  {LibFunc::msvc_new_array_int,         {OpNewLike,   1, 0,  -1}}, // new[](unsigned int)
+  {LibFunc::msvc_new_array_int_nothrow, {MallocLike,  2, 0,  -1}}, // new[](unsigned int, nothrow)
+  {LibFunc::msvc_new_array_longlong,         {OpNewLike,   1, 0,  -1}}, // new[](unsigned long long)
+  {LibFunc::msvc_new_array_longlong_nothrow, {MallocLike,  2, 0,  -1}}, // new[](unsigned long long, nothrow)
+  {LibFunc::calloc,              {CallocLike,  2, 0,   1}},
+  {LibFunc::realloc,             {ReallocLike, 2, 1,  -1}},
+  {LibFunc::reallocf,            {ReallocLike, 2, 1,  -1}},
+  {LibFunc::strdup,              {StrDupLike,  1, -1, -1}},
+  {LibFunc::strndup,             {StrDupLike,  2, 1,  -1}}
   // TODO: Handle "int posix_memalign(void **, size_t, size_t)"
 };
 
 
 static Function *getCalledFunction(const Value *V, bool LookThroughBitCast) {
   if (LookThroughBitCast)
     V = V->stripPointerCasts();
 
   CallSite CS(const_cast<Value*>(V));
   if (!CS.getInstruction())
     return nullptr;
 
   if (CS.isNoBuiltin())
     return nullptr;
 
   Function *Callee = CS.getCalledFunction();
   if (!Callee || !Callee->isDeclaration())
     return nullptr;
   return Callee;
 }
 
-/// \brief Returns the allocation data for the given value if it is a call to a
-/// known allocation function, and NULL otherwise.
-static const AllocFnsTy *getAllocationData(const Value *V, AllocType AllocTy,
+/// Returns the allocation data for the given value if it's either a call to a
+/// known allocation function, or a call to a function with the allocsize
+/// attribute.
+static Optional<AllocFnsTy> getAllocationData(const Value *V, AllocType AllocTy,
                                               const TargetLibraryInfo *TLI,
                                               bool LookThroughBitCast = false) {
   // Skip intrinsics
   if (isa<IntrinsicInst>(V))
-    return nullptr;
+    return None;
 
-  Function *Callee = getCalledFunction(V, LookThroughBitCast);
+  const Function *Callee = getCalledFunction(V, LookThroughBitCast);
   if (!Callee)
-    return nullptr;
+    return None;
+
+  // If it has allocsize, we can skip checking if it's a known function.
+  //
+  // MallocLike is chosen here because allocsize makes no guarantees about the
+  // nullness of the result of the function, nor does it deal with strings, nor
+  // does it require that the memory returned is zeroed out.
+  LLVM_CONSTEXPR auto AllocSizeAllocTy = MallocLike;
+  if ((AllocTy & AllocSizeAllocTy) == AllocSizeAllocTy &&
+      Callee->hasFnAttribute(Attribute::AllocSize)) {
+    Attribute Attr = Callee->getFnAttribute(Attribute::AllocSize);
+    std::pair<unsigned, Optional<unsigned>> Args = Attr.getAllocSizeArgs();
+
+    AllocFnsTy Result;
+    Result.AllocTy = AllocSizeAllocTy;
+    Result.NumParams = Callee->getNumOperands();
+    Result.FstParam = Args.first;
+    Result.SndParam = Args.second.getValueOr(-1);
+    return Result;
+  }
 
   // Make sure that the function is available.
   StringRef FnName = Callee->getName();
   LibFunc::Func TLIFn;
   if (!TLI || !TLI->getLibFunc(FnName, TLIFn) || !TLI->has(TLIFn))
-    return nullptr;
+    return None;
 
-  const AllocFnsTy *FnData =
+  const auto *Iter =
       std::find_if(std::begin(AllocationFnData), std::end(AllocationFnData),
-                   [TLIFn](const AllocFnsTy &Fn) { return Fn.Func == TLIFn; });
+                   [TLIFn](const std::pair<LibFunc::Func, AllocFnsTy> &P) {
+                     return P.first == TLIFn;
+                   });
 
-  if (FnData == std::end(AllocationFnData))
-    return nullptr;
+  if (Iter == std::end(AllocationFnData))
+    return None;
 
+  const AllocFnsTy *FnData = &Iter->second;
   if ((FnData->AllocTy & AllocTy) != FnData->AllocTy)
-    return nullptr;
+    return None;
 
   // Check function prototype.
   int FstParam = FnData->FstParam;
   int SndParam = FnData->SndParam;
   FunctionType *FTy = Callee->getFunctionType();
 
   if (FTy->getReturnType() == Type::getInt8PtrTy(FTy->getContext()) &&
       FTy->getNumParams() == FnData->NumParams &&
       (FstParam < 0 ||
        (FTy->getParamType(FstParam)->isIntegerTy(32) ||
         FTy->getParamType(FstParam)->isIntegerTy(64))) &&
       (SndParam < 0 ||
        FTy->getParamType(SndParam)->isIntegerTy(32) ||
        FTy->getParamType(SndParam)->isIntegerTy(64)))
-    return FnData;
-  return nullptr;
+    return *FnData;
+  return None;
 }
 
 static bool hasNoAliasAttr(const Value *V, bool LookThroughBitCast) {
   ImmutableCallSite CS(LookThroughBitCast ? V->stripPointerCasts() : V);
   return CS && CS.paramHasAttr(AttributeSet::ReturnIndex, Attribute::NoAlias);
 }
 
 
 /// \brief Tests if a value is a call or invoke to a library function that
 /// allocates or reallocates memory (either malloc, calloc, realloc, or strdup
 /// like).
 bool llvm::isAllocationFn(const Value *V, const TargetLibraryInfo *TLI,
                           bool LookThroughBitCast) {
-  return getAllocationData(V, AnyAlloc, TLI, LookThroughBitCast);
+  return getAllocationData(V, AnyAlloc, TLI, LookThroughBitCast).hasValue();
 }
 
 /// \brief Tests if a value is a call or invoke to a function that returns a
 /// NoAlias pointer (including malloc/calloc/realloc/strdup-like functions).
 bool llvm::isNoAliasFn(const Value *V, const TargetLibraryInfo *TLI,
                        bool LookThroughBitCast) {
   // it's safe to consider realloc as noalias since accessing the original
   // pointer is undefined behavior
   return isAllocationFn(V, TLI, LookThroughBitCast) ||
          hasNoAliasAttr(V, LookThroughBitCast);
 }
 
 /// \brief Tests if a value is a call or invoke to a library function that
 /// allocates uninitialized memory (such as malloc).
 bool llvm::isMallocLikeFn(const Value *V, const TargetLibraryInfo *TLI,
                           bool LookThroughBitCast) {
-  return getAllocationData(V, MallocLike, TLI, LookThroughBitCast);
+  return getAllocationData(V, MallocLike, TLI, LookThroughBitCast).hasValue();
 }
 
 /// \brief Tests if a value is a call or invoke to a library function that
 /// allocates zero-filled memory (such as calloc).
 bool llvm::isCallocLikeFn(const Value *V, const TargetLibraryInfo *TLI,
                           bool LookThroughBitCast) {
-  return getAllocationData(V, CallocLike, TLI, LookThroughBitCast);
+  return getAllocationData(V, CallocLike, TLI, LookThroughBitCast).hasValue();
 }
 
 /// \brief Tests if a value is a call or invoke to a library function that
 /// allocates memory (either malloc, calloc, or strdup like).
 bool llvm::isAllocLikeFn(const Value *V, const TargetLibraryInfo *TLI,
                          bool LookThroughBitCast) {
-  return getAllocationData(V, AllocLike, TLI, LookThroughBitCast);
+  return getAllocationData(V, AllocLike, TLI, LookThroughBitCast).hasValue();
 }
 
 /// extractMallocCall - Returns the corresponding CallInst if the instruction
 /// is a malloc call.  Since CallInst::CreateMalloc() only creates calls, we
 /// ignore InvokeInst here.
 const CallInst *llvm::extractMallocCall(const Value *I,
                                         const TargetLibraryInfo *TLI) {
   return isMallocLikeFn(I, TLI) ? dyn_cast<CallInst>(I) : nullptr;
if (!A.hasByValOrInAllocaAttr()) {
     return unknown();
   }
   PointerType *PT = cast<PointerType>(A.getType());
   APInt Size(IntTyBits, DL.getTypeAllocSize(PT->getElementType()));
   return std::make_pair(align(Size, A.getParamAlignment()), Zero);
 }
 
 SizeOffsetType ObjectSizeOffsetVisitor::visitCallSite(CallSite CS) {
-  const AllocFnsTy *FnData = getAllocationData(CS.getInstruction(), AnyAlloc,
-                                               TLI);
+  Optional<AllocFnsTy> FnData =
+      getAllocationData(CS.getInstruction(), AnyAlloc, TLI);
   if (!FnData)
     return unknown();
 
   // handle strdup-like functions separately
   if (FnData->AllocTy == StrDupLike) {
     APInt Size(IntTyBits, GetStringLength(CS.getArgument(0)));
     if (!Size)
       return unknown();
 
     // strndup limits strlen
     if (FnData->FstParam > 0) {
-      ConstantInt *Arg= dyn_cast<ConstantInt>(CS.getArgument(FnData->FstParam));
+      ConstantInt *Arg =
+          dyn_cast<ConstantInt>(CS.getArgument(FnData->FstParam));
       if (!Arg)
         return unknown();
 
       APInt MaxSize = Arg->getValue().zextOrSelf(IntTyBits);
       if (Size.ugt(MaxSize))
         Size = MaxSize + 1;
     }
     return std::make_pair(Size, Zero);
   }
 
   ConstantInt *Arg = dyn_cast<ConstantInt>(CS.getArgument(FnData->FstParam));
   if (!Arg)
     return unknown();
 
-  APInt Size = Arg->getValue().zextOrSelf(IntTyBits);
+  // When we're compiling N-bit code, and the user uses parameters that are
+  // greater than N bits (e.g. uint64_t on a 32-bit build), we can run into
+  // trouble with APInt size issues. This function handles resizing + overflow
+  // checks for us.
+  auto CheckedZextOrTrunc = [&](APInt &I) {
+    // More bits than we can handle. Checking the bit width isn't necessary, but
+    // it's faster than checking active bits, and should give `false` in the
+    // vast majority of cases.
+    if (I.getBitWidth() > IntTyBits && I.getActiveBits() > IntTyBits)
+      return false;
+    if (I.getBitWidth() != IntTyBits)
+      I = I.zextOrTrunc(IntTyBits);
+    return true;
+  };
+
+  APInt Size = Arg->getValue();
+  if (!CheckedZextOrTrunc(Size))
+    return unknown();
+
   // size determined by just 1 parameter
   if (FnData->SndParam < 0)
     return std::make_pair(Size, Zero);
 
   Arg = dyn_cast<ConstantInt>(CS.getArgument(FnData->SndParam));
   if (!Arg)
     return unknown();
 
-  Size *= Arg->getValue().zextOrSelf(IntTyBits);
-  return std::make_pair(Size, Zero);
+  APInt NumElems = Arg->getValue();
+  if (!CheckedZextOrTrunc(NumElems))
+    return unknown();
+
+  bool Overflow;
+  Size = Size.umul_ov(NumElems, Overflow);
+  return Overflow ? unknown() : std::make_pair(Size, Zero);
 
   // TODO: handle more standard functions (+ wchar cousins):
   // - strdup / strndup
   // - strcpy / strncpy
   // - strcat / strncat
   // - memcpy / memmove
   // - strcat / strncat
   // - memset
SizeOffsetEvalType ObjectSizeOffsetEvaluator::visitAllocaInst(AllocaInst &I) {
   Value *ArraySize = I.getArraySize();
   Value *Size = ConstantInt::get(ArraySize->getType(),
                                  DL.getTypeAllocSize(I.getAllocatedType()));
   Size = Builder.CreateMul(Size, ArraySize);
   return std::make_pair(Size, Zero);
 }
 
 SizeOffsetEvalType ObjectSizeOffsetEvaluator::visitCallSite(CallSite CS) {
-  const AllocFnsTy *FnData = getAllocationData(CS.getInstruction(), AnyAlloc,
-                                               TLI);
+  Optional<AllocFnsTy> FnData =
+      getAllocationData(CS.getInstruction(), AnyAlloc, TLI);
   if (!FnData)
     return unknown();
 
   // handle strdup-like functions separately
   if (FnData->AllocTy == StrDupLike) {
     // TODO
     return unknown();
   }

llvm/trunk/lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(amdgpu_cs);
 
   KEYWORD(cc);
   KEYWORD(c);
 
   KEYWORD(attributes);
 
   KEYWORD(alwaysinline);
+  KEYWORD(allocsize);
   KEYWORD(argmemonly);
   KEYWORD(builtin);
   KEYWORD(byval);
   KEYWORD(inalloca);
   KEYWORD(cold);
   KEYWORD(convergent);
   KEYWORD(dereferenceable);
   KEYWORD(dereferenceable_or_null);

llvm/trunk/lib/AsmParser/LLParser.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_ASMPARSER_LLPARSER_H
 #define LLVM_LIB_ASMPARSER_LLPARSER_H
 
 #include "LLLexer.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/ValueHandle.h"
 #include <map>
private:
     bool ParseOptionalAlignment(unsigned &Alignment);
     bool ParseOptionalDerefAttrBytes(lltok::Kind AttrKind, uint64_t &Bytes);
     bool ParseScopeAndOrdering(bool isAtomic, SynchronizationScope &Scope,
                                AtomicOrdering &Ordering);
     bool ParseOrdering(AtomicOrdering &Ordering);
     bool ParseOptionalStackAlignment(unsigned &Alignment);
     bool ParseOptionalCommaAlign(unsigned &Alignment, bool &AteExtraComma);
     bool ParseOptionalCommaInAlloca(bool &IsInAlloca);
-    bool ParseIndexList(SmallVectorImpl<unsigned> &Indices,bool &AteExtraComma);
+    bool parseAllocSizeArguments(unsigned &ElemSizeArg,
+                                 Optional<unsigned> &HowManyArg);
+    bool ParseIndexList(SmallVectorImpl<unsigned> &Indices,
+                        bool &AteExtraComma);
     bool ParseIndexList(SmallVectorImpl<unsigned> &Indices) {
       bool AteExtraComma;
       if (ParseIndexList(Indices, AteExtraComma)) return true;
       if (AteExtraComma)
         return TokError("expected index");
       return false;
     }
 

llvm/trunk/lib/AsmParser/LLParser.cpp

 //===-- LLParser.cpp - Parser Class ---------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 //  This file defines the parser class for .ll files.
 //
 //===----------------------------------------------------------------------===//
 
 #include "LLParser.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/AsmParser/SlotMapping.h"
 #include "llvm/IR/AutoUpgrade.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/InlineAsm.h"
case lltok::kw_alignstack: {
           return true;
       } else {
         if (ParseOptionalStackAlignment(Alignment))
           return true;
       }
       B.addStackAlignmentAttr(Alignment);
       continue;
     }
+    case lltok::kw_allocsize: {
+      unsigned ElemSizeArg;
+      Optional<unsigned> NumElemsArg;
+      // inAttrGrp doesn't matter; we only support allocsize(a[, b])
+      if (parseAllocSizeArguments(ElemSizeArg, NumElemsArg))
+        return true;
+      B.addAllocSizeAttr(ElemSizeArg, NumElemsArg);
+      continue;
+    }
     case lltok::kw_alwaysinline: B.addAttribute(Attribute::AlwaysInline); break;
     case lltok::kw_argmemonly: B.addAttribute(Attribute::ArgMemOnly); break;
     case lltok::kw_builtin: B.addAttribute(Attribute::Builtin); break;
     case lltok::kw_cold: B.addAttribute(Attribute::Cold); break;
     case lltok::kw_convergent: B.addAttribute(Attribute::Convergent); break;
     case lltok::kw_inaccessiblememonly:
       B.addAttribute(Attribute::InaccessibleMemOnly); break;
     case lltok::kw_inaccessiblemem_or_argmemonly:
if (Lex.getKind() != lltok::kw_align)
       return Error(Lex.getLoc(), "expected metadata or 'align'");
 
     if (ParseOptionalAlignment(Alignment)) return true;
   }
 
   return false;
 }
 
+bool LLParser::parseAllocSizeArguments(unsigned &BaseSizeArg,
+                                       Optional<unsigned> &HowManyArg) {
+  Lex.Lex();
+
+  auto StartParen = Lex.getLoc();
+  if (!EatIfPresent(lltok::lparen))
+    return Error(StartParen, "expected '('");
+
+  if (ParseUInt32(BaseSizeArg))
+    return true;
+
+  if (EatIfPresent(lltok::comma)) {
+    auto HowManyAt = Lex.getLoc();
+    unsigned HowMany;
+    if (ParseUInt32(HowMany))
+      return true;
+    if (HowMany == BaseSizeArg)
+      return Error(HowManyAt,
+                   "'allocsize' indices can't refer to the same parameter");
+    HowManyArg = HowMany;
+  } else
+    HowManyArg = None;
+
+  auto EndParen = Lex.getLoc();
+  if (!EatIfPresent(lltok::rparen))
+    return Error(EndParen, "expected ')'");
+  return false;
+}
+
 /// ParseScopeAndOrdering
 ///   if isAtomic: ::= 'singlethread'? AtomicOrdering
 ///   else: ::=
 ///
 /// This sets Scope and Ordering to the parsed values.
 bool LLParser::ParseScopeAndOrdering(bool isAtomic, SynchronizationScope &Scope,
                                      AtomicOrdering &Ordering) {
   if (!isAtomic)

llvm/trunk/lib/AsmParser/LLToken.h

enum Kind {
     kw_cxx_fast_tlscc,
     kw_amdgpu_vs,
     kw_amdgpu_gs,
     kw_amdgpu_ps,
     kw_amdgpu_cs,
 
     // Attributes:
     kw_attributes,
+    kw_allocsize,
     kw_alwaysinline,
     kw_argmemonly,
     kw_sanitize_address,
     kw_builtin,
     kw_byval,
     kw_inalloca,
     kw_cold,
     kw_convergent,

llvm/trunk/lib/Bitcode/Reader/BitcodeReader.cpp

static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
   case bitc::ATTR_KIND_NON_LAZY_BIND:
     return Attribute::NonLazyBind;
   case bitc::ATTR_KIND_NON_NULL:
     return Attribute::NonNull;
   case bitc::ATTR_KIND_DEREFERENCEABLE:
     return Attribute::Dereferenceable;
   case bitc::ATTR_KIND_DEREFERENCEABLE_OR_NULL:
     return Attribute::DereferenceableOrNull;
+  case bitc::ATTR_KIND_ALLOC_SIZE:
+    return Attribute::AllocSize;
   case bitc::ATTR_KIND_NO_RED_ZONE:
     return Attribute::NoRedZone;
   case bitc::ATTR_KIND_NO_RETURN:
     return Attribute::NoReturn;
   case bitc::ATTR_KIND_NO_UNWIND:
     return Attribute::NoUnwind;
   case bitc::ATTR_KIND_OPTIMIZE_FOR_SIZE:
     return Attribute::OptimizeForSize;
case bitc::PARAMATTR_GRP_CODE_ENTRY: { // ENTRY: [grpid, idx, a0, a1, ...]
           if (Kind == Attribute::Alignment)
             B.addAlignmentAttr(Record[++i]);
           else if (Kind == Attribute::StackAlignment)
             B.addStackAlignmentAttr(Record[++i]);
           else if (Kind == Attribute::Dereferenceable)
             B.addDereferenceableAttr(Record[++i]);
           else if (Kind == Attribute::DereferenceableOrNull)
             B.addDereferenceableOrNullAttr(Record[++i]);
+          else if (Kind == Attribute::AllocSize)
+            B.addAllocSizeAttrFromRawRepr(Record[++i]);
         } else {                     // String attribute
           assert((Record[i] == 3 || Record[i] == 4) &&
                  "Invalid attribute group entry");
           bool HasValue = (Record[i++] == 4);
           SmallString<64> KindStr;
           SmallString<64> ValStr;
 
           while (Record[i] != 0 && i != e)

llvm/trunk/lib/Bitcode/Writer/BitcodeWriter.cpp

static void WriteStringRecord(unsigned Code, StringRef Str,
   // Emit the finished record.
   Stream.EmitRecord(Code, Vals, AbbrevToUse);
 }
 
 static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
   switch (Kind) {
   case Attribute::Alignment:
     return bitc::ATTR_KIND_ALIGNMENT;
+  case Attribute::AllocSize:
+    return bitc::ATTR_KIND_ALLOC_SIZE;
   case Attribute::AlwaysInline:
     return bitc::ATTR_KIND_ALWAYS_INLINE;
   case Attribute::ArgMemOnly:
     return bitc::ATTR_KIND_ARGMEMONLY;
   case Attribute::Builtin:
     return bitc::ATTR_KIND_BUILTIN;
   case Attribute::ByVal:
     return bitc::ATTR_KIND_BY_VAL;

llvm/trunk/lib/IR/AttributeImpl.h

 /// LLVMContextImpl for creating and managing attributes.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_IR_ATTRIBUTEIMPL_H
 #define LLVM_LIB_IR_ATTRIBUTEIMPL_H
 
 #include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/TrailingObjects.h"
 #include <climits>
 #include <string>
 
 namespace llvm {
 
class IntAttributeImpl : public EnumAttributeImpl {
   void anchor() override;
   uint64_t Val;
 
 public:
   IntAttributeImpl(Attribute::AttrKind Kind, uint64_t Val)
       : EnumAttributeImpl(IntAttrEntry, Kind), Val(Val) {
     assert((Kind == Attribute::Alignment || Kind == Attribute::StackAlignment ||
             Kind == Attribute::Dereferenceable ||
-            Kind == Attribute::DereferenceableOrNull) &&
+            Kind == Attribute::DereferenceableOrNull ||
+            Kind == Attribute::AllocSize) &&
            "Wrong kind for int attribute!");
   }
 
   uint64_t getValue() const { return Val; }
 };
 
 class StringAttributeImpl : public AttributeImpl {
   virtual void anchor();
public:
 
   Attribute getAttribute(Attribute::AttrKind Kind) const;
   Attribute getAttribute(StringRef Kind) const;
 
   unsigned getAlignment() const;
   unsigned getStackAlignment() const;
   uint64_t getDereferenceableBytes() const;
   uint64_t getDereferenceableOrNullBytes() const;
+  std::pair<unsigned, Optional<unsigned>> getAllocSizeArgs() const;
   std::string getAsString(bool InAttrGrp) const;
 
   typedef const Attribute *iterator;
   iterator begin() const { return getTrailingObjects<Attribute>(); }
   iterator end() const { return begin() + NumAttrs; }
 
   void Profile(FoldingSetNodeID &ID) const {
     Profile(ID, makeArrayRef(begin(), end()));

llvm/trunk/lib/IR/Attributes.cpp

 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
 // Attribute Construction Methods
 //===----------------------------------------------------------------------===//
 
+// allocsize has two integer arguments, but because they're both 32 bits, we can
+// pack them into one 64-bit value, at the cost of making said value
+// nonsensical.
+//
+// In order to do this, we need to reserve one value of the second (optional)
+// allocsize argument to signify "not present."
+LLVM_CONSTEXPR static unsigned AllocSizeNumElemsNotPresent =
+    std::numeric_limits<unsigned>::max();
+
+static uint64_t packAllocSizeArgs(unsigned ElemSizeArg,
+                                  const Optional<unsigned> &NumElemsArg) {
+  assert((!NumElemsArg.hasValue() ||
+          *NumElemsArg != AllocSizeNumElemsNotPresent) &&
+         "Attempting to pack a reserved value");
+
+  return uint64_t(ElemSizeArg) << 32 |
+         NumElemsArg.getValueOr(AllocSizeNumElemsNotPresent);
+}
+
+static std::pair<unsigned, Optional<unsigned>>
+unpackAllocSizeArgs(uint64_t Num) {
+  unsigned NumElems = Num & std::numeric_limits<unsigned>::max();
+  unsigned ElemSizeArg = Num >> 32;
+
+  Optional<unsigned> NumElemsArg;
+  if (NumElems != AllocSizeNumElemsNotPresent)
+    NumElemsArg = NumElems;
+  return std::make_pair(ElemSizeArg, NumElemsArg);
+}
+
 Attribute Attribute::get(LLVMContext &Context, Attribute::AttrKind Kind,
                          uint64_t Val) {
   LLVMContextImpl *pImpl = Context.pImpl;
   FoldingSetNodeID ID;
   ID.AddInteger(Kind);
   if (Val) ID.AddInteger(Val);
 
   void *InsertPoint;

 }
 
 Attribute Attribute::getWithDereferenceableOrNullBytes(LLVMContext &Context,
                                                        uint64_t Bytes) {
   assert(Bytes && "Bytes must be non-zero.");
   return get(Context, DereferenceableOrNull, Bytes);
 }
 
+Attribute
+Attribute::getWithAllocSizeArgs(LLVMContext &Context, unsigned ElemSizeArg,
+                                const Optional<unsigned> &NumElemsArg) {
+  assert(!(ElemSizeArg == 0 && NumElemsArg && *NumElemsArg == 0) &&
+         "Invalid allocsize arguments -- given allocsize(0, 0)");
+  return get(Context, AllocSize, packAllocSizeArgs(ElemSizeArg, NumElemsArg));
+}
+
 //===----------------------------------------------------------------------===//
 // Attribute Accessor Methods
 //===----------------------------------------------------------------------===//
 
 bool Attribute::isEnumAttribute() const {
   return pImpl && pImpl->isEnumAttribute();
 }
 

 
 uint64_t Attribute::getDereferenceableOrNullBytes() const {
   assert(hasAttribute(Attribute::DereferenceableOrNull) &&
          "Trying to get dereferenceable bytes from "
          "non-dereferenceable attribute!");
   return pImpl->getValueAsInt();
 }
 
+std::pair<unsigned, Optional<unsigned>> Attribute::getAllocSizeArgs() const {
+  assert(hasAttribute(Attribute::AllocSize) &&
+         "Trying to get allocsize args from non-allocsize attribute");
+  return unpackAllocSizeArgs(pImpl->getValueAsInt());
+}
+
 std::string Attribute::getAsString(bool InAttrGrp) const {
   if (!pImpl) return "";
 
   if (hasAttribute(Attribute::SanitizeAddress))
     return "sanitize_address";
   if (hasAttribute(Attribute::AlwaysInline))
     return "alwaysinline";
   if (hasAttribute(Attribute::ArgMemOnly))
if (hasAttribute(Attribute::StackAlignment))
     return AttrWithBytesToString("alignstack");
 
   if (hasAttribute(Attribute::Dereferenceable))
     return AttrWithBytesToString("dereferenceable");
 
   if (hasAttribute(Attribute::DereferenceableOrNull))
     return AttrWithBytesToString("dereferenceable_or_null");
 
+  if (hasAttribute(Attribute::AllocSize)) {
+    unsigned ElemSize;
+    Optional<unsigned> NumElems;
+    std::tie(ElemSize, NumElems) = getAllocSizeArgs();
+
+    std::string Result = "allocsize(";
+    Result += utostr(ElemSize);
+    if (NumElems.hasValue()) {
+      Result += ',';
+      Result += utostr(*NumElems);
+    }
+    Result += ')';
+    return Result;
+  }
+
   // Convert target-dependent attributes to strings of the form:
   //
   //   "kind"
   //   "kind" = "value"
   //
   if (isStringAttribute()) {
     std::string Result;
     Result += (Twine('"') + getKindAsString() + Twine('"')).str();
case Attribute::Dereferenceable:
     break;
   case Attribute::DereferenceableOrNull:
     llvm_unreachable("dereferenceable_or_null attribute not supported in raw "
                      "format");
     break;
   case Attribute::ArgMemOnly:
     llvm_unreachable("argmemonly attribute not supported in raw format");
     break;
+  case Attribute::AllocSize:
+    llvm_unreachable("allocsize not supported in raw format");
+    break;
   }
   llvm_unreachable("Unsupported attribute type");
 }
 
 //===----------------------------------------------------------------------===//
 // AttributeSetNode Definition
 //===----------------------------------------------------------------------===//
 

 
 uint64_t AttributeSetNode::getDereferenceableOrNullBytes() const {
   for (iterator I = begin(), E = end(); I != E; ++I)
     if (I->hasAttribute(Attribute::DereferenceableOrNull))
       return I->getDereferenceableOrNullBytes();
   return 0;
 }
 
+std::pair<unsigned, Optional<unsigned>>
+AttributeSetNode::getAllocSizeArgs() const {
+  for (iterator I = begin(), E = end(); I != E; ++I)
+    if (I->hasAttribute(Attribute::AllocSize))
+      return I->getAllocSizeArgs();
+  return std::make_pair(0, 0);
+}
+
 std::string AttributeSetNode::getAsString(bool InAttrGrp) const {
   std::string Str;
   for (iterator I = begin(), E = end(); I != E; ++I) {
     if (I != begin())
       Str += ' ';
     Str += I->getAsString(InAttrGrp);
   }
   return Str;
for (AttributeSetNode::iterator II = ASN->begin(),
       Attribute::AttrKind Kind = Attr.getKindAsEnum();
 
       if (Kind == Attribute::Alignment)
         Mask |= (Log2_32(ASN->getAlignment()) + 1) << 16;
       else if (Kind == Attribute::StackAlignment)
         Mask |= (Log2_32(ASN->getStackAlignment()) + 1) << 26;
       else if (Kind == Attribute::Dereferenceable)
         llvm_unreachable("dereferenceable not supported in bit mask");
+      else if (Kind == Attribute::AllocSize)
+        llvm_unreachable("allocsize not supported in bit mask");
       else
         Mask |= AttributeImpl::getAttrMask(Kind);
     }
 
     return Mask;
   }
 
   return 0;
for (Attribute::AttrKind Kind = Attribute::None;
     case Attribute::Dereferenceable:
       Attr = Attribute::getWithDereferenceableBytes(
           C, B.getDereferenceableBytes());
       break;
     case Attribute::DereferenceableOrNull:
       Attr = Attribute::getWithDereferenceableOrNullBytes(
           C, B.getDereferenceableOrNullBytes());
       break;
+    case Attribute::AllocSize: {
+      auto A = B.getAllocSizeArgs();
+      Attr = Attribute::getWithAllocSizeArgs(C, A.first, A.second);
+      break;
+    }
     default:
       Attr = Attribute::get(C, Kind);
     }
     Attrs.push_back(std::make_pair(Index, Attr));
   }
 
   // Add target-dependent (string) attributes.
   for (const AttrBuilder::td_type &TDA : B.td_attrs())

 AttributeSet AttributeSet::addDereferenceableOrNullAttr(LLVMContext &C,
                                                         unsigned Index,
                                                         uint64_t Bytes) const {
   llvm::AttrBuilder B;
   B.addDereferenceableOrNullAttr(Bytes);
   return addAttributes(C, Index, AttributeSet::get(C, Index, B));
 }
 
+AttributeSet
+AttributeSet::addAllocSizeAttr(LLVMContext &C, unsigned Index,
+                               unsigned ElemSizeArg,
+                               const Optional<unsigned> &NumElemsArg) {
+  llvm::AttrBuilder B;
+  B.addAllocSizeAttr(ElemSizeArg, NumElemsArg);
+  return addAttributes(C, Index, AttributeSet::get(C, Index, B));
+}
+
 //===----------------------------------------------------------------------===//
 // AttributeSet Accessor Methods
 //===----------------------------------------------------------------------===//
 
 LLVMContext &AttributeSet::getContext() const {
   return pImpl->getContext();
 }
 
uint64_t AttributeSet::getDereferenceableBytes(unsigned Index) const {
   return ASN ? ASN->getDereferenceableBytes() : 0;
 }
 
 uint64_t AttributeSet::getDereferenceableOrNullBytes(unsigned Index) const {
   AttributeSetNode *ASN = getAttributes(Index);
   return ASN ? ASN->getDereferenceableOrNullBytes() : 0;
 }
 
-std::string AttributeSet::getAsString(unsigned Index,
-                                      bool InAttrGrp) const {
+std::pair<unsigned, Optional<unsigned>>
+AttributeSet::getAllocSizeArgs(unsigned Index) const {
+  AttributeSetNode *ASN = getAttributes(Index);
+  return ASN ? ASN->getAllocSizeArgs() : std::make_pair(0, 0);
+}
+
+std::string AttributeSet::getAsString(unsigned Index, bool InAttrGrp) const {
   AttributeSetNode *ASN = getAttributes(Index);
   return ASN ? ASN->getAsString(InAttrGrp) : std::string("");
 }
 
 AttributeSetNode *AttributeSet::getAttributes(unsigned Index) const {
   if (!pImpl) return nullptr;
 
   // Loop through to find the attribute node we want.

 }
 
 //===----------------------------------------------------------------------===//
 // AttrBuilder Method Implementations
 //===----------------------------------------------------------------------===//
 
 AttrBuilder::AttrBuilder(AttributeSet AS, unsigned Index)
     : Attrs(0), Alignment(0), StackAlignment(0), DerefBytes(0),
-      DerefOrNullBytes(0) {
+      DerefOrNullBytes(0), AllocSizeArgs(0) {
   AttributeSetImpl *pImpl = AS.pImpl;
   if (!pImpl) return;
 
   for (unsigned I = 0, E = pImpl->getNumAttributes(); I != E; ++I) {
     if (pImpl->getSlotIndex(I) != Index) continue;
 
     for (AttributeSetImpl::iterator II = pImpl->begin(I),
            IE = pImpl->end(I); II != IE; ++II)
       addAttribute(*II);
 
     break;
   }
 }
 
 void AttrBuilder::clear() {
   Attrs.reset();
   TargetDepAttrs.clear();
   Alignment = StackAlignment = DerefBytes = DerefOrNullBytes = 0;
+  AllocSizeArgs = 0;
 }
 
 AttrBuilder &AttrBuilder::addAttribute(Attribute::AttrKind Val) {
   assert((unsigned)Val < Attribute::EndAttrKinds && "Attribute out of range!");
   assert(Val != Attribute::Alignment && Val != Attribute::StackAlignment &&
-         Val != Attribute::Dereferenceable &&
+         Val != Attribute::Dereferenceable && Val != Attribute::AllocSize &&
          "Adding integer attribute without adding a value!");
   Attrs[Val] = true;
   return *this;
 }
 
 AttrBuilder &AttrBuilder::addAttribute(Attribute Attr) {
   if (Attr.isStringAttribute()) {
     addAttribute(Attr.getKindAsString(), Attr.getValueAsString());
     return *this;
   }
 
   Attribute::AttrKind Kind = Attr.getKindAsEnum();
   Attrs[Kind] = true;
 
   if (Kind == Attribute::Alignment)
     Alignment = Attr.getAlignment();
   else if (Kind == Attribute::StackAlignment)
     StackAlignment = Attr.getStackAlignment();
   else if (Kind == Attribute::Dereferenceable)
     DerefBytes = Attr.getDereferenceableBytes();
   else if (Kind == Attribute::DereferenceableOrNull)
     DerefOrNullBytes = Attr.getDereferenceableOrNullBytes();
+  else if (Kind == Attribute::AllocSize)
+    AllocSizeArgs = Attr.getValueAsInt();
   return *this;
 }
 
 AttrBuilder &AttrBuilder::addAttribute(StringRef A, StringRef V) {
   TargetDepAttrs[A] = V;
   return *this;
 }
 
 AttrBuilder &AttrBuilder::removeAttribute(Attribute::AttrKind Val) {
   assert((unsigned)Val < Attribute::EndAttrKinds && "Attribute out of range!");
   Attrs[Val] = false;
 
   if (Val == Attribute::Alignment)
     Alignment = 0;
   else if (Val == Attribute::StackAlignment)
     StackAlignment = 0;
   else if (Val == Attribute::Dereferenceable)
     DerefBytes = 0;
   else if (Val == Attribute::DereferenceableOrNull)
     DerefOrNullBytes = 0;
+  else if (Val == Attribute::AllocSize)
+    AllocSizeArgs = 0;
 
   return *this;
 }
 
 AttrBuilder &AttrBuilder::removeAttributes(AttributeSet A, uint64_t Index) {
   unsigned Slot = ~0U;
   for (unsigned I = 0, E = A.getNumSlots(); I != E; ++I)
     if (A.getSlotIndex(I) == Index) {

 
 AttrBuilder &AttrBuilder::removeAttribute(StringRef A) {
   std::map<std::string, std::string>::iterator I = TargetDepAttrs.find(A);
   if (I != TargetDepAttrs.end())
     TargetDepAttrs.erase(I);
   return *this;
 }
 
+std::pair<unsigned, Optional<unsigned>> AttrBuilder::getAllocSizeArgs() const {
+  return unpackAllocSizeArgs(AllocSizeArgs);
+}
+
 AttrBuilder &AttrBuilder::addAlignmentAttr(unsigned Align) {
   if (Align == 0) return *this;
 
   assert(isPowerOf2_32(Align) && "Alignment must be a power of two.");
   assert(Align <= 0x40000000 && "Alignment too large.");
 
   Attrs[Attribute::Alignment] = true;
   Alignment = Align;
AttrBuilder &AttrBuilder::addDereferenceableOrNullAttr(uint64_t Bytes) {
   if (Bytes == 0)
     return *this;
 
   Attrs[Attribute::DereferenceableOrNull] = true;
   DerefOrNullBytes = Bytes;
   return *this;
 }
 
+AttrBuilder &AttrBuilder::addAllocSizeAttr(unsigned ElemSize,
+                                           const Optional<unsigned> &NumElems) {
+  return addAllocSizeAttrFromRawRepr(packAllocSizeArgs(ElemSize, NumElems));
+}
+
+AttrBuilder &AttrBuilder::addAllocSizeAttrFromRawRepr(uint64_t RawArgs) {
+  // (0, 0) is our "not present" value, so we need to check for it here.
+  assert(RawArgs && "Invalid allocsize arguments -- given allocsize(0, 0)");
+
+  Attrs[Attribute::AllocSize] = true;
+  // Reuse existing machinery to store this as a single 64-bit integer so we can
+  // save a few bytes over using a pair<unsigned, Optional<unsigned>>.
+  AllocSizeArgs = RawArgs;
+  return *this;
+}
+
 AttrBuilder &AttrBuilder::merge(const AttrBuilder &B) {
   // FIXME: What if both have alignments, but they don't match?!
   if (!Alignment)
     Alignment = B.Alignment;
 
   if (!StackAlignment)
     StackAlignment = B.StackAlignment;
 
   if (!DerefBytes)
     DerefBytes = B.DerefBytes;
 
   if (!DerefOrNullBytes)
     DerefOrNullBytes = B.DerefOrNullBytes;
 
+  if (!AllocSizeArgs)
+    AllocSizeArgs = B.AllocSizeArgs;
+
   Attrs |= B.Attrs;
 
   for (auto I : B.td_attrs())
     TargetDepAttrs[I.first] = I.second;
 
   return *this;
 }
 
 AttrBuilder &AttrBuilder::remove(const AttrBuilder &B) {
   // FIXME: What if both have alignments, but they don't match?!
   if (B.Alignment)
     Alignment = 0;
 
   if (B.StackAlignment)
     StackAlignment = 0;
 
   if (B.DerefBytes)
     DerefBytes = 0;
 
   if (B.DerefOrNullBytes)
     DerefOrNullBytes = 0;
 
+  if (B.AllocSizeArgs)
+    AllocSizeArgs = 0;
+
   Attrs &= ~B.Attrs;
 
   for (auto I : B.td_attrs())
     TargetDepAttrs.erase(I.first);
 
   return *this;
 }
 

 AttrBuilder &AttrBuilder::addRawValue(uint64_t Val) {
   // FIXME: Remove this in 4.0.
   if (!Val) return *this;
 
   for (Attribute::AttrKind I = Attribute::None; I != Attribute::EndAttrKinds;
        I = Attribute::AttrKind(I + 1)) {
     if (I == Attribute::Dereferenceable ||
         I == Attribute::DereferenceableOrNull ||
-        I == Attribute::ArgMemOnly)
+        I == Attribute::ArgMemOnly ||
+        I == Attribute::AllocSize)
       continue;
     if (uint64_t A = (Val & AttributeImpl::getAttrMask(I))) {
       Attrs[I] = true;
  
       if (I == Attribute::Alignment)
         Alignment = 1ULL << ((A >> 16) - 1);
       else if (I == Attribute::StackAlignment)
         StackAlignment = 1ULL << ((A >> 26)-1);

llvm/trunk/lib/IR/Verifier.cpp

if (I->getKindAsEnum() == Attribute::NoReturn ||
         I->getKindAsEnum() == Attribute::NoBuiltin ||
         I->getKindAsEnum() == Attribute::Cold ||
         I->getKindAsEnum() == Attribute::OptimizeNone ||
         I->getKindAsEnum() == Attribute::JumpTable ||
         I->getKindAsEnum() == Attribute::Convergent ||
         I->getKindAsEnum() == Attribute::ArgMemOnly ||
         I->getKindAsEnum() == Attribute::NoRecurse ||
         I->getKindAsEnum() == Attribute::InaccessibleMemOnly ||
-        I->getKindAsEnum() == Attribute::InaccessibleMemOrArgMemOnly) {
+        I->getKindAsEnum() == Attribute::InaccessibleMemOrArgMemOnly ||
+        I->getKindAsEnum() == Attribute::AllocSize) {
       if (!isFunction) {
         CheckFailed("Attribute '" + I->getAsString() +
                     "' only applies to functions!", V);
         return;
       }
     } else if (I->getKindAsEnum() == Attribute::ReadOnly ||
                I->getKindAsEnum() == Attribute::ReadNone) {
       if (Idx == 0) {
void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeSet Attrs,
   }
 
   if (Attrs.hasAttribute(AttributeSet::FunctionIndex,
                          Attribute::JumpTable)) {
     const GlobalValue *GV = cast<GlobalValue>(V);
     Assert(GV->hasUnnamedAddr(),
            "Attribute 'jumptable' requires 'unnamed_addr'", V);
   }
+
+  if (Attrs.hasAttribute(AttributeSet::FunctionIndex, Attribute::AllocSize)) {
+    std::pair<unsigned, Optional<unsigned>> Args =
+        Attrs.getAllocSizeArgs(AttributeSet::FunctionIndex);
+
+    auto CheckParam = [&](StringRef Name, unsigned ParamNo) {
+      if (ParamNo >= FT->getNumParams()) {
+        CheckFailed("'allocsize' " + Name + " argument is out of bounds", V);
+        return false;
+      }
+
+      if (!FT->getParamType(ParamNo)->isIntegerTy()) {
+        CheckFailed("'allocsize' " + Name +
+                        " argument must refer to an integer parameter",
+                    V);
+        return false;
+      }
+
+      return true;
+    };
+
+    if (!CheckParam("element size", Args.first))
+      return;
+
+    if (Args.second && !CheckParam("number of elements", *Args.second))
+      return;
+  }
 }
 
 void Verifier::verifyFunctionMetadata(
     const SmallVector<std::pair<unsigned, MDNode *>, 4> MDs) {
   if (MDs.empty())
     return;
 
   for (const auto &Pair : MDs) {

llvm/trunk/lib/Transforms/InstCombine/InstCombineCalls.cpp

auto SimplifyDemandedVectorEltsLow = [this](Value *Op, unsigned Width,
     APInt DemandedElts = APInt::getLowBitsSet(Width, DemandedWidth);
     return SimplifyDemandedVectorElts(Op, DemandedElts, UndefElts);
   };
 
   switch (II->getIntrinsicID()) {
   default: break;
   case Intrinsic::objectsize: {
     uint64_t Size;
-    if (getObjectSize(II->getArgOperand(0), Size, DL, TLI))
-      return replaceInstUsesWith(CI, ConstantInt::get(CI.getType(), Size));
+    if (getObjectSize(II->getArgOperand(0), Size, DL, TLI)) {
+      APInt APSize(II->getType()->getIntegerBitWidth(), Size);
+      // Equality check to be sure that `Size` can fit in a value of type
+      // `II->getType()`
+      if (APSize == Size)
+        return replaceInstUsesWith(CI, ConstantInt::get(II->getType(), APSize));
+    }
     return nullptr;
   }
   case Intrinsic::bswap: {
     Value *IIOperand = II->getArgOperand(0);
     Value *X = nullptr;
 
     // bswap(bswap(x)) -> x
     if (match(IIOperand, m_BSwap(m_Value(X))))

llvm/trunk/test/Bitcode/attributes.ll

 }
 
 declare void @nobuiltin()
 
 define void @f34()
 ; CHECK: define void @f34()
 {
         call void @nobuiltin() nobuiltin
-; CHECK: call void @nobuiltin() #30
+; CHECK: call void @nobuiltin() #32
         ret void;
 }
 
 define void @f35() optnone noinline
 ; CHECK: define void @f35() #23
 {
         ret void;
 }

 define float @f53(i8* %error_ref) {
 entry:
   %error_ptr_ref = alloca swifterror %swift_error*
   store %swift_error* null, %swift_error** %error_ptr_ref
   %call = call float @foo(%swift_error** swifterror %error_ptr_ref)
   ret float 1.0
 }
 
+; CHECK: define i8* @f54(i32) #30
+define i8* @f54(i32) allocsize(0) {
+  ret i8* null
+}
+
+; CHECK: define i8* @f55(i32, i32) #31
+define i8* @f55(i32, i32) allocsize(0, 1) {
+  ret i8* null
+}
+
 ; CHECK: attributes #0 = { noreturn }
 ; CHECK: attributes #1 = { nounwind }
 ; CHECK: attributes #2 = { readnone }
 ; CHECK: attributes #3 = { readonly }
 ; CHECK: attributes #4 = { noinline }
 ; CHECK: attributes #5 = { alwaysinline }
 ; CHECK: attributes #6 = { optsize }
 ; CHECK: attributes #7 = { ssp }

 ; CHECK: attributes #22 = { minsize }
 ; CHECK: attributes #23 = { noinline optnone }
 ; CHECK: attributes #24 = { jumptable }
 ; CHECK: attributes #25 = { convergent }
 ; CHECK: attributes #26 = { argmemonly }
 ; CHECK: attributes #27 = { norecurse }
 ; CHECK: attributes #28 = { inaccessiblememonly }
 ; CHECK: attributes #29 = { inaccessiblemem_or_argmemonly }
-; CHECK: attributes #30 = { nobuiltin }
+; CHECK: attributes #30 = { allocsize(0) }
+; CHECK: attributes #31 = { allocsize(0,1) }
+; CHECK: attributes #32 = { nobuiltin }

llvm/trunk/test/Transforms/InstCombine/allocsize-32.ll

+; RUN: opt < %s -instcombine -S | FileCheck %s
+;
+; The idea is that we want to have sane semantics (e.g. not assertion failures)
+; when given an allocsize function that takes a 64-bit argument in the face of
+; 32-bit pointers.
+
+target datalayout="e-p:32:32:32"
+
+declare i8* @my_malloc(i8*, i64) allocsize(1)
+
+define void @test_malloc(i8** %p, i32* %r) {
+  %1 = call i8* @my_malloc(i8* null, i64 100)
+  store i8* %1, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  %2 = call i32 @llvm.objectsize.i32.p0i8(i8* %1, i1 false)
+  ; CHECK: store i32 100
+  store i32 %2, i32* %r, align 8
+
+  ; Big number is 5 billion.
+  %3 = call i8* @my_malloc(i8* null, i64 5000000000)
+  store i8* %3, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  ; CHECK: call i32 @llvm.objectsize
+  %4 = call i32 @llvm.objectsize.i32.p0i8(i8* %3, i1 false)
+  store i32 %4, i32* %r, align 8
+  ret void
+}
+
+declare i32 @llvm.objectsize.i32.p0i8(i8*, i1)

llvm/trunk/test/Transforms/InstCombine/allocsize.ll

+; RUN: opt < %s -instcombine -S | FileCheck %s
+;
+; Test that instcombine folds allocsize function calls properly.
+; Dummy arguments are inserted to verify that allocsize is picking the right
+; args, and to prove that arbitrary unfoldable values don't interfere with
+; allocsize if they're not used by allocsize.
+
+declare i8* @my_malloc(i8*, i32) allocsize(1)
+declare i8* @my_calloc(i8*, i8*, i32, i32) allocsize(2, 3)
+
+; CHECK-LABEL: define void @test_malloc
+define void @test_malloc(i8** %p, i64* %r) {
+  %1 = call i8* @my_malloc(i8* null, i32 100)
+  store i8* %1, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  %2 = call i64 @llvm.objectsize.i64.p0i8(i8* %1, i1 false)
+  ; CHECK: store i64 100
+  store i64 %2, i64* %r, align 8
+  ret void
+}
+
+; CHECK-LABEL: define void @test_calloc
+define void @test_calloc(i8** %p, i64* %r) {
+  %1 = call i8* @my_calloc(i8* null, i8* null, i32 100, i32 5)
+  store i8* %1, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  %2 = call i64 @llvm.objectsize.i64.p0i8(i8* %1, i1 false)
+  ; CHECK: store i64 500
+  store i64 %2, i64* %r, align 8
+  ret void
+}
+
+; Failure cases with non-constant values...
+; CHECK-LABEL: define void @test_malloc_fails
+define void @test_malloc_fails(i8** %p, i64* %r, i32 %n) {
+  %1 = call i8* @my_malloc(i8* null, i32 %n)
+  store i8* %1, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  ; CHECK: @llvm.objectsize.i64.p0i8
+  %2 = call i64 @llvm.objectsize.i64.p0i8(i8* %1, i1 false)
+  store i64 %2, i64* %r, align 8
+  ret void
+}
+
+; CHECK-LABEL: define void @test_calloc_fails
+define void @test_calloc_fails(i8** %p, i64* %r, i32 %n) {
+  %1 = call i8* @my_calloc(i8* null, i8* null, i32 %n, i32 5)
+  store i8* %1, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  ; CHECK: @llvm.objectsize.i64.p0i8
+  %2 = call i64 @llvm.objectsize.i64.p0i8(i8* %1, i1 false)
+  store i64 %2, i64* %r, align 8
+
+
+  %3 = call i8* @my_calloc(i8* null, i8* null, i32 100, i32 %n)
+  store i8* %3, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  ; CHECK: @llvm.objectsize.i64.p0i8
+  %4 = call i64 @llvm.objectsize.i64.p0i8(i8* %3, i1 false)
+  store i64 %4, i64* %r, align 8
+  ret void
+}
+
+declare i8* @my_malloc_outofline(i8*, i32) #0
+declare i8* @my_calloc_outofline(i8*, i8*, i32, i32) #1
+
+; Verifying that out of line allocsize is parsed correctly
+; CHECK-LABEL: define void @test_outofline
+define void @test_outofline(i8** %p, i64* %r) {
+  %1 = call i8* @my_malloc_outofline(i8* null, i32 100)
+  store i8* %1, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  %2 = call i64 @llvm.objectsize.i64.p0i8(i8* %1, i1 false)
+  ; CHECK: store i64 100
+  store i64 %2, i64* %r, align 8
+
+
+  %3 = call i8* @my_calloc_outofline(i8* null, i8* null, i32 100, i32 5)
+  store i8* %3, i8** %p, align 8 ; To ensure objectsize isn't killed
+
+  %4 = call i64 @llvm.objectsize.i64.p0i8(i8* %3, i1 false)
+  ; CHECK: store i64 500
+  store i64 %4, i64* %r, align 8
+  ret void
+}
+
+declare i8* @my_malloc_i64(i8*, i64) #0
+declare i8* @my_tiny_calloc(i8*, i8*, i8, i8) #1
+declare i8* @my_varied_calloc(i8*, i8*, i32, i8) #1
+
+; CHECK-LABEL: define void @test_overflow
+define void @test_overflow(i8** %p, i32* %r) {
+  %r64 = bitcast i32* %r to i64*
+
+  ; (2**31 + 1) * 2 > 2**31. So overflow. Yay.
+  %big_malloc = call i8* @my_calloc(i8* null, i8* null, i32 2147483649, i32 2)
+  store i8* %big_malloc, i8** %p, align 8
+
+  ; CHECK: @llvm.objectsize
+  %1 = call i32 @llvm.objectsize.i32.p0i8(i8* %big_malloc, i1 false)
+  store i32 %1, i32* %r, align 4
+
+
+  %big_little_malloc = call i8* @my_tiny_calloc(i8* null, i8* null, i8 127, i8 4)
+  store i8* %big_little_malloc, i8** %p, align 8
+
+  ; CHECK: store i32 508
+  %2 = call i32 @llvm.objectsize.i32.p0i8(i8* %big_little_malloc, i1 false)
+  store i32 %2, i32* %r, align 4
+
+
+  ; malloc(2**33)
+  %big_malloc_i64 = call i8* @my_malloc_i64(i8* null, i64 8589934592)
+  store i8* %big_malloc_i64, i8** %p, align 8
+
+  ; CHECK: @llvm.objectsize
+  %3 = call i32 @llvm.objectsize.i32.p0i8(i8* %big_malloc_i64, i1 false)
+  store i32 %3, i32* %r, align 4
+
+
+  %4 = call i64 @llvm.objectsize.i64.p0i8(i8* %big_malloc_i64, i1 false)
+  ; CHECK: store i64 8589934592
+  store i64 %4, i64* %r64, align 8
+
+
+  ; Just intended to ensure that we properly handle args of different types...
+  %varied_calloc = call i8* @my_varied_calloc(i8* null, i8* null, i32 1000, i8 5)
+  store i8* %varied_calloc, i8** %p, align 8
+
+  ; CHECK: store i32 5000
+  %5 = call i32 @llvm.objectsize.i32.p0i8(i8* %varied_calloc, i1 false)
+  store i32 %5, i32* %r, align 4
+
+  ret void
+}
+
+attributes #0 = { allocsize(1) }
+attributes #1 = { allocsize(2, 3) }
+
+declare i32 @llvm.objectsize.i32.p0i8(i8*, i1)
+declare i64 @llvm.objectsize.i64.p0i8(i8*, i1)

llvm/trunk/test/Verifier/alloc-size-failedparse.ll

+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+;
+; We handle allocsize with identical args in the parser, rather than the
+; verifier. So, a seperate test is needed.
+
+; CHECK: 'allocsize' indices can't refer to the same parameter
+declare i8* @a(i32, i32) allocsize(0, 0)

llvm/trunk/test/Verifier/allocsize.ll

+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: 'allocsize' element size argument is out of bounds
+declare i8* @a(i32) allocsize(1)
+
+; CHECK: 'allocsize' element size argument must refer to an integer parameter
+declare i8* @b(i32*) allocsize(0)
+
+; CHECK: 'allocsize' number of elements argument is out of bounds
+declare i8* @c(i32) allocsize(0, 1)
+
+; CHECK: 'allocsize' number of elements argument must refer to an integer parameter
+declare i8* @d(i32, i32*) allocsize(0, 1)
+
+; CHECK: 'allocsize' number of elements argument is out of bounds
+declare i8* @e(i32, i32) allocsize(1, 2)