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

================
@@ -13451,41 +13701,124 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
     if (!AllocationSize)
       return indicatePessimisticFixpoint();
 
-    // For zero sized allocations, we give up.
-    // Since we can't reduce further
+    // For zero sized allocations, we give up
+    // because we cannot reduce them any further.
     if (*AllocationSize == 0)
       return indicatePessimisticFixpoint();
 
-    int64_t BinSize = PI->numOffsetBins();
-
-    // TODO: implement for multiple bins
-    if (BinSize > 1)
-      return indicatePessimisticFixpoint();
-
-    if (BinSize == 0) {
-      auto NewAllocationSize = std::make_optional<TypeSize>(0, false);
+    int64_t NumBins = PI->numOffsetBins();
+    if (NumBins == 0) {
+      std::optional<TypeSize> NewAllocationSize =
+          std::optional<TypeSize>(TypeSize(0, false));
       if (!changeAllocationSize(NewAllocationSize))
         return ChangeStatus::UNCHANGED;
       return ChangeStatus::CHANGED;
     }
 
-    // TODO: refactor this to be part of multiple bin case
-    const auto &It = PI->begin();
+    // Algorithm:
+    // For all the accessed ranges in AAPointerInfo, we need to find the minimum
+    // longest range that contains all the accessed offsets.
+    // In case, disjoint ranges exist, we want to merge them since we want to
+    // make a packed Alloca. For this we need to calculate the minimum size of
+    // the new allocation. Then adjust all the old offsets and map them to their
+    // new offsets.
+
+    // A tuple to store a cluster, a cluster is a maximal unique range.
+    // Different Clusters are meant to be disjoint, but we eventually merge them
+    // together. A bigger cluster can subsume a smaller cluster inside it.
+    using ClusterTy = std::tuple<int64_t, int64_t>;
+    using ClustersTy = SmallVector<ClusterTy, 4>;
+
+    // BinInterval is an interval to keep track of the Start and
+    // End of a RangeTy. Since RangeTy stores the Offset and the
+    // Corresponding size but not the end of the range.
+    //(StartOffset, EndOffset, RangeTy struct)
+    using AccessedInterval = std::tuple<int64_t, int64_t, AA::RangeTy>;
+    using AccessedIntervals = SmallVector<AccessedInterval, 4>;
+
+    // Obtain all the offset bins that exists in AAPointerInfo.
+    SmallVector<AA::RangeTy, 4> OldBins;
+    for (AAPointerInfo::OffsetBinsTy::const_iterator It = PI->begin();
+         It != PI->end(); It++) {
+      const AA::RangeTy &Bin = It->getFirst();
+      // In case unknown or unassigned bins exists,
+      // We don't want to change the allocation size and return a
+      // pessimistic fixpoint.
+      if (Bin.offsetOrSizeAreUnknown() || Bin.isUnassigned())
+        return indicatePessimisticFixpoint();
+      OldBins.push_back(Bin);
+    }
 
-    // TODO: handle if Offset is not zero
-    if (It->first.Offset != 0)
-      return indicatePessimisticFixpoint();
+    AccessedIntervals Intervals;
+    // Obtain the intervals from the ranges.
+    for (AA::RangeTy &OldBin : OldBins) {
+      AccessedInterval Interval =
+          std::make_tuple(OldBin.Offset, OldBin.Offset + OldBin.Size, OldBin);
+      Intervals.push_back(Interval);
+    }
 
-    uint64_t SizeOfBin = It->first.Offset + It->first.Size;
+    // sort the intervals to push the largest interval to the beginning.
+    llvm::sort(Intervals,
+               [](const AccessedInterval &A, const AccessedInterval &B) {
+                 return get<0>(A) < get<0>(B);
+               });
 
-    if (SizeOfBin >= *AllocationSize)
-      return indicatePessimisticFixpoint();
+    ClustersTy Clusters;
+    int64_t CurrentClusterStart = get<0>(Intervals[0]);
+    int64_t CurrentClusterEnd = get<1>(Intervals[0]);
 
-    auto NewAllocationSize = std::make_optional<TypeSize>(SizeOfBin * 8, false);
+    for (size_t I = 1; I < Intervals.size(); I++) {
+      if (get<0>(Intervals[I]) <= CurrentClusterEnd) {
+        CurrentClusterEnd = std::max(CurrentClusterEnd, get<1>(Intervals[I]));
+      } else {
+        Clusters.push_back({CurrentClusterStart, CurrentClusterEnd});
+        CurrentClusterStart = get<0>(Intervals[I]);
+        CurrentClusterEnd = get<1>(Intervals[I]);
+      }
+    }
+    Clusters.push_back({CurrentClusterStart, CurrentClusterEnd});
+
+    int64_t PackedRangeSize = 0;
+    // Sum up each cluster to get the total size of the packed Alloca.
+    for (ClusterTy &Cluster : Clusters)
+      PackedRangeSize += (get<1>(Cluster) - get<0>(Cluster));
+
+    // For each access bin we compute its new start offset
+    // and store the results in a new map (NewOffsetBins).
+    // NewOffsetsBins is a Map from AA::RangeTy OldRange to AA::RangeTy
+    // NewRange
+    int64_t PackedCursorStart = 0;
+    bool ChangedOffsets = false;
+    for (ClusterTy &Cluster : Clusters) {
+      const long &ClusterStart = get<0>(Cluster);
+      const long &ClusterEnd = get<1>(Cluster);
+      for (AccessedInterval &Interval : Intervals) {
+        const long &InterValStart = get<0>(Interval);
+        const long &InterValEnd = get<1>(Interval);
+        if (InterValStart >= ClusterStart && InterValEnd <= ClusterEnd) {
+          int64_t NewRangeStart =
+              PackedCursorStart + (InterValStart - ClusterStart);
+          int64_t NewRangeSize = InterValEnd - InterValStart;
+          const AA::RangeTy &OldRange = get<2>(Interval);
+
+          ChangedOffsets |= setNewOffsets(OldRange, OldRange.Offset,
+                                          NewRangeStart, NewRangeSize);
+        }
+      }
+      PackedCursorStart += ClusterEnd - ClusterStart;
+    }
+
+    // Set the new size of the allocation. The new size of the Allocation should
+    // be the size of PrevBinEndOffset * 8 in bits.
+    std::optional<TypeSize> NewAllocationSize =
----------------
shiltian wrote:

Same here

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