[llvm] [Attributor] Change allocation size and load/store offsets using AAPointerInfo for Alloca instructions (PR #72029)

[Shilei Tian via llvm-commits]

================
@@ -12748,39 +12735,227 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
     assert(isValidState() &&
            "Manifest should only be called if the state is valid.");
 
-    Instruction *I = getIRPosition().getCtxI();
+    bool Changed = false;
+    const IRPosition &IRP = getIRPosition();
+    Instruction *I = IRP.getCtxI();
+
+    // check if simplified values exist
+    if (simplifiedValuesExists(A, I))
+      return ChangeStatus::UNCHANGED;
+
+    if (getAllocatedSize() == HasNoAllocationSize)
+      return ChangeStatus::UNCHANGED;
 
-    auto FixedAllocatedSizeInBits = getAllocatedSize()->getFixedValue();
+    const AAPointerInfo *PI =
+        A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::REQUIRED);
+
+    if (!PI)
+      return ChangeStatus::UNCHANGED;
+
+    if (!PI->getState().isValidState())
+      return ChangeStatus::UNCHANGED;
+
+    // Store a map where each instruction maps to a set of bins accessed by that
+    // instruction
+    DenseMap<Instruction *, DenseMap<AA::RangeTy, AA::RangeTy>>
+        AccessedInstructionsToBinsMap;
+
+    const auto &NewOffsetsMap = getNewOffsets();
+    const auto &OffsetInfoMap = PI->getOffsetInfoMap();
+
+    // Map Instructions to accessed bins.
+    for (AAPointerInfo::OffsetBinsTy::const_iterator It = PI->begin();
+         It != PI->end(); It++) {
+
+      const auto &OldOffsetRange = It->getFirst();
+
+      // If the OldOffsetRange is not in the map, offsets for that bin did not
+      // change. We should just continue and skip changing the offsets in that
+      // case
+      if (!NewOffsetsMap.contains(OldOffsetRange))
+        continue;
+
+      const auto &NewOffsetRange = NewOffsetsMap.lookup(OldOffsetRange);
+
+      for (const auto AccIndex : It->getSecond()) {
+        const auto &AccessInstruction = PI->getBinAccess(AccIndex);
+        Instruction *LocalInst = AccessInstruction.getLocalInst();
+
+        // TODO: handle case for a similified value
+        // Right now we don't change the value and give up
+        // on modifying the size and offsets of the allocation
+        // this may be sub-optimal
+        if (simplifiedValuesExists(A, LocalInst))
+          return ChangeStatus::UNCHANGED;
+
+        // BackTrack and check if there are multiple bins for instructions in
+        // the
+        // chain
+        std::vector<Instruction *> ReadyList;
+        DenseMap<Instruction *, bool> Visited;
+        ReadyList.push_back(LocalInst);
+        while (!ReadyList.empty()) {
+          Instruction *GetBack = ReadyList.back();
+          ReadyList.pop_back();
+          // check if the Instruction has multiple bins, if so give up
+          // for calls it is okay to have multiple bins
+          // TODO: handle when one instruction has multiple bins
+          auto OffsetsVecArg = OffsetInfoMap.lookup(GetBack).Offsets;
+          if (GetBack->getOpcode() != Instruction::Call &&
+              OffsetsVecArg.size() > 1)
+            return ChangeStatus::UNCHANGED;
+
+          for (auto *It = GetBack->op_begin(); It != GetBack->op_end(); It++) {
+            if (Instruction *Ins = dyn_cast<Instruction>(*It)) {
+              if (!Visited[Ins]) {
+                ReadyList.push_back(Ins);
+              }
+            }
+          }
----------------
shiltian wrote:

```suggestion
          for (auto *It = GetBack->op_begin(); It != GetBack->op_end(); It++)
            if (Instruction *Ins = dyn_cast<Instruction>(*It))
              if (!Visited[Ins])
                ReadyList.push_back(Ins);
```

Can you revise your code based on the standard on how to use `{}`?

https://github.com/llvm/llvm-project/pull/72029