[llvm] [Attributor] Fix Load/Store offsets in case of multiple access bins when an allocation size is changed. (PR #72029)

[Vidush Singhal via llvm-commits]

================
@@ -12816,16 +12830,102 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
           new AllocaInst(CharType, AI->getAddressSpace(), NumBytesToValue,
                          AI->getAlign(), AI->getName(), AI->getNextNode());
 
-      if (A.changeAfterManifest(IRPosition::inst(*AI), *NewAllocaInst))
-        return ChangeStatus::CHANGED;
-
+      Changed |= A.changeAfterManifest(IRPosition::inst(*AI), *NewAllocaInst);
       break;
     }
     default:
       break;
     }
 
-    return ChangeStatus::UNCHANGED;
+    const AAPointerInfo *PI =
+        A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::REQUIRED);
+
+    if (!PI)
+      return ChangeStatus::UNCHANGED;
+
+    if (!PI->getState().isValidState())
+      return ChangeStatus::UNCHANGED;
+
+    const auto &NewOffsetsMap = getNewOffsets();
+    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);
+        auto *LocalInst = AccessInstruction.getLocalInst();
+
+        switch (LocalInst->getOpcode()) {
+        case Instruction::Load: {
+          LoadInst *OldLoadInst = cast<LoadInst>(LocalInst);
+          Value *PointerOperand = OldLoadInst->getPointerOperand();
+
+          // We need to shift the old start offset by the difference between the
+          // Old offset and the New offset.
+          Value *ShiftValueForPointerOperand = ConstantInt::get(
+              OldLoadInst->getContext(),
+              APInt(32, (OldOffsetRange.Offset - NewOffsetRange.Offset)));
+          // Cast ptr to 32 bit integer
+          CastInst *PtrOperandToInt = PtrToIntInst::CreatePointerCast(
+              PointerOperand, Type::getInt32Ty(OldLoadInst->getContext()), "",
+              OldLoadInst);
+          // Subtract shift amount from old offset
+          BinaryOperator *NewPointerOperandInt = BinaryOperator::Create(
+              Instruction::Sub, PtrOperandToInt, ShiftValueForPointerOperand,
+              "", OldLoadInst);
+          // Convert Int to ptr.
+          CastInst *IntToPtr = new IntToPtrInst(
+              NewPointerOperandInt, OldLoadInst->getPointerOperandType(), "",
+              OldLoadInst);
----------------
vidsinghal wrote:

It would be easier to use ptrtoint + sub + inttoptr. 
If a GEP is created, we would have to backtrack the load instruction and potentially correct the backtracked instructions for correctness. 
whereas with pointer arithmetic, we only need to change the load if we know the new offsets. 

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