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

[Vidush Singhal via llvm-commits]

================
@@ -6285,12 +6343,106 @@ struct AAAllocationInfo : public StateWrapper<BooleanState, AbstractAttribute> {
     return AbstractAttribute::isValidIRPositionForInit(A, IRP);
   }
 
+  // A helper function to check if simplified values exists for the current
+  // instruction.
+  // TODO: handle case for a similified value
+  // Right now we don't change the value and give up
+  // on modifying the size and offsets of the allocation
+  // this may be sub-optimal
+  bool checkIfSimplifiedValuesExists(Attributor &A, Instruction *LocalInst) {
+
+    // If there are potential values that replace the accessed instruction, we
+    // should use those instead
+    bool UsedAssumedInformation = false;
+    SmallVector<AA::ValueAndContext> Values;
+    if (A.getAssumedSimplifiedValues(IRPosition::inst(*LocalInst), *this,
+                                     Values, AA::AnyScope,
+                                     UsedAssumedInformation))
+
+      for (auto &ValAndContext : Values)
+        // don't modify instruction if any simplified value exists
+        if (ValAndContext.getValue() && ValAndContext.getValue() != LocalInst)
+          return true;
+
+    return false;
+  }
+
+  // A helper function to back track the pointer operand to
+  // a GEP that calculates the pointer operand from the original allocation.
+  // So that we can correct the type of the GEP.
+  // If such a GEP is not found, we return nullptr.
+  Instruction *
+  backTrackPointerOperandToGepFromAllocation(Instruction *PointerOperand,
+                                             Instruction *Allocation) {
+
+    SmallVector<Instruction *> ReadyList;
+    ReadyList.push_back(PointerOperand);
+    while (!ReadyList.empty()) {
+      Instruction *Back = ReadyList.back();
+      ReadyList.pop_back();
+
+      if (Back == Allocation)
+        continue;
+
+      for (auto *It = Back->op_begin(); It != Back->op_end(); It++) {
+        Instruction *OperandInstruction = dyn_cast<Instruction>(It);
+
+        if (!OperandInstruction)
+          continue;
+
+        // This is the GEP instruction whose type and offsets we can change
+        // without any illegal memory access or poison result.
+        if (Back->getOpcode() == Instruction::GetElementPtr &&
+            OperandInstruction == Allocation)
+          return Back;
+
+        ReadyList.push_back(OperandInstruction);
+      }
+    }
+    return nullptr;
+  }
+
+  bool checkIfAccessChainUsesMultipleBins(
+      Attributor &A, Instruction *LocalInst,
+      const DenseMap<Value *, AAPointerInfo::OffsetInfo> &OffsetInfoMap) {
+
+    // BackTrack and check if there are multiple bins for instructions in
+    // the
+    // chain
+    std::vector<Instruction *> ReadyList;
+    DenseMap<Instruction *, bool> Visited;
+    ReadyList.push_back(LocalInst);
+    while (!ReadyList.empty()) {
+      Instruction *GetBack = ReadyList.back();
+      ReadyList.pop_back();
+      // check if the Instruction has multiple bins, if so give up
+      // for calls it is okay to have multiple bins
+      // TODO: handle when one instruction has multiple bins
+      auto OffsetsVecArg = OffsetInfoMap.lookup(GetBack).Offsets;
+      if (GetBack->getOpcode() != Instruction::Call && OffsetsVecArg.size() > 1)
+        return true;
----------------
vidsinghal wrote:

Calls are OK, since it can have different arguments accessing the allocation at different offsets. 
If any one of those arguments accesses multiple offsets that we give up atm. 

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