[AAPointerInfo] refactor how offsets and Access objects are tracked

[sameerds]

Summary

AAPointerInfo now maintains a list of all Access objects that it owns, along
with the following maps:

• OffsetBins: OffsetAndSize -> { Access }
• InstTupleMap: RemoteI x LocalI -> Access

A RemoteI is any instruction that accesses memory. RemoteI is different from
LocalI if and only if LocalI is a call; then RemoteI is some instruction in the
callgraph starting from LocalI.

Motivation: When AAPointerInfo recomputes the offset for an instruction, it sets
the value to Unknown if the new offset is not the same as the old offset. The
instruction must now be moved from its current bin to the bin corresponding to
the new offset. This happens for example, when:

• A PHINode has operands that result in different offsets.
• The same remote inst is reachable from the same local inst via different paths in the callgraph:

  A (local inst)
  |
  B
 / \
C1  C2
 \ /
  D (remote inst)

This fixes a bug where a store is incorrectly eliminated in a lit test.

[jdoerfert]

Can you point out the test that this fixes? Maybe a standalone reproducer would be good too. I read only the commit message and I'm not sure I understand the problem.

[sameerds]

In D136526#3877140, @jdoerfert wrote:

Can you point out the test that this fixes? Maybe a standalone reproducer would be good too. I read only the commit message and I'm not sure I understand the problem.

I did leave comments on the Phab page for that other commit ... https://reviews.llvm.org/rGad98ef8be409 . But I guess it's incorrect to say that the other commit introduced the bug. It's a latent bug that was exposed by tests like phi_no_store_2.

Essentially, the PHI %p in phi_no_store_2 is visited twice. The first time, it is assigned the offset 0, and the second time it is set to unknown. This results in the following bins being printed in the debug log:

Accesses by bin after update:
[0-1] : 1

• 10 - store i8 1, i8* %p, align 1
  • c: i8 1

[-2--1] : 1

• 10 - store i8 1, i8* %p, align 1
  • c: i8 1

Then forallInterferingAccesses comes along and sees only the first bin with the assumption that one access can only be in one bin. Hence it concludes that the store does not have a potential copy in load %l21, and hence eliminates the store. This is incorrect, since the load is still there and it should see the value written by that store. The store can be eliminated only by recognizing the potential copy and first propagating the value to the load.

The root cause is that if an access is already in a bin, and its offset is updated to unknown, then the access needs to first be removed from its current bin. This change ensures this behaviour by improving the plumbing of how accesses are tracked in bin, so that clients don't have to worry about these things.

[jdoerfert]

In D136526#3877410, @sameerds wrote:

The root cause is that if an access is already in a bin, and its offset is updated to unknown, then the access needs to first be removed from its current bin. This change ensures this behaviour by improving the plumbing of how accesses are tracked in bin, so that clients don't have to worry about these things.

I see. The error description makes sense. Can we split this into a fix and the rewrite though? The fix would be to remove the entry from the bin as soon as the offset changes (which can only become invalid I assume). The rewrite gets rid of the bins and keeps lists of instructions per access, right? There is a change in one of the tests that I think is caused by the rewrite and should not be. We also need tests to capture the benefits of the rewrite then. Also, are the bins itself bad or simply bad that we have on offset per access?

llvm/test/Transforms/Attributor/value-simplify-pointer-info.ll

746

Why do we miss out on these propagations? There are no PHIs involved, right? Something is amiss.

[sameerds]

The refactor now accounts for remote/local instructions. This restores all the other affected lit tests.

[sameerds]

In D136526#3880473, @jdoerfert wrote:

In D136526#3877410, @sameerds wrote:

The root cause is that if an access is already in a bin, and its offset is updated to unknown, then the access needs to first be removed from its current bin. This change ensures this behaviour by improving the plumbing of how accesses are tracked in bin, so that clients don't have to worry about these things.

I see. The error description makes sense. Can we split this into a fix and the rewrite though? The fix would be to remove the entry from the bin as soon as the offset changes (which can only become invalid I assume). The rewrite gets rid of the bins and keeps lists of instructions per access, right? There is a change in one of the tests that I think is caused by the rewrite and should not be. We also need tests to capture the benefits of the rewrite then. Also, are the bins itself bad or simply bad that we have on offset per access?

It's the one-offset-per-access that causes loss of information. I've updated the change description to talk more about this. Hope that clarifies the situation! Also fixed the side effect on other places. I had failed to distinguish between multiple local insts pointing to the same remote inst.

[sameerds]

llvm/test/Transforms/Attributor/value-simplify-pointer-info.ll

746

When I don't distinguish between calls that reach the same remote inst, that remote inst is conservatively treated as unknown. I thought it would be okay to "improve" this later, but then realized that the solution is very relevant to the refactoring. The new change does it right, and this side-effect disappears.

[jdoerfert]

Test impact looks good now. I have two issues with the patch but we can resolve both. One is the update of the OAS which needs to also trigger an update of the access kind and value, e.g., if we loose precise information we cannot keep the MUST kind and value around. Second, we should probably avoid nested DenseMaps, I tried to propose some alternatives.

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

271

At this point we need to signal the access it is a May not Must access anymore if the offset or size changed from something that was not unassigned, right? We also need to give up on the value I suppose.

273

Not super happy about first and second. Should we wrap it in setOffset setSize ?

5099

Documentation, for all of these. See the surrounding code.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

832

We could use Access* here to get stable addresses.

835

I'm not sure nested DenseMaps are a good idea. I think it usually ends up using way too much memory.
Let's try to use an encoding without it.

878

This looks like we could map instructions to a list of OAS instead. Or directly the Access*?

896

Also here, if we have RemoteI -> list<Access*> we should be fine, no?

[sameerds]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

832

We could, but then we have to worry about memory allocation on the stack. In fact, I am still wondering about the copy constructor for PointerInfo::State. It seems to make a shallow copy of AccessBins. Doesn't that result in a double free?

If we keep a list of Access*, instead of a list of Access, it looks like more code with no clear benefit. TBH, I am a big fan of "never use new". For now, the unsigned index into the list of Access is sufficient for its very local use. If it needs to be exposed to clients, maybe could wrap it in a custom iterator instead of exposing stable pointers?

Of course I am open to using a list of Access* if you think that is beneficial.

[sameerds]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

832

s/stack/heap

[jdoerfert]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

832

"never use new" is not a good argument. For one, we have a bump allocator for the Attributor so the "cost" argument doesn't apply at all. Second, even if we don't use the bump allocator, new is not inherently costly. DenseMaps of DenseMaps (or similar nested structures) can hover grow fast. They also will inherently use new, so all you do here is to hide the heap allocations and cause more of them.

We could, but then we have to worry about memory allocation on the stack. In fact, I am still wondering about the copy constructor for PointerInfo::State. It seems to make a shallow copy of AccessBins. Doesn't that result in a double free?

As far as I can tell, PpinterInfo::State did not have a copy constructor. The move constructor did not cause double free problems as it cleared the other access bin after moving the content. If the new version requires adequate handling for copy/move constructor, that should be easy to add, no?

Of course I am open to using a list of Access* if you think that is beneficial.

As I said, I had mixed experience with nested maps. This is not exactly the same case but the first one I could find: https://reviews.llvm.org/rG14cb0bdf2b6ca0b7befbb07fe9f73dad5786f59b

[sameerds]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

832

Right, there is no copy constructor in state, but there is a copy assignment operator which does a shallow copy. I guess it didn't matter because it only got used in places where the state is empty.

We don't seem to be talking about the same thing here. I totally agree with not using DenseMap, and I am working on using a list instead. My comment was a reply to this granparent comment:

We could use Access* here to get stable addresses.

Where in this specific case do you see the need for stable addresses?

To me, stable addresses means something like this:

SmallVector<Access*> AccessList

auto *Acc = new Access;
AccessList.push_back(Acc)

Which means now I have to remember to destroy those Access objects later and also worry about deep copies.

If the new version requires adequate handling for copy/move constructor, that should be easy to add, no?

Memory management may be easy, but it is also easily missed. I am quite the fan of coding practices that simply avoid the hassle. For example, I see a commit in git history where the destruction on AccessBins was missed.

I would much prefer the following:

SmallVector<Access> AccessList;
AccessList.emplace_back({....});

Here, AccessList manages the allocation of the Access object and I don't have to worry about destruction and copying.

This is what "never use new" means. To repeat, it has nothing to do with the cost of DenseMap. It also has nothing to do with custom Allocators. The point is that application programmers should refrain from using the new operator, and instead rely on containers to take care of the allocation. The simplest such container is unique_ptr, for example. All containers take Allocators as template arguments, so the question of which Allocator to use is completely orthogonal. "never use new" is a code hygiene concept, not a performance concept.

So do you really see a reason to have stable pointers to Access objects? Otherwise I would much prefer keeping them directly in a SmallVector.

[jdoerfert]

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

271

^ this one we need to handle as well.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

832

I'm happy with alternatives. Stable addresses would have allowed to do avoid all the indices and double lookups, though, that is not by itself a problem. I thought below it would be easier if we could store Access*, if you think we don't need to, that's generally OK.

[sameerds]

Addressed most of the review comments. Replaced the nested DenseMap with a list.

This diff includes https://reviews.llvm.org/D136745 (OffsetAndSize without the
std::pair) merged into it, so that I don't have to rebase the current diff.
Those changes will get filtered out when submitting to tip of main.

[sameerds]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

878

Should we do this in a separate change? Latest diff addresses pretty much all the other comments.

[tschuett]

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

223

Would a class with public and private make this safer to use?

[jdoerfert]

I think this looks good. Fixes an error and I hope doesn't introduce new ones.
Could you, in a follow up, add some tests for the multi access per instruction support you introduced now?

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

223

We use it as a POD, basically a pair with helper functions and nicer names. I doubt there is much value in private members and public setters.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

878

This doesn't need addressing per se. Let's keep it as is for now.

[sameerds]

• Rebased to tip of main.
• Added test for callgraph from the description: a leaf instruction may have Unknown offset if different offsets reach along different paths in the callgraph.

NOTE: I intend to commit this on Sunday night, IST.

[sameerds]

llvm/test/Transforms/Attributor/value-simplify-pointer-info.ll

1707

Oops! This was not supposed to happen. Checking.

[sameerds]

rebased to include a recent fix; added lit tests

[sameerds]

Fixes https://github.com/llvm/llvm-project/issues/58774

When merging to Access objects, don't call reset() on the content, but it needs
to be set to nullptr instead. The latter means that the content was determined
to be unknown. Added an assertion to getWrittenValue() to which would have
caught this more quickly.

[ye-luo]

Pass tests on my side. Thanks

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

struct OffsetAndSize : public std::pair<int64_t, int64_t> {
   using BaseTy = std::pair<int64_t, int64_t>;
   OffsetAndSize(int64_t Offset, int64_t Size) : BaseTy(Offset, Size) {}
   OffsetAndSize(const BaseTy &P) : BaseTy(P) {}
   int64_t getOffset() const { return first; }
   int64_t getSize() const { return second; }
   static OffsetAndSize getUnknown() { return OffsetAndSize(Unknown, Unknown); }
   static OffsetAndSize getUnassigned() {
     return OffsetAndSize(Unassigned, Unassigned);
   }
 
   /// Return true if offset or size are unknown.
   bool offsetOrSizeAreUnknown() const {
     return getOffset() == OffsetAndSize::Unknown ||
            getSize() == OffsetAndSize::Unknown;
   }
 
   /// Return true if offset and size are unknown, thus this is the default
if (offsetOrSizeAreUnknown() || OAS.offsetOrSizeAreUnknown())
       return true;
 
     // Check if one offset point is in the other interval [offset,
     // offset+size].
     return OAS.getOffset() + OAS.getSize() > getOffset() &&
            OAS.getOffset() < getOffset() + getSize();
   }
 
+  OffsetAndSize &operator&=(const OffsetAndSize &R) {
+    if (getOffset() == Unassigned)
+      first = R.getOffset();
+    else if (R.getOffset() != Unassigned && R.getOffset() != getOffset())
+      first = Unknown;
+
+    if (getSize() == Unassigned)
+      second = R.getSize();
+    else if (getSize() == Unknown || R.getSize() == Unknown)
+      second = Unknown;
+    else if (R.getSize() != Unassigned)
+      second = std::max(second, R.getSize());
+
+    return *this;
+  }
+
   /// Constants used to represent special offsets or sizes.
   static constexpr int64_t Unassigned = -1;
   static constexpr int64_t Unknown = -2;
 };
 
 /// Return the initial value of \p Obj with type \p Ty if that is a constant.
 Constant *getInitialValueForObj(Value &Obj, Type &Ty,
                                 const TargetLibraryInfo *TLI,
enum AccessKind {
     AK_MAY_READ_WRITE = AK_MAY | AK_R | AK_W,
     AK_MUST_READ = AK_MUST | AK_R,
     AK_MUST_WRITE = AK_MUST | AK_W,
     AK_MUST_READ_WRITE = AK_MUST | AK_R | AK_W,
   };
 
   /// An access description.
   struct Access {
-    Access(Instruction *I, Optional<Value *> Content, AccessKind Kind, Type *Ty)
-        : LocalI(I), RemoteI(I), Content(Content), Kind(Kind), Ty(Ty) {
+    Access(Instruction *I, int64_t Offset, int64_t Size,
+           Optional<Value *> Content, AccessKind Kind, Type *Ty)
+        : LocalI(I), RemoteI(I), Content(Content), OAS(Offset, Size),
+          Kind(Kind), Ty(Ty) {
       verify();
     }
-    Access(Instruction *LocalI, Instruction *RemoteI, Optional<Value *> Content,
-           AccessKind Kind, Type *Ty)
-        : LocalI(LocalI), RemoteI(RemoteI), Content(Content), Kind(Kind),
-          Ty(Ty) {
+    Access(Instruction *LocalI, Instruction *RemoteI, int64_t Offset,
+           int64_t Size, Optional<Value *> Content, AccessKind Kind, Type *Ty)
+        : LocalI(LocalI), RemoteI(RemoteI), Content(Content), OAS(Offset, Size),
+          Kind(Kind), Ty(Ty) {
       verify();
     }
     Access(const Access &Other) = default;
     Access(const Access &&Other)
         : LocalI(Other.LocalI), RemoteI(Other.RemoteI), Content(Other.Content),
-          Kind(Other.Kind), Ty(Other.Ty) {}
+          OAS(Other.OAS), Kind(Other.Kind), Ty(Other.Ty) {}
 
     Access &operator=(const Access &Other) = default;
     bool operator==(const Access &R) const {
-      return LocalI == R.LocalI && RemoteI == R.RemoteI &&
+      return LocalI == R.LocalI && RemoteI == R.RemoteI && OAS == R.OAS &&
              Content == R.Content && Kind == R.Kind;
     }
     bool operator!=(const Access &R) const { return !(*this == R); }
 
     Access &operator&=(const Access &R) {
       assert(RemoteI == R.RemoteI && "Expected same instruction!");
+      assert(LocalI == R.LocalI && "Expected same instruction!");
       Content =
           AA::combineOptionalValuesInAAValueLatice(Content, R.Content, Ty);
       Kind = AccessKind(Kind | R.Kind);
+      OAS &= R.OAS;
       return *this;
     }
 
     void verify() {
       assert(isMustAccess() + isMayAccess() == 1 &&
              "Expect must or may access, not both.");
     }
 
struct Access {
     /// isWrittenValueYetUndetermined return true this function shall not be
     /// called.
     Value *getWrittenValue() const { return *Content; }
 
     /// Return the written value which can be `llvm::null` if it is not yet
     /// determined.
     Optional<Value *> getContent() const { return Content; }
 
+    int64_t getOffset() const { return OAS.getOffset(); }
+    int64_t getSize() const { return OAS.getSize(); }
+
   private:
     /// The instruction responsible for the access with respect to the local
     /// scope of the associated attribute.
     Instruction *LocalI;
 
     /// The instruction responsible for the access.
     Instruction *RemoteI;
 
     /// The value written, if any. `llvm::none` means "not known yet", `nullptr`
     /// cannot be determined.
     Optional<Value *> Content;
 
+    AA::OffsetAndSize OAS = AA::OffsetAndSize::getUnassigned();
+
     /// The access kind, e.g., READ, as bitset (could be more than one).
     AccessKind Kind;
 
     /// The type of the content, thus the type read/written, can be null if not
     /// available.
     Type *Ty;
   };
 
struct AAPointerInfo : public AbstractAttribute {
   /// 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,
       function_ref<bool(const Access &, bool)> CB, bool &HasBeenWrittenTo,
-      AA::OffsetAndSize *OASPtr = nullptr) const = 0;
+      AA::OffsetAndSize &OAS) const = 0;
 
   /// This function should return true if the type of the \p AA is AAPointerInfo
   static bool classof(const AbstractAttribute *AA) {
     return (AA->getIdAddr() == &ID);
   }
 
   /// Unique ID (due to the unique address)
   static const char ID;

llvm/lib/Transforms/IPO/Attributor.cpp

for (Value *Obj : Objects) {
     // If the value has been written to we don't need the initial value of the
     // object.
     bool HasBeenWrittenTo = false;
 
     AA::OffsetAndSize OAS = AA::OffsetAndSize::getUnassigned();
     auto &PI = A.getAAFor<AAPointerInfo>(QueryingAA, IRPosition::value(*Obj),
                                          DepClassTy::NONE);
     if (!PI.forallInterferingAccesses(A, QueryingAA, I, CheckAccess,
-                                      HasBeenWrittenTo, &OAS)) {
+                                      HasBeenWrittenTo, OAS)) {
       LLVM_DEBUG(
           dbgs()
           << "Failed to verify all interfering accesses for underlying object: "
           << *Obj << "\n");
       return false;
     }
 
     if (IsLoad && !HasBeenWrittenTo && !OAS.isUnassigned()) {

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

struct AccessAsInstructionInfo : DenseMapInfo<Instruction *> {
   static bool isEqual(const Access &LHS, const Access &RHS);
 };
 
 } // namespace llvm
 
 /// A type to track pointer/struct usage and accesses for AAPointerInfo.
 struct AA::PointerInfo::State : public AbstractState {
 
-  ~State() {
-    // We do not delete the Accesses objects but need to destroy them still.
-    for (auto &It : AccessBins)
-      It.second->~Accesses();
-  }
-
   /// Return the best possible representable state.
   static State getBestState(const State &SIS) { return State(); }
 
   /// Return the worst possible representable state.
   static State getWorstState(const State &SIS) {
     State R;
     R.indicatePessimisticFixpoint();
     return R;
   }
 
   State() = default;
-  State(State &&SIS) : AccessBins(std::move(SIS.AccessBins)) {
-    SIS.AccessBins.clear();
-  }
+  State(State &&SIS) = default;
 
   const State &getAssumed() const { return *this; }
 
   /// See AbstractState::isValidState().
   bool isValidState() const override { return BS.isValidState(); }
 
   /// See AbstractState::isAtFixpoint().
   bool isAtFixpoint() const override { return BS.isAtFixpoint(); }
ChangeStatus indicatePessimisticFixpoint() override {
     BS.indicatePessimisticFixpoint();
     return ChangeStatus::CHANGED;
   }
 
   State &operator=(const State &R) {
     if (this == &R)
       return *this;
     BS = R.BS;
-    AccessBins = R.AccessBins;
+    AccessList = R.AccessList;
+    OffsetBins = R.OffsetBins;
+    InstTupleMap = R.InstTupleMap;
     return *this;
   }
 
   State &operator=(State &&R) {
     if (this == &R)
       return *this;
     std::swap(BS, R.BS);
-    std::swap(AccessBins, R.AccessBins);
+    std::swap(AccessList, R.AccessList);
+    std::swap(OffsetBins, R.OffsetBins);
+    std::swap(InstTupleMap, R.InstTupleMap);
     return *this;
   }
 
-  bool operator==(const State &R) const {
-    if (BS != R.BS)
-      return false;
-    if (AccessBins.size() != R.AccessBins.size())
-      return false;
-    auto It = begin(), RIt = R.begin(), E = end();
-    while (It != E) {
-      if (It->getFirst() != RIt->getFirst())
-        return false;
-      auto &Accs = It->getSecond();
-      auto &RAccs = RIt->getSecond();
-      if (Accs->size() != RAccs->size())
-        return false;
-      for (const auto &ZipIt : llvm::zip(*Accs, *RAccs))
-        if (std::get<0>(ZipIt) != std::get<1>(ZipIt))
-          return false;
-      ++It;
-      ++RIt;
-    }
-    return true;
-  }
-  bool operator!=(const State &R) const { return !(*this == R); }
-
-  /// We store accesses in a set with the instruction as key.
-  struct Accesses {
-    SmallVector<AAPointerInfo::Access, 4> Accesses;
-    DenseMap<const Instruction *, unsigned> Map;
-
-    unsigned size() const { return Accesses.size(); }
-
-    using vec_iterator = decltype(Accesses)::iterator;
-    vec_iterator begin() { return Accesses.begin(); }
-    vec_iterator end() { return Accesses.end(); }
-
-    using iterator = decltype(Map)::const_iterator;
-    iterator find(AAPointerInfo::Access &Acc) {
-      return Map.find(Acc.getRemoteInst());
-    }
-    iterator find_end() { return Map.end(); }
-
-    AAPointerInfo::Access &get(iterator &It) {
-      return Accesses[It->getSecond()];
-    }
-
-    void insert(AAPointerInfo::Access &Acc) {
-      Map[Acc.getRemoteInst()] = Accesses.size();
-      Accesses.push_back(Acc);
-    }
-  };
-
-  /// We store all accesses in bins denoted by their offset and size.
-  using AccessBinsTy = DenseMap<AA::OffsetAndSize, Accesses *>;
-
-  AccessBinsTy::const_iterator begin() const { return AccessBins.begin(); }
-  AccessBinsTy::const_iterator end() const { return AccessBins.end(); }
-
-protected:
-  /// The bins with all the accesses for the associated pointer.
-  AccessBinsTy AccessBins;
-
-  /// Add a new access to the state at offset \p Offset and with size \p Size.
+  /// Add a new Access to the state at offset \p Offset and with size \p Size.
   /// The access is associated with \p I, writes \p Content (if anything), and
-  /// is of kind \p Kind.
+  /// is of kind \p Kind. If an Access already exists for the same \p I and same
+  /// \p RemoteI, the two are combined, potentially losing information about
+  /// offset and size. The resulting access must now be moved from its original
+  /// OffsetBin to the bin for its new offset.
+  ///
   /// \Returns CHANGED, if the state changed, UNCHANGED otherwise.
   ChangeStatus addAccess(Attributor &A, int64_t Offset, int64_t Size,
                          Instruction &I, Optional<Value *> Content,
                          AAPointerInfo::AccessKind Kind, Type *Ty,
-                         Instruction *RemoteI = nullptr,
-                         Accesses *BinPtr = nullptr) {
-    AA::OffsetAndSize Key{Offset, Size};
-    Accesses *&Bin = BinPtr ? BinPtr : AccessBins[Key];
-    if (!Bin)
-      Bin = new (A.Allocator) Accesses;
-    AAPointerInfo::Access Acc(&I, RemoteI ? RemoteI : &I, Content, Kind, Ty);
-    // Check if we have an access for this instruction in this bin, if not,
-    // simply add it.
-    auto It = Bin->find(Acc);
-    if (It == Bin->find_end()) {
-      Bin->insert(Acc);
-      return ChangeStatus::CHANGED;
-    }
-    // If the existing access is the same as then new one, nothing changed.
-    AAPointerInfo::Access &Current = Bin->get(It);
-    AAPointerInfo::Access Before = Current;
-    // The new one will be combined with the existing one.
-    Current &= Acc;
-    return Current == Before ? ChangeStatus::UNCHANGED : ChangeStatus::CHANGED;
+                         Instruction *RemoteI = nullptr);
+
+  using OffsetBinsTy = DenseMap<OffsetAndSize, 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(); }
+
+  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:
+  //
+  // - OffsetBins: OffsetAndSize -> { Access }
+  // - InstTupleMap: RemoteI x LocalI -> Access
+  //
+  // A RemoteI is any instruction that accesses memory. RemoteI is different
+  // from LocalI if and only if LocalI is a call; then RemoteI is some
+  // instruction in the callgraph starting from LocalI. Multiple paths in the
+  // callgraph from LocalI to RemoteI may produce multiple accesses, but these
+  // are all combined into a single Access object. This may result in loss of
+  // information in OffsetAndSize in the Access object.
+  SmallVector<AAPointerInfo::Access, 4> AccessList;
+  OffsetBinsTy OffsetBins;
+  DenseMap<const Instruction *, DenseMap<const Instruction *, unsigned>>
+      InstTupleMap;
+
   /// See AAPointerInfo::forallInterferingAccesses.
   bool forallInterferingAccesses(
       AA::OffsetAndSize OAS,
       function_ref<bool(const AAPointerInfo::Access &, bool)> CB) const {
     if (!isValidState())
       return false;
 
-    for (const auto &It : AccessBins) {
+    for (const auto &It : OffsetBins) {
       AA::OffsetAndSize ItOAS = It.getFirst();
       if (!OAS.mayOverlap(ItOAS))
         continue;
       bool IsExact = OAS == ItOAS && !OAS.offsetOrSizeAreUnknown();
-      for (auto &Access : *It.getSecond())
+      for (auto Index : It.getSecond()) {
+        auto &Access = AccessList[Index];
         if (!CB(Access, IsExact))
           return false;
       }
+    }
     return true;
   }
 
   /// See AAPointerInfo::forallInterferingAccesses.
   bool forallInterferingAccesses(
       Instruction &I,
       function_ref<bool(const AAPointerInfo::Access &, bool)> CB,
-      AA::OffsetAndSize *OASPtr) const {
+      AA::OffsetAndSize &OAS) const {
     if (!isValidState())
       return false;
 
-    // First find the offset and size of I.
-    AA::OffsetAndSize OAS = AA::OffsetAndSize::getUnassigned();
-    for (const auto &It : AccessBins) {
-      for (auto &Access : *It.getSecond()) {
-        if (Access.getRemoteInst() == &I) {
-          OAS = It.getFirst();
-          break;
-        }
-      }
-      if (OAS.getSize() != AA::OffsetAndSize::Unassigned)
-        break;
+    auto LocalMap = InstTupleMap.find(&I);
+    if (LocalMap == InstTupleMap.end()) {
+      return true;
     }
 
-    if (OASPtr)
-      *OASPtr = OAS;
-
-    // No access for I was found, we are done.
-    if (OAS.getSize() == AA::OffsetAndSize::Unassigned)
-      return true;
+    bool Changed = false;
+    for (auto LI : LocalMap->getSecond()) {
+      auto &Access = AccessList[LI.getSecond()];
+      OAS &= {Access.getOffset(), Access.getSize()};
+      Changed |= forallInterferingAccesses(OAS, CB);
+    }
 
-    // Now that we have an offset and size, find all overlapping ones and use
-    // the callback on the accesses.
-    return forallInterferingAccesses(OAS, CB);
+    return Changed;
   }
 
 private:
   /// State to track fixpoint and validity.
   BooleanState BS;
 };
 
+ChangeStatus AA::PointerInfo::State::addAccess(Attributor &A, int64_t Offset,
+                                               int64_t Size, Instruction &I,
+                                               Optional<Value *> Content,
+                                               AAPointerInfo::AccessKind Kind,
+                                               Type *Ty, Instruction *RemoteI) {
+  RemoteI = RemoteI ? RemoteI : &I;
+  AAPointerInfo::Access Acc(&I, RemoteI, Offset, Size, Content, Kind, Ty);
+
+  // Check if we have an access for this instruction, if not, simply add it.
+  auto &LocalMap = InstTupleMap[RemoteI];
+  auto [AccIt, AccInserted] = LocalMap.try_emplace(&I, AccessList.size());
+  auto Index = AccIt->getSecond();
+  if (AccInserted) {
+    assert(Index == AccessList.size() && "Expected to append new access.");
+    AccessList.push_back(Acc);
+  } else {
+    // The new one will be combined with the existing one.
+    auto &Current = AccessList[Index];
+    auto Before = Current;
+    Current &= Acc;
+    if (Current == Before)
+      return ChangeStatus::UNCHANGED;
+
+    LLVM_DEBUG(dbgs() << "Access changed from " << Before << " to " << Current
+                      << "\n");
+    Acc = Current;
+    AA::OffsetAndSize Key{Before.getOffset(), Before.getSize()};
+    assert(OffsetBins.count(Key) && "Existing Access must be in some bin.");
+    auto &Bin = OffsetBins[Key];
+    assert(Bin.count(Index) && "Expected bin to actually contain the Access.");
+    LLVM_DEBUG(dbgs() << "[AAPointerInfo] Removing Access " << AccessList[Index]
+                      << " with key " << '{' << Key.getOffset() << ','
+                      << Key.getSize() << '}' << "\n");
+    Bin.erase(Index);
+  }
+
+  AA::OffsetAndSize Key{Acc.getOffset(), Acc.getSize()};
+  LLVM_DEBUG(dbgs() << "[AAPointerInfo] Inserting Access " << Acc
+                    << " with key " << '{' << Key.getOffset() << ','
+                    << Key.getSize() << '}' << "\n");
+  OffsetBins[Key].insert(Index);
+  return ChangeStatus::CHANGED;
+}
+
 namespace {
 struct AAPointerInfoImpl
     : public StateWrapper<AA::PointerInfo::State, AAPointerInfo> {
   using BaseTy = StateWrapper<AA::PointerInfo::State, AAPointerInfo>;
   AAPointerInfoImpl(const IRPosition &IRP, Attributor &A) : BaseTy(IRP) {}
 
   /// See AbstractAttribute::getAsStr().
   const std::string getAsStr() const override {
     return std::string("PointerInfo ") +
            (isValidState() ? (std::string("#") +
-                              std::to_string(AccessBins.size()) + " bins")
+                              std::to_string(OffsetBins.size()) + " bins")
                            : "<invalid>");
   }
 
   /// See AbstractAttribute::manifest(...).
   ChangeStatus manifest(Attributor &A) override {
     return AAPointerInfo::manifest(A);
   }
 
   bool forallInterferingAccesses(
       AA::OffsetAndSize OAS,
       function_ref<bool(const AAPointerInfo::Access &, bool)> CB)
       const override {
     return State::forallInterferingAccesses(OAS, CB);
   }
 
   bool forallInterferingAccesses(
       Attributor &A, const AbstractAttribute &QueryingAA, Instruction &I,
       function_ref<bool(const Access &, bool)> UserCB, bool &HasBeenWrittenTo,
-      AA::OffsetAndSize *OASPtr = nullptr) const override {
+      AA::OffsetAndSize &OAS) const override {
     HasBeenWrittenTo = false;
 
     SmallPtrSet<const Access *, 8> DominatingWrites;
     SmallVector<std::pair<const Access *, bool>, 8> InterferingAccesses;
 
     Function &Scope = *I.getFunction();
     const auto &NoSyncAA = A.getAAFor<AANoSync>(
         QueryingAA, IRPosition::function(Scope), DepClassTy::OPTIONAL);
auto AccessCB = [&](const Access &Acc, bool Exact) {
       // work yet if we have threading effects, or the access is complicated.
       if (CanUseCFGResoning && Dominates && UseDominanceReasoning &&
           IsSameThreadAsLoad(Acc))
         DominatingWrites.insert(&Acc);
 
       InterferingAccesses.push_back({&Acc, Exact});
       return true;
     };
-    if (!State::forallInterferingAccesses(I, AccessCB, OASPtr))
+    if (!State::forallInterferingAccesses(I, AccessCB, OAS))
       return false;
 
     if (HasBeenWrittenTo) {
       const Function *ScopePtr = &Scope;
       IsLiveInCalleeCB = [ScopePtr](const Function &Fn) {
         return ScopePtr != &Fn;
       };
     }
if (!OtherAA.getState().isValidState() || !isValidState())
       return indicatePessimisticFixpoint();
 
     const auto &OtherAAImpl = static_cast<const AAPointerInfoImpl &>(OtherAA);
     bool IsByval =
         FromCallee && OtherAAImpl.getAssociatedArgument()->hasByValAttr();
 
     // Combine the accesses bin by bin.
     ChangeStatus Changed = ChangeStatus::UNCHANGED;
-    for (const auto &It : OtherAAImpl.getState()) {
+    const auto &State = OtherAAImpl.getState();
+    for (const auto &It : State) {
       AA::OffsetAndSize OAS = AA::OffsetAndSize::getUnknown();
       if (Offset != AA::OffsetAndSize::Unknown)
         OAS = AA::OffsetAndSize(It.first.getOffset() + Offset,
                                 It.first.getSize());
-      Accesses *Bin = AccessBins.lookup(OAS);
-      for (const AAPointerInfo::Access &RAcc : *It.second) {
+      for (auto Index : It.getSecond()) {
+        const auto &RAcc = State.getAccess(Index);
         if (IsByval && !RAcc.isRead())
           continue;
         bool UsedAssumedInformation = false;
         AccessKind AK = RAcc.getKind();
         Optional<Value *> Content = RAcc.getContent();
         if (FromCallee) {
           Content = A.translateArgumentToCallSiteContent(
               RAcc.getContent(), CB, *this, UsedAssumedInformation);
           AK =
               AccessKind(AK & (IsByval ? AccessKind::AK_R : AccessKind::AK_RW));
           AK = AccessKind(AK | (RAcc.isMayAccess() ? AK_MAY : AK_MUST));
         }
         Changed =
             Changed | addAccess(A, OAS.getOffset(), OAS.getSize(), CB, Content,
-                                AK, RAcc.getType(), RAcc.getRemoteInst(), Bin);
+                                AK, RAcc.getType(), RAcc.getRemoteInst());
       }
     }
     return Changed;
   }
 
   /// Statistic tracking for all AAPointerInfo implementations.
   /// See AbstractAttribute::trackStatistics().
   void trackPointerInfoStatistics(const IRPosition &IRP) const {}
 
   /// Dump the state into \p O.
   void dumpState(raw_ostream &O) {
-    for (auto &It : AccessBins) {
+    for (auto &It : OffsetBins) {
       O << "[" << It.first.getOffset() << "-"
         << It.first.getOffset() + It.first.getSize()
-        << "] : " << It.getSecond()->size() << "\n";
-      for (auto &Acc : *It.getSecond()) {
+        << "] : " << It.getSecond().size() << "\n";
+      for (auto AccIndex : It.getSecond()) {
+        auto &Acc = AccessList[AccIndex];
         O << "     - " << Acc.getKind() << " - " << *Acc.getLocalInst() << "\n";
         if (Acc.getLocalInst() != Acc.getRemoteInst())
           O << "     -->                         " << *Acc.getRemoteInst()
             << "\n";
         if (!Acc.isWrittenValueYetUndetermined()) {
           if (Acc.getWrittenValue())
             O << "       - c: " << *Acc.getWrittenValue() << "\n";
           else

llvm/test/Transforms/Attributor/value-simplify-pointer-info.ll

 ; TUNIT-NEXT:    store float 0x40119999A0000000, float* [[F2]], align 4, !tbaa [[TBAA10]]
 ; TUNIT-NEXT:    [[F3:%.*]] = getelementptr inbounds [[STRUCT_S]], %struct.S* [[AGG_RESULT]], i64 0, i32 5
 ; TUNIT-NEXT:    store float 0x40119999A0000000, float* [[F3]], align 4, !tbaa [[TBAA11]]
 ; TUNIT-NEXT:    [[I1:%.*]] = getelementptr inbounds [[STRUCT_S]], %struct.S* [[AGG_RESULT]], i64 0, i32 0
 ; TUNIT-NEXT:    store i32 1, i32* [[I1]], align 4, !tbaa [[TBAA12]]
 ; TUNIT-NEXT:    [[I2:%.*]] = getelementptr inbounds [[STRUCT_S]], %struct.S* [[AGG_RESULT]], i64 0, i32 1
 ; TUNIT-NEXT:    store i32 4, i32* [[I2]], align 4, !tbaa [[TBAA13]]
 ; TUNIT-NEXT:    [[I3:%.*]] = getelementptr inbounds [[STRUCT_S]], %struct.S* [[AGG_RESULT]], i64 0, i32 2
 ; TUNIT-NEXT:    store i32 4, i32* [[I3]], align 4, !tbaa [[TBAA14]]
 ; TUNIT-NEXT:    ret void
 ;
 ; CGSCC: Function Attrs: nofree nosync nounwind willreturn
 ; CGSCC-LABEL: define {{[^@]+}}@static_global_simplifiable_1
 ; CGSCC-SAME: (%struct.S* noalias nocapture nofree nonnull writeonly sret([[STRUCT_S:%.*]]) align 4 dereferenceable(24) [[AGG_RESULT:%.*]]) #[[ATTR3]] {
 ; CGSCC-NEXT:  entry:
 ; CGSCC-NEXT:    store float 0x3FF19999A0000000, float* getelementptr inbounds ([[STRUCT_S]], %struct.S* @Gs1, i64 0, i32 3), align 4, !tbaa [[TBAA7]]
 ; CGSCC-NEXT:    store float 0x40019999A0000000, float* getelementptr inbounds ([[STRUCT_S]], %struct.S* @Gs1, i64 0, i32 4), align 4, !tbaa [[TBAA10]]

 ; CHECK-NEXT:    [[L2:%.*]] = load i32, i32* [[GEP2]], align 4
 ; CHECK-NEXT:    [[L3:%.*]] = load i32, i32* [[GEP3]], align 4
 ; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[L1]], [[L2]]
 ; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[ADD1]], [[L3]]
 ; CHECK-NEXT:    ret i32 [[ADD2]]
 ;
 entry:
   %s = alloca %struct.S, align 4
   %bc = bitcast %struct.S* %s to i32*
   %gep1 = getelementptr inbounds %struct.S, %struct.S* %s, i64 0, i32 2
   %gep2 = getelementptr inbounds i32, i32* %bc, i32 %arg1
   %gep3 = getelementptr inbounds i32, i32* %bc, i32 %arg2
   store i32 7, i32* %gep1
   store i32 7, i32* %gep2
   store i32 7, i32* %gep3
   %l1 = load i32, i32* %gep1
   %l2 = load i32, i32* %gep2

 ; TUNIT: Function Attrs: nofree norecurse nosync nounwind willreturn
 ; TUNIT-LABEL: define {{[^@]+}}@phi_no_store_2
 ; TUNIT-SAME: () #[[ATTR3]] {
 ; TUNIT-NEXT:  entry:
 ; TUNIT-NEXT:    br label [[LOOP:%.*]]
 ; TUNIT:       loop:
 ; TUNIT-NEXT:    [[P:%.*]] = phi i8* [ bitcast (i32* @a2 to i8*), [[ENTRY:%.*]] ], [ [[G:%.*]], [[LOOP]] ]
 ; TUNIT-NEXT:    [[I:%.*]] = phi i8 [ 0, [[ENTRY]] ], [ [[O:%.*]], [[LOOP]] ]
+; TUNIT-NEXT:    store i8 1, i8* [[P]], align 2
 ; TUNIT-NEXT:    [[G]] = getelementptr i8, i8* bitcast (i32* @a2 to i8*), i64 2
 ; TUNIT-NEXT:    [[O]] = add nsw i8 [[I]], 1
 ; TUNIT-NEXT:    [[C:%.*]] = icmp eq i8 [[O]], 7
 ; TUNIT-NEXT:    br i1 [[C]], label [[END:%.*]], label [[LOOP]]
 ; TUNIT:       end:
 ; TUNIT-NEXT:    [[L21:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a2 to i8*), i64 2), align 2
 ; TUNIT-NEXT:    [[L22:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a2 to i8*), i64 3), align 1
 ; TUNIT-NEXT:    [[ADD:%.*]] = add i8 [[L21]], [[L22]]
 ; TUNIT-NEXT:    ret i8 [[ADD]]
 ;
 ; CGSCC: Function Attrs: nofree norecurse nosync nounwind willreturn
 ; CGSCC-LABEL: define {{[^@]+}}@phi_no_store_2
 ; CGSCC-SAME: () #[[ATTR5]] {
 ; CGSCC-NEXT:  entry:
 ; CGSCC-NEXT:    br label [[LOOP:%.*]]
 ; CGSCC:       loop:
 ; CGSCC-NEXT:    [[P:%.*]] = phi i8* [ bitcast (i32* @a2 to i8*), [[ENTRY:%.*]] ], [ [[G:%.*]], [[LOOP]] ]
 ; CGSCC-NEXT:    [[I:%.*]] = phi i8 [ 0, [[ENTRY]] ], [ [[O:%.*]], [[LOOP]] ]
+; CGSCC-NEXT:    store i8 1, i8* [[P]], align 2
 ; CGSCC-NEXT:    [[G]] = getelementptr i8, i8* bitcast (i32* @a2 to i8*), i64 2
 ; CGSCC-NEXT:    [[O]] = add nsw i8 [[I]], 1
 ; CGSCC-NEXT:    [[C:%.*]] = icmp eq i8 [[O]], 7
 ; CGSCC-NEXT:    br i1 [[C]], label [[END:%.*]], label [[LOOP]]
 ; CGSCC:       end:
 ; CGSCC-NEXT:    [[L21:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a2 to i8*), i64 2), align 2
 ; CGSCC-NEXT:    [[L22:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a2 to i8*), i64 3), align 1
 ; CGSCC-NEXT:    [[ADD:%.*]] = add i8 [[L21]], [[L22]]

 ; TUNIT-LABEL: define {{[^@]+}}@phi_no_store_3
 ; TUNIT-SAME: () #[[ATTR3]] {
 ; TUNIT-NEXT:  entry:
 ; TUNIT-NEXT:    store i8 0, i8* getelementptr (i8, i8* bitcast (i32* @a3 to i8*), i64 3), align 1
 ; TUNIT-NEXT:    br label [[LOOP:%.*]]
 ; TUNIT:       loop:
 ; TUNIT-NEXT:    [[P:%.*]] = phi i8* [ bitcast (i32* @a3 to i8*), [[ENTRY:%.*]] ], [ [[G:%.*]], [[LOOP]] ]
 ; TUNIT-NEXT:    [[I:%.*]] = phi i8 [ 0, [[ENTRY]] ], [ [[O:%.*]], [[LOOP]] ]
+; TUNIT-NEXT:    store i8 1, i8* [[P]], align 2
 ; TUNIT-NEXT:    [[G]] = getelementptr i8, i8* bitcast (i32* @a3 to i8*), i64 2
 ; TUNIT-NEXT:    [[O]] = add nsw i8 [[I]], 1
 ; TUNIT-NEXT:    [[C:%.*]] = icmp eq i8 [[O]], 7
 ; TUNIT-NEXT:    br i1 [[C]], label [[END:%.*]], label [[LOOP]]
 ; TUNIT:       end:
 ; TUNIT-NEXT:    [[L31:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a3 to i8*), i64 2), align 2
 ; TUNIT-NEXT:    [[L32:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a3 to i8*), i64 3), align 1
 ; TUNIT-NEXT:    [[ADD:%.*]] = add i8 [[L31]], [[L32]]
 ; TUNIT-NEXT:    [[L34:%.*]] = load i8, i8* bitcast (i32* getelementptr inbounds (i32, i32* @a3, i64 1) to i8*), align 4
 ; TUNIT-NEXT:    [[ADD2:%.*]] = add i8 [[ADD]], [[L34]]
 ; TUNIT-NEXT:    ret i8 [[ADD2]]
 ;
 ; CGSCC: Function Attrs: nofree norecurse nosync nounwind willreturn
 ; CGSCC-LABEL: define {{[^@]+}}@phi_no_store_3
 ; CGSCC-SAME: () #[[ATTR5]] {
 ; CGSCC-NEXT:  entry:
 ; CGSCC-NEXT:    store i8 0, i8* getelementptr (i8, i8* bitcast (i32* @a3 to i8*), i64 3), align 1
 ; CGSCC-NEXT:    br label [[LOOP:%.*]]
 ; CGSCC:       loop:
 ; CGSCC-NEXT:    [[P:%.*]] = phi i8* [ bitcast (i32* @a3 to i8*), [[ENTRY:%.*]] ], [ [[G:%.*]], [[LOOP]] ]
 ; CGSCC-NEXT:    [[I:%.*]] = phi i8 [ 0, [[ENTRY]] ], [ [[O:%.*]], [[LOOP]] ]
+; CGSCC-NEXT:    store i8 1, i8* [[P]], align 2
 ; CGSCC-NEXT:    [[G]] = getelementptr i8, i8* bitcast (i32* @a3 to i8*), i64 2
 ; CGSCC-NEXT:    [[O]] = add nsw i8 [[I]], 1
 ; CGSCC-NEXT:    [[C:%.*]] = icmp eq i8 [[O]], 7
 ; CGSCC-NEXT:    br i1 [[C]], label [[END:%.*]], label [[LOOP]]
 ; CGSCC:       end:
 ; CGSCC-NEXT:    [[L31:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a3 to i8*), i64 2), align 2
 ; CGSCC-NEXT:    [[L32:%.*]] = load i8, i8* getelementptr (i8, i8* bitcast (i32* @a3 to i8*), i64 3), align 1
 ; CGSCC-NEXT:    [[ADD:%.*]] = add i8 [[L31]], [[L32]]