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

[Vidush Singhal via llvm-commits]

================
@@ -12752,39 +12732,284 @@ 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();
 
-    auto FixedAllocatedSizeInBits = getAllocatedSize()->getFixedValue();
+    // check if simplified values exist
+    if (checkIfSimplifiedValuesExists(A, I))
+      return ChangeStatus::UNCHANGED;
 
-    unsigned long NumBytesToAllocate = (FixedAllocatedSizeInBits + 7) / 8;
+    if (getAllocatedSize() == HasNoAllocationSize)
+      return ChangeStatus::UNCHANGED;
 
-    switch (I->getOpcode()) {
-    // TODO: add case for malloc like calls
-    case Instruction::Alloca: {
+    const AAPointerInfo *PI =
+        A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::REQUIRED);
 
-      AllocaInst *AI = cast<AllocaInst>(I);
+    if (!PI)
+      return ChangeStatus::UNCHANGED;
 
-      Type *CharType = Type::getInt8Ty(I->getContext());
+    if (!PI->getState().isValidState())
+      return ChangeStatus::UNCHANGED;
 
-      auto *NumBytesToValue =
-          ConstantInt::get(I->getContext(), APInt(32, NumBytesToAllocate));
+    // Store a map where each instruction maps to a set of bins accessed by that
+    // instruction
+    DenseMap<Instruction *, DenseMap<AA::RangeTy, AA::RangeTy>>
+        AccessedInstructionsToBinsMap;
 
-      BasicBlock::iterator insertPt = AI->getIterator();
-      insertPt = std::next(insertPt);
-      AllocaInst *NewAllocaInst =
-          new AllocaInst(CharType, AI->getAddressSpace(), NumBytesToValue,
-                         AI->getAlign(), AI->getName(), insertPt);
+    auto AddBins =
+        [](DenseMap<Instruction *, DenseMap<AA::RangeTy, AA::RangeTy>> &Map,
+           Instruction *LocalInst, const AA::RangeTy &OldRange,
+           const AA::RangeTy &NewRange) {
+          DenseMap<AA::RangeTy, AA::RangeTy> &NewBinsForInstruction =
+              Map.getOrInsertDefault(LocalInst);
 
-      if (A.changeAfterManifest(IRPosition::inst(*AI), *NewAllocaInst))
-        return ChangeStatus::CHANGED;
+          NewBinsForInstruction.insert(std::make_pair(OldRange, NewRange));
+        };
+
+    const auto &NewOffsetsMap = getNewOffsets();
+    const auto &OffsetInfoMap = PI->getOffsetInfoMap();
+
+    // Map access causing instructions to tuple of (Old, New) bins.
+    // The access causing instruction contains the pointer operand
+    // which comes from the allocation
+    // We may want to backtrack that pointer operand, there are 2 cases that may
+    // arise:
+    // A) A GEP exists that calculates the pointer operand from the original
+    // allocation instruction: I
+    // B) A GEP does not exists
+    // in which case we need to insert a GEP just before the access causing
+    // instruction with the shift value from the original offset.
+    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;
----------------
vidsinghal wrote:

The NewOffsetsMap is a Map of type 
`DenseMap<AA::RangeTy, AA::RangeTy>`

It maps the old byte range to its new byte range in the allocation. 
In the example you pointed out, the 2 offset ranges are [4:8) and [8:12)
If they are moved by 4, their new offsets are [0:4), [4:8)

The map will store them as two different entries {[4:8) : [0:4) , [8:12) : [4:8)}
In that for loop we are going over all the original access causing offsets and checking if they exist in the map. It they do not exist in the map, we continue. In this case since both old offsets will exist in the map, we will not skip them and correctly change their offsets. 









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