[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:

Using ptrtoint + sub + inttoptr means that LLVM will give up on any further optimization of these accesses. It will assume that everything can alias everything. You need to use GEP to preserve provenance information.

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