[Attributor] New attribute to identify what byte ranges are alive for an allocation

[vidsinghal]

Summary

Changes the size of allocations automatically.
Only implements the case when a single range from start of the allocation is alive.

[jdoerfert]

High level: Merge the test in one file. Use descriptive names for functions. remove the printf and other stuff that is not useful.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12310

Manifest is only called if the state is valid at the end. It looks like you duplicate lots of checks here that should not be necessary.

12412–12417

Assertion, this should never be the case, right?

12419–12421

The lambda is called only once, right? Why use a lambda at all?
Also, if the alloca cannot be changed, indicate a pessimistic fixpoint and give up.

[vidsinghal]

[Attributor] Revision for AAAllocationInfo attribute based on Review

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

3551–3553

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

874

Call it numBins or sth, to make it clear. Size could be a lot of things in this context.

12208

no this->

12210

This is just getIRPosition().

12218

Why would we want to have the same value twice, use the pointer.

12220–12222

Style, also elsewhere.

12224

We should not cast things to the Impl. Expose the functions you need in the AAPointerInfo interface instead.

12240

I think A.getDataLayout() is a thing too.

12265–12267

Make this an assertion, manifest will not be called on invalid states.

[vidsinghal]

Fixed some comments.

[vidsinghal]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12224

I think it might make sense to expose the state in the AAPointerInfo struct.
AAAllocation Info might need more information from AAPointerInfo for more complex cases.
But I wanted to get an opinion before doing that.
Right now, State is defined in the AAPointerInfo namespace in llvm/lib/Transforms/IPO/AttributorAttributes.cpp
However, the AAPointerInfo struct does not store it, so that information cannot be accessed without casting to AAPointerInfoImpl.

[vidsinghal]

Edited unit tests

[vidsinghal]

Fix typo in assert.
Add a positive test case for malloc with a TODO.

[vidsinghal]

This comment was removed by vidsinghal.

[jdoerfert]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12220

12228

Don't cast it to an "Impl". There is no need anymore anyway.

12229

State could be anything. At least call it PIState.

12233–12239

This is basically BinSize != 1, return, no?

12252

Why does this not trigger for the ptr/i64 case? Is one of them in bytes the other in bits?

12255

This is fine for OffsetEnd = 2,4,8, but then it becomes a little iffy. Check for non default sizes and go with a char array instead.

12273

Same as above.

llvm/test/Transforms/Attributor/value-simplify-local-remote.ll

82 ↗

(On Diff #551856)

This doesn't actually make it smaller, we should avoid such changes.

135 ↗

(On Diff #551856)

Same here.

[vidsinghal]

Revision based on comments

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12233–12239

umm, I think BinSize == 0

12252

if its a ptr

public: std::optional<TypeSize> getAllocationSize(const DataLayout &DL)

returns a size of 0 which is why that if condition is not triggered.
At the LLVM IR level, I think the size of pointer is just 0.

[vidsinghal]

Fix bugs

[jdoerfert]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12261

Why do you give up on struct type like {i32, i32}?
I don't understand why we need this conditional at all.

[vidsinghal]

Remove un-necessary if condition for struct with only one element.

[vidsinghal]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12261

I was trying to indicate a fix point for structs like
type S = {ptr}
where the struct has only one element in it.
But I think that OffsetEnd == AllocationSize check should work for this now (after fixing the size of pointer issue) I'll remove this check.
Previously this was not working which is why I had to had this check and i think
getNumContainedTypes should have been getStructNumElements instead.

[vidsinghal]

check for zero alloca and indicate fixpoint.

[tianshilei1992]

llvm/lib/Transforms/IPO/Attributor.cpp

3558

Better to use [[maybe_unused]] now.

[jdoerfert]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

12223–12224

As I said before, no casting to implementations of the interface. put the method you need into AAPointerInfo.

12272–12274

llvm/test/Transforms/Attributor/value-simplify-local-remote.ll

384 ↗

(On Diff #552461)

unrelated, just keep the old version of this file.

[vidsinghal]

1.) Insert a new instruction instead of modiying the old one (Char array)
2.) Restore value-simplify-local-remote.ll
3.) Add a new positive test to reduce size of array alloca

[vidsinghal]

1.) Put method to get offsetBins in AAPointerInfo
2.) work on other comments.

[vidsinghal]

add maybe_unused pragma

[vidsinghal]

edits

[jdoerfert]

llvm/include/llvm/Transforms/IPO/Attributor.h

5953

This is not a good idea. What happend to the old API you had?
We do not want to leak out the underlying map if we don't have to.
All you want is a size and begin/end, no?

[vidsinghal]

llvm/include/llvm/Transforms/IPO/Attributor.h

5953

Yeah I see.
The underlying map is protected in the state I think.
Yeah for now, size begin and end should be good.
I think I will try making those virtual pure in AAPointerInfo and then override them in the state?

[vidsinghal]

Add API for begin, end size in AAPointerInfo

[vidsinghal]

Add a positive test with TODO

[vidsinghal]

Make switch to indentify alloca vs malloc instruction

[tianshilei1992]

llvm/lib/Transforms/IPO/Attributor.cpp

3558

Then there is no need (void)Success;.

[vidsinghal]

remove (void)success

[vidsinghal]

Support for modifying malloc instructions
TODO: The original malloc call is not being removed by changeAfterManifest even though it is dead.

[vidsinghal]

edits

[vidsinghal]

fix name variable

[vidsinghal]

edit

[vidsinghal]

Delete malloc call after attributor update

[vidsinghal]

fix example IR

[vidsinghal]

edits

[vidsinghal]

Don't call AAPointerInfoGetOrCreate other than Alloca and Malloc Instruction. Results in infinite loop for instance when calling a ptr

[vidsinghal]

Clean up code

[vidsinghal]

style

[vidsinghal]

Add new test where argument to malloc isn't a constant.
Use AAPotentialConstantValues to get the constant value.

[vidsinghal]

Add null check for IntOperand

[vidsinghal]

New example and question about what the optimization should do for an example like @baz in allocator.ll

[vidsinghal]

Remove run commands in depgraph.ll that were not being parsed

llvm/include/llvm/Transforms/IPO/Attributor.h

 #include "llvm/IR/Value.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/DOTGraphTraits.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ModRef.h"
 #include "llvm/Support/TimeProfiler.h"
+#include "llvm/Support/TypeSize.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 
 #include <limits>
 #include <map>
 #include <optional>
 
 namespace llvm {
struct AAPointerInfo : public AbstractAttribute {
   /// all of them, false otherwise. An access interferes with an offset-size
   /// pair if it might read or write that memory region.
   virtual bool forallInterferingAccesses(
       AA::RangeTy Range, function_ref<bool(const Access &, bool)> CB) const = 0;
 
   /// Call \p CB on all accesses that might interfere with \p I and
   /// return true if all such accesses were known and the callback returned true
   /// for all of them, false otherwise. In contrast to forallInterferingAccesses
   /// this function will perform reasoning to exclude write accesses that cannot
   /// affect the load even if they on the surface look as if they would. The
   /// flag \p HasBeenWrittenTo will be set to true if we know that \p I does not
   /// read the intial value of the underlying memory.
   virtual bool forallInterferingAccesses(
       Attributor &A, const AbstractAttribute &QueryingAA, Instruction &I,
       bool FindInterferingWrites, bool FindInterferingReads,
       function_ref<bool(const Access &, bool)> CB, bool &HasBeenWrittenTo,
       AA::RangeTy &Range) const = 0;
struct AAAddressSpace : public StateWrapper<BooleanState, AbstractAttribute> {
 
   // No address space which indicates the associated value is dead.
   static const int32_t NoAddressSpace = -1;
 
   /// Unique ID (due to the unique address)
   static const char ID;
 };
 
+struct AAAllocationInfo : public StateWrapper<BooleanState, AbstractAttribute> {
+  AAAllocationInfo(const IRPosition &IRP, Attributor &A)
+      : StateWrapper<BooleanState, AbstractAttribute>(IRP) {}
+
+  /// See AbstractAttribute::isValidIRPositionForInit
+  static bool isValidIRPositionForInit(Attributor &A, const IRPosition &IRP) {
+    if (!IRP.getAssociatedType()->isPtrOrPtrVectorTy())
+      return false;
+    return AbstractAttribute::isValidIRPositionForInit(A, IRP);
+  }
+
+  /// Create an abstract attribute view for the position \p IRP.
+  static AAAllocationInfo &createForPosition(const IRPosition &IRP,
+                                             Attributor &A);
+
+  virtual std::optional<TypeSize> getAllocatedSize() const = 0;
+
+  /// See AbstractAttribute::getName()
+  const std::string getName() const override { return "AAAllocationInfo"; }
+
+  /// See AbstractAttribute::getIdAddr()
+  const char *getIdAddr() const override { return &ID; }
+
+  /// This function should return true if the type of the \p AA is
+  /// AAAllocationInfo
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  // No AllocatedSpace which means it doesn't allocated memory.
+  constexpr static const std::optional<TypeSize> NOAllocatedSize =
+      std::optional<TypeSize>(TypeSize(-1, true));
+
+  static const char ID;
+};
+
 raw_ostream &operator<<(raw_ostream &, const AAPointerInfo::Access &);
 
 /// Run options, used by the pass manager.
 enum AttributorRunOption {
   NONE = 0,
   MODULE = 1 << 0,
   CGSCC = 1 << 1,
   ALL = MODULE | CGSCC

llvm/lib/Transforms/IPO/Attributor.cpp

auto LoadStorePred = [&](Instruction &I) -> bool {
     return true;
   };
   Success = checkForAllInstructionsImpl(
       nullptr, OpcodeInstMap, LoadStorePred, nullptr, nullptr,
       {(unsigned)Instruction::Load, (unsigned)Instruction::Store},
       UsedAssumedInformation);
   (void)Success;
   assert(Success && "Expected the check call to be successful!");
+
+  // AllocaInstPredicate
+  auto AllocaInstPred = [&](Instruction &I) -> bool {
+    if (auto *AI = dyn_cast<AllocaInst>(&I)) {
+      getOrCreateAAFor<AAAllocationInfo>(IRPosition::value(*AI));
+    }
+    return true;
+  };
+
+  Success = checkForAllInstructionsImpl(
+      nullptr, OpcodeInstMap, AllocaInstPred, nullptr, nullptr,
+      {(unsigned)Instruction::Alloca}, UsedAssumedInformation);
+  (void)Success;
+  assert(Success && "Expected the check call to be successful!");
 }
 
 /// Helpers to ease debugging through output streams and print calls.
 ///
 ///{
 raw_ostream &llvm::operator<<(raw_ostream &OS, ChangeStatus S) {
   return OS << (S == ChangeStatus::CHANGED ? "changed" : "unchanged");
 }

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/GraphWriter.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/TypeSize.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <cassert>
 #include <numeric>
 #include <optional>
 
 using namespace llvm;

 PIPE_OPERATOR(AANoUndef)
 PIPE_OPERATOR(AANoFPClass)
 PIPE_OPERATOR(AACallEdges)
 PIPE_OPERATOR(AAInterFnReachability)
 PIPE_OPERATOR(AAPointerInfo)
 PIPE_OPERATOR(AAAssumptionInfo)
 PIPE_OPERATOR(AAUnderlyingObjects)
 PIPE_OPERATOR(AAAddressSpace)
+PIPE_OPERATOR(AAAllocationInfo)
 
 #undef PIPE_OPERATOR
 
 template <>
 ChangeStatus clampStateAndIndicateChange<DerefState>(DerefState &S,
                                                      const DerefState &R) {
   ChangeStatus CS0 =
       clampStateAndIndicateChange(S.DerefBytesState, R.DerefBytesState);
ChangeStatus addAccess(Attributor &A, const AAPointerInfo::RangeList &Ranges,
                          AAPointerInfo::AccessKind Kind, Type *Ty,
                          Instruction *RemoteI = nullptr);
 
   using OffsetBinsTy = DenseMap<RangeTy, SmallSet<unsigned, 4>>;
 
   using const_bin_iterator = OffsetBinsTy::const_iterator;
   const_bin_iterator begin() const { return OffsetBins.begin(); }
   const_bin_iterator end() const { return OffsetBins.end(); }
+  int64_t size() const { return OffsetBins.size(); }
 
   const AAPointerInfo::Access &getAccess(unsigned Index) const {
     return AccessList[Index];
   }
 
 protected:
   // Every memory instruction results in an Access object. We maintain a list of
   // all Access objects that we own, along with the following maps:
struct AAAddressSpaceCallSiteArgument final : AAAddressSpaceImpl {
   }
 
   void trackStatistics() const override {
     STATS_DECLTRACK_CSARG_ATTR(addrspace);
   }
 };
 } // namespace
 
+/// ----------- Allocation Info ----------
+namespace {
+struct AAAllocationInfoImpl : public AAAllocationInfo {
+  AAAllocationInfoImpl(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfo(IRP, A) {}
+
+  std::optional<TypeSize> getAllocatedSize() const override {
+    assert(isValidState() && "the AA is invalid");
+    return AssumedAllocatedSize;
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+    auto MakeChange = [&](Instruction &I) {
+      AllocaInst &AI = cast<AllocaInst>(I);
+      const IRPosition &IRP = IRPosition::inst(AI);
+      auto *PI =
+          A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::OPTIONAL);
+
+      if (!PI) {
+        LLVM_DEBUG(dbgs() << "[AAAllocationInfo, updateImpl] could not get "
+                             "AAPointerInfo object"
+                          << AI << "\n");
+        return false;
+      }
+
+      LLVM_DEBUG(dbgs() << "[AAAllocationInfo, updateImpl] Was able to create "
+                           "AAPointer Info object "
+                        << *PI << "\n");
+
+      const AAPointerInfo &OtherAA = *PI;
+
+      if (!OtherAA.getState().isValidState()) {
+        LLVM_DEBUG(dbgs() << "[AAAllocationInfo, updateImpl] AAPointerInfo not "
+                             "in valid state."
+                          << "\n");
+        return false;
+      }
+
+      const auto &OtherAAImpl = static_cast<const AAPointerInfoImpl &>(*PI);
+      const auto &State = OtherAAImpl.getState();
+
+      int BinSize = State.size();
+
+      if (BinSize > 1) {
+        LLVM_DEBUG(dbgs() << "[AAAllocationInfo, updateImpl] AAPointerInfo Bin "
+                             "Size is greater than one. "
+                             "Not supported yet!"
+                          << "\n");
+        return false;
+      }
+
+      const auto &It = State.begin();
+
+      if (It->getFirst().Offset != 0) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, updateImpl] AAPointerInfo Access not "
+                      "starting at 0th byte in Alloca."
+                      "Not supported yet!"
+                   << "\n");
+        return false;
+      }
+
+      LLVM_DEBUG(dbgs() << "[AAAllocationInfo, updateImpl] The first offset is "
+                           "the start of the Alloca."
+                        << "\n");
+
+      int64_t OffsetEnd = It->getFirst().Offset + It->getFirst().Size;
+      const DataLayout &DL = AI.getModule()->getDataLayout();
+      const auto &AllocationSize = AI.getAllocationSize(DL);
+
+      if (!AllocationSize) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, updateImpl] Could not get allocation "
+                      "size from Alloca."
+                   << "\n");
+        return false;
+      }
+
+      if (OffsetEnd == AllocationSize) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, updateImpl] Offset size cannot be "
+                      "reduced in Alloca."
+                   << "\n");
+        return false;
+      }
+
+      Type *IntegerSizeOfOffsetEnd =
+          Type::getIntNTy(AI.getContext(), OffsetEnd * 8);
+
+      LLVM_DEBUG(dbgs() << "[AAAllocationInfo] Successfully reduced allocation "
+                           "size of alloca based on Bin."
+                        << "\n");
+
+      Changed = ChangeStatus::CHANGED;
+
+      changeAllocationSize(IntegerSizeOfOffsetEnd->getPrimitiveSizeInBits());
+
+      return true;
+    };
+
+    auto *I = this->getIRPosition().getCtxI();
+    if (!I) {
+      return ChangeStatus::UNCHANGED;
+    }
+
+    if (!MakeChange(*I)) {
+      return indicatePessimisticFixpoint();
+    }
+
+    return Changed;
+  }
+
+  ChangeStatus Changed = ChangeStatus::UNCHANGED;
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+    auto MakeChange = [&](Instruction &I) {
+      AllocaInst &AI = cast<AllocaInst>(I);
+      const IRPosition &IRP = IRPosition::inst(AI);
+      auto *PI =
+          A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::OPTIONAL);
+
+      if (!PI) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, manifest] Failed to verify all "
+                      "interfering accesses for Instruction "
+                   << AI << "\n");
+        return false;
+      }
+
+      LLVM_DEBUG(dbgs() << "[AAAllocationInfo, manifest] Was able to create "
+                           "AAPointer Info object "
+                        << *PI << "\n");
+
+      const AAPointerInfo &OtherAA = *PI;
+
+      if (!OtherAA.getState().isValidState()) {
+        LLVM_DEBUG(
+            dbgs()
+            << "[AAAllocationInfo, manifest] AAPointerInfo not in valid state."
+            << "\n");
+        return false;
+      }
+
+      const auto &OtherAAImpl = static_cast<const AAPointerInfoImpl &>(*PI);
+      const auto &State = OtherAAImpl.getState();
+
+      int BinSize = State.size();
+
+      if (BinSize > 1) {
+        LLVM_DEBUG(dbgs() << "[AAAllocationInfo, manifest] AAPointerInfo Bin "
+                             "Size is greater than one. "
+                             "Not supported yet!"
+                          << "\n");
+        return false;
+      }
+
+      const auto &It = State.begin();
+
+      if (It->getFirst().Offset != 0) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, manifest] AAPointerInfo Access not "
+                      "starting at 0th byte in Alloca."
+                      "Not supported yet!"
+                   << "\n");
+        return false;
+      }
+
+      LLVM_DEBUG(dbgs() << "[AAAllocationInfo, manifest] The first offset is "
+                           "the start of the Alloca."
+                        << "\n");
+
+      int64_t OffsetEnd = It->getFirst().Offset + It->getFirst().Size;
+      const DataLayout &DL = AI.getModule()->getDataLayout();
+      const auto &AllocationSize = AI.getAllocationSize(DL);
+
+      if (!AllocationSize) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, manifest] Could not get allocation "
+                      "size from Alloca."
+                   << "\n");
+        return false;
+      }
+
+      if (OffsetEnd == AllocationSize) {
+        LLVM_DEBUG(
+            dbgs() << "[AAAllocationInfo, manifest] Offset size cannot be "
+                      "reduced in Alloca."
+                   << "\n");
+        return false;
+      }
+
+      if (getAllocatedSize() == AllocationSize) {
+        return false;
+      }
+
+      // if the Bin Size is greater then current state then update state.
+      if (getAllocatedSize() < OffsetEnd) {
+        changeAllocationSize(
+            std::optional<TypeSize>(TypeSize(OffsetEnd, true)));
+      }
+
+      Type *IntegerSizeOfOffsetEnd =
+          Type::getIntNTy(AI.getContext(), OffsetEnd);
+
+      AI.setAllocatedType(IntegerSizeOfOffsetEnd);
+
+      LLVM_DEBUG(
+          dbgs()
+          << "[AAAllocationInfo, manifest] Successfully reduced allocation "
+             "size of alloca based on Bin."
+          << AI << "\n");
+
+      Changed = ChangeStatus::CHANGED;
+      return true;
+    };
+
+    auto *I = this->getIRPosition().getCtxI();
+    if (!I) {
+      LLVM_DEBUG(
+          dbgs() << "[AAAllocationInfo, manifest] Instruction not valid.\n");
+      return ChangeStatus::UNCHANGED;
+    }
+
+    if (!MakeChange(*I)) {
+      return ChangeStatus::UNCHANGED;
+    }
+
+    LLVM_DEBUG(dbgs() << "[AAAllocationInfo, manifest] return from "
+                         "AAAllocation info manifest.\n");
+    return Changed;
+  }
+
+  /// See AbstractAttribute::getAsStr().
+  const std::string getAsStr(Attributor *A) const override {
+    if (!isValidState())
+      return "allocationinfo(<invalid>)";
+    return "allocationinfo(" +
+           (AssumedAllocatedSize == NOAllocatedSize
+                ? "none"
+                : std::to_string(AssumedAllocatedSize->getFixedValue())) +
+           ")";
+  }
+
+private:
+  std::optional<TypeSize> AssumedAllocatedSize = NOAllocatedSize;
+
+  bool changeAllocationSize(std::optional<TypeSize> Size) {
+    if (AssumedAllocatedSize == NOAllocatedSize ||
+        AssumedAllocatedSize != Size) {
+      AssumedAllocatedSize = Size;
+      return true;
+    }
+    return false;
+  }
+};
+
+struct AAAllocationInfoFloating : AAAllocationInfoImpl {
+  AAAllocationInfoFloating(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfoImpl(IRP, A) {}
+
+  void trackStatistics() const override {
+    STATS_DECLTRACK_FLOATING_ATTR(allocationinfo);
+  }
+};
+
+struct AAAllocationInfoReturned : AAAllocationInfoImpl {
+  AAAllocationInfoReturned(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfoImpl(IRP, A) {}
+
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+    // TODO: we don't rewrite function argument for now because it will need to
+    // rewrite the function signature and all call sites
+    (void)indicatePessimisticFixpoint();
+  }
+
+  void trackStatistics() const override {
+    STATS_DECLTRACK_FNRET_ATTR(allocationinfo);
+  }
+};
+
+struct AAAllocationInfoCallSiteReturned : AAAllocationInfoImpl {
+  AAAllocationInfoCallSiteReturned(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfoImpl(IRP, A) {}
+
+  void trackStatistics() const override {
+    STATS_DECLTRACK_CSRET_ATTR(allocationinfo);
+  }
+};
+
+struct AAAllocationInfoArgument : AAAllocationInfoImpl {
+  AAAllocationInfoArgument(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfoImpl(IRP, A) {}
+
+  void trackStatistics() const override {
+    STATS_DECLTRACK_ARG_ATTR(allocationinfo);
+  }
+};
+
+struct AAAllocationInfoCallSiteArgument : AAAllocationInfoImpl {
+  AAAllocationInfoCallSiteArgument(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfoImpl(IRP, A) {}
+
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+
+    (void)indicatePessimisticFixpoint();
+  }
+
+  void trackStatistics() const override {
+    STATS_DECLTRACK_CSARG_ATTR(allocationinfo);
+  }
+};
+} // namespace
+
 const char AANoUnwind::ID = 0;
 const char AANoSync::ID = 0;
 const char AANoFree::ID = 0;
 const char AANonNull::ID = 0;
 const char AAMustProgress::ID = 0;
 const char AANoRecurse::ID = 0;
 const char AANonConvergent::ID = 0;
 const char AAWillReturn::ID = 0;

 const char AANoUndef::ID = 0;
 const char AANoFPClass::ID = 0;
 const char AACallEdges::ID = 0;
 const char AAInterFnReachability::ID = 0;
 const char AAPointerInfo::ID = 0;
 const char AAAssumptionInfo::ID = 0;
 const char AAUnderlyingObjects::ID = 0;
 const char AAAddressSpace::ID = 0;
+const char AAAllocationInfo::ID = 0;
 
 // Macro magic to create the static generator function for attributes that
 // follow the naming scheme.
 
 #define SWITCH_PK_INV(CLASS, PK, POS_NAME)                                     \
   case IRPosition::PK:                                                         \
     llvm_unreachable("Cannot create " #CLASS " for a " POS_NAME " position!");
 

 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoCapture)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAValueConstantRange)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAPotentialConstantValues)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAPotentialValues)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoUndef)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoFPClass)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAPointerInfo)
 CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAddressSpace)
+CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAllocationInfo)
 
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAValueSimplify)
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAIsDead)
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoFree)
 CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAUnderlyingObjects)
 
 CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAHeapToStack)
 CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAUndefinedBehavior)

llvm/test/Transforms/Attributor/allocator-start-range-only-negative-1.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --check-globals --version 2
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal  -attributor-annotate-decl-cs  -S < %s | FileCheck %s --check-prefixes=CHECK,TUNIT
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,CGSCC
+
+%struct.Foo = type { i32, i8, ptr }
+
+; Function Attrs: noinline nounwind uwtable
+define dso_local ptr @foo(i32 noundef %val) #0 {
+; CHECK-LABEL: define dso_local ptr @foo
+; CHECK-SAME: (i32 noundef [[VAL:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VAL_ADDR:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[F:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    store i32 [[VAL]], ptr [[VAL_ADDR]], align 4
+; CHECK-NEXT:    [[CALL:%.*]] = call noalias ptr @malloc(i64 noundef 16)
+; CHECK-NEXT:    store ptr [[CALL]], ptr [[F]], align 8
+; CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr [[F]], align 8
+; CHECK-NEXT:    store i32 2, ptr [[TMP0]], align 8
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 10, [[VAL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load ptr, ptr [[F]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr [[TMP1]], align 8
+; CHECK-NEXT:    [[ADD2:%.*]] = add nsw i32 [[TMP2]], [[ADD]]
+; CHECK-NEXT:    store i32 [[ADD2]], ptr [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP3:%.*]] = load ptr, ptr [[F]], align 8
+; CHECK-NEXT:    ret ptr [[TMP3]]
+;
+entry:
+  %val.addr = alloca i32, align 4
+  %f = alloca ptr, align 8
+  store i32 %val, ptr %val.addr, align 4
+  %call = call noalias ptr @malloc(i64 noundef 16) #2
+  store ptr %call, ptr %f, align 8
+  %0 = load ptr, ptr %f, align 8
+  %field1 = getelementptr inbounds %struct.Foo, ptr %0, i32 0, i32 0
+  store i32 2, ptr %field1, align 8
+  %1 = load i32, ptr %val.addr, align 4
+  %add = add nsw i32 10, %1
+  %2 = load ptr, ptr %f, align 8
+  %field11 = getelementptr inbounds %struct.Foo, ptr %2, i32 0, i32 0
+  %3 = load i32, ptr %field11, align 8
+  %add2 = add nsw i32 %3, %add
+  store i32 %add2, ptr %field11, align 8
+  %4 = load ptr, ptr %f, align 8
+  ret ptr %4
+}
+
+; Function Attrs: nounwind allocsize(0)
+declare noalias ptr @malloc(i64 noundef) #1
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CGSCC: {{.*}}
+; TUNIT: {{.*}}

llvm/test/Transforms/Attributor/allocator-start-range-only-negative-2.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --check-globals --version 2
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal  -attributor-annotate-decl-cs  -S < %s | FileCheck %s --check-prefixes=CHECK,TUNIT
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,CGSCC
+
+%struct.Foo = type { i32, i8, ptr }
+
+; Function Attrs: noinline nounwind uwtable
+define dso_local { i64, ptr } @foo(i32 noundef %val) #0 {
+; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
+; CHECK-LABEL: define dso_local { i64, ptr } @foo
+; CHECK-SAME: (i32 noundef [[VAL:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RETVAL:%.*]] = alloca [[STRUCT_FOO:%.*]], align 8
+; CHECK-NEXT:    [[VAL_ADDR:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 [[VAL]], ptr [[VAL_ADDR]], align 4
+; CHECK-NEXT:    store i32 2, ptr [[RETVAL]], align 8
+; CHECK-NEXT:    [[FIELD3:%.*]] = getelementptr inbounds [[STRUCT_FOO]], ptr [[RETVAL]], i32 0, i32 2
+; CHECK-NEXT:    store ptr [[VAL_ADDR]], ptr [[FIELD3]], align 8
+; CHECK-NEXT:    [[TMP0:%.*]] = load { i64, ptr }, ptr [[RETVAL]], align 8
+; CHECK-NEXT:    ret { i64, ptr } [[TMP0]]
+;
+entry:
+  %retval = alloca %struct.Foo, align 8
+  %val.addr = alloca i32, align 4
+  store i32 %val, ptr %val.addr, align 4
+  %field1 = getelementptr inbounds %struct.Foo, ptr %retval, i32 0, i32 0
+  store i32 2, ptr %field1, align 8
+  %field3 = getelementptr inbounds %struct.Foo, ptr %retval, i32 0, i32 2
+  store ptr %val.addr, ptr %field3, align 8
+  %0 = load { i64, ptr }, ptr %retval, align 8
+  ret { i64, ptr } %0
+}
+;.
+; CHECK: attributes #[[ATTR0]] = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) }
+;.
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CGSCC: {{.*}}
+; TUNIT: {{.*}}

llvm/test/Transforms/Attributor/allocator-start-range-only-positive-1.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --check-globals --version 2
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal  -attributor-annotate-decl-cs  -S < %s | FileCheck %s --check-prefixes=CHECK,TUNIT
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,CGSCC
+
+%struct.Foo = type { i32, i32, i8 }
+
+@.str = private unnamed_addr constant [15 x i8] c"Field 1 is %d\0A\00", align 1
+
+;.
+; CHECK: @[[_STR:[a-zA-Z0-9_$"\\.-]+]] = private unnamed_addr constant [15 x i8] c"Field 1 is %d\0A\00", align 1
+;.
+define dso_local void @foo(i32 noundef %val) #0 {
+; CHECK-LABEL: define dso_local void @foo
+; CHECK-SAME: (i32 noundef [[VAL:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VAL_ADDR:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[F:%.*]] = alloca [[STRUCT_FOO:%.*]], align 4
+; CHECK-NEXT:    store i32 [[VAL]], ptr [[VAL_ADDR]], align 4
+; CHECK-NEXT:    store i32 10, ptr [[F]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[F]], align 4
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    store i32 [[ADD]], ptr [[F]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr [[F]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[TMP1]], [[VAL]]
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 (ptr, ...) @printf(ptr noundef nonnull dereferenceable(15) @.str, i32 noundef [[ADD3]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %val.addr = alloca i32, align 4
+  %f = alloca %struct.Foo, align 4
+  store i32 %val, ptr %val.addr, align 4
+  %field1 = getelementptr inbounds %struct.Foo, ptr %f, i32 0, i32 0
+  store i32 10, ptr %field1, align 4
+  %field11 = getelementptr inbounds %struct.Foo, ptr %f, i32 0, i32 0
+  %0 = load i32, ptr %field11, align 4
+  %add = add nsw i32 %0, 1
+  store i32 %add, ptr %field11, align 4
+  %field12 = getelementptr inbounds %struct.Foo, ptr %f, i32 0, i32 0
+  %1 = load i32, ptr %field12, align 4
+  %2 = load i32, ptr %val.addr, align 4
+  %add3 = add nsw i32 %1, %2
+  %call = call i32 (ptr, ...) @printf(ptr noundef @.str, i32 noundef %add3)
+  ret void
+}
+
+declare i32 @printf(ptr noundef, ...) #1
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CGSCC: {{.*}}
+; TUNIT: {{.*}}