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

[Shilei Tian via llvm-commits]

================
@@ -12752,39 +12734,315 @@ 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();
+
+    // check if simplified values exist
+    if (checkIfSimplifiedValuesExists(A, I))
+      return ChangeStatus::UNCHANGED;
+
+    if (getAllocatedSize() == HasNoAllocationSize)
+      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;
+
+    // Store a map where each instruction maps to a set of bins accessed by that
+    // instruction
+    DenseMap<Instruction *, DenseMap<AA::RangeTy, AA::RangeTy>>
+        AccessedInstructionsToBinsMap;
+
+    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);
+
+          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;
+
+      const auto &NewOffsetRange = NewOffsetsMap.lookup(OldOffsetRange);
+
+      for (const auto 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 the access causing instruction
+        switch (LocalInst->getOpcode()) {
+        case Instruction::Load: {
+          LoadInst *Load = cast<LoadInst>(LocalInst);
+          Instruction *PointerOperand =
+              cast<Instruction>(Load->getPointerOperand());
+          Instruction *BackTrackedGEP =
+              backTrackPointerOperandToGepFromAllocation(PointerOperand, I);
+
+          if (!BackTrackedGEP) {
+            AddBins(AccessedInstructionsToBinsMap, LocalInst, OldOffsetRange,
+                    NewOffsetRange);
+            break;
+          }
+
+          if (checkIfSimplifiedValuesExists(A, BackTrackedGEP))
+            return ChangeStatus::UNCHANGED;
+
+          AddBins(AccessedInstructionsToBinsMap, BackTrackedGEP, OldOffsetRange,
+                  NewOffsetRange);
+          break;
+        }
+        case Instruction::Store: {
+          StoreInst *Store = cast<StoreInst>(LocalInst);
+          Instruction *PointerOperand =
+              cast<Instruction>(Store->getPointerOperand());
+          Instruction *BackTrackedGEP =
+              backTrackPointerOperandToGepFromAllocation(PointerOperand, I);
+
+          if (!BackTrackedGEP) {
+            AddBins(AccessedInstructionsToBinsMap, LocalInst, OldOffsetRange,
+                    NewOffsetRange);
+            break;
+          }
+
+          if (checkIfSimplifiedValuesExists(A, BackTrackedGEP))
+            return ChangeStatus::UNCHANGED;
+
+          AddBins(AccessedInstructionsToBinsMap, BackTrackedGEP, OldOffsetRange,
+                  NewOffsetRange);
+          break;
+        }
+        case Instruction::Call: {
+          CallInst *CallInstruction = cast<CallInst>(LocalInst);
+
+          for (auto *It = CallInstruction->op_begin();
+               It != CallInstruction->op_end(); It++) {
 
----------------
shiltian wrote:

Is it necessary to have an empty line at the beginning of each kind of block?

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