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

[Nikita Popov 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);
----------------
nikic wrote:

Use GEP instead of ptrtoint + sub + inttoptr.

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