[llvm] [Attributor] Change allocation size and load/store offsets using AAPointerInfo for Alloca instructions and keep track of instructions causing an Access (PR #72029)

[Shilei Tian via llvm-commits]

================
@@ -13495,39 +13828,329 @@ 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());
+    assert(PI->getState().isValidState() &&
+           "[AAAllocationinfo]: AAPointerinfo was not in valid state!");
 
-      auto *NumBytesToValue =
-          ConstantInt::get(I->getContext(), APInt(32, NumBytesToAllocate));
+    // Store a map where each instruction is mapped to a map containing
+    // old bins accessed by that instruction to the corresponding new
+    // bins in the allocation.
+    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) {
+          auto [It, Inserted] = Map.try_emplace(LocalInst);
+          DenseMap<AA::RangeTy, AA::RangeTy> &NewBinsForInstruction =
+              It->second;
 
-      if (A.changeAfterManifest(IRPosition::inst(*AI), *NewAllocaInst))
-        return ChangeStatus::CHANGED;
+          NewBinsForInstruction.insert(std::make_pair(OldRange, NewRange));
+        };
+
+    const NewOffsetsTy &NewOffsetsMap = getNewOffsets();
+    const AAPointerInfo::OffsetInfoMapTy &OffsetInfoMap =
+        PI->getOffsetInfoMap();
+
+    // Map access causing instructions to a 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 AA::RangeTy &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 AA::RangeTy &NewOffsetRange = NewOffsetsMap.lookup(OldOffsetRange);
+      for (const unsigned int &AccIndex : It->getSecond()) {
+        const auto &AccessInstruction = PI->getBinAccess(AccIndex);
+        Instruction *LocalInst = AccessInstruction.getLocalInst();
+
+        if (checkIfSimplifiedValuesExists(A, LocalInst))
+          return ChangeStatus::UNCHANGED;
+
+        if (checkIfAccessChainUsesMultipleBins(A, LocalInst, OffsetInfoMap))
+          return ChangeStatus::UNCHANGED;
+
+        // Check if we can backtrack the access causing instruction to a GEP
+        // from the original allocation, if yes, then we prefer to change the
+        // GEP rather than inserting a new one just before the access causing
+        // instruction.
+        switch (LocalInst->getOpcode()) {
+        case Instruction::Call: {
+          CallInst *CallInstruction = cast<CallInst>(LocalInst);
+          for (Use *It = CallInstruction->op_begin();
+               It != CallInstruction->op_end(); It++) {
+            if (Instruction *OperandInstruction = dyn_cast<Instruction>(It)) {
----------------
shiltian wrote:

Similar here

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