[Attributor] Fix Load/Store Offsets if multiple bins are present for a pointer allocation.

Author: vidsinghal

llvm/include/llvm/Transforms/IPO/Attributor.h

@@ -5143,9 +5143,7 @@ struct DenormalFPMathState : public AbstractState {
       return Mode != Other.Mode || ModeF32 != Other.ModeF32;
     }
 
-    bool isValid() const {
-      return Mode.isValid() && ModeF32.isValid();
-    }
+    bool isValid() const { return Mode.isValid() && ModeF32.isValid(); }
 
     static DenormalMode::DenormalModeKind
     unionDenormalKind(DenormalMode::DenormalModeKind Callee,
@@ -5185,9 +5183,7 @@ struct DenormalFPMathState : public AbstractState {
   // state.
   DenormalState getAssumed() const { return Known; }
 
-  bool isValidState() const override {
-    return Known.isValid();
-  }
+  bool isValidState() const override { return Known.isValid(); }
 
   /// Return true if there are no dynamic components to the denormal mode worth
   /// specializing.
@@ -5198,9 +5194,7 @@ struct DenormalFPMathState : public AbstractState {
            Known.ModeF32.Output != DenormalMode::Dynamic;
   }
 
-  bool isAtFixpoint() const override {
-    return IsAtFixedpoint;
-  }
+  bool isAtFixpoint() const override { return IsAtFixedpoint; }
 
   ChangeStatus indicateFixpoint() {
     bool Changed = !IsAtFixedpoint;
@@ -6122,6 +6116,8 @@ struct AAPointerInfo : public AbstractAttribute {
   virtual const_bin_iterator begin() const = 0;
   virtual const_bin_iterator end() const = 0;
   virtual int64_t numOffsetBins() const = 0;
+  virtual void dumpState(raw_ostream &O) const = 0;
+  virtual const Access &getBinAccess(unsigned Index) const = 0;
 
   /// Call \p CB on all accesses that might interfere with \p Range and return
   /// true if all such accesses were known and the callback returned true for
@@ -6293,6 +6289,9 @@ struct AAAllocationInfo : public StateWrapper<BooleanState, AbstractAttribute> {
 
   virtual std::optional<TypeSize> getAllocatedSize() const = 0;
 
+  using NewOffsetsTy = DenseMap<AA::RangeTy, AA::RangeTy>;
+  virtual const NewOffsetsTy &getNewOffsets() const = 0;
+
   /// See AbstractAttribute::getName()
   const std::string getName() const override { return "AAAllocationInfo"; }
 

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -419,7 +419,8 @@ struct AAReturnedFromReturnedValues : public BaseType {
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override {
     StateType S(StateType::getBestState(this->getState()));
-    clampReturnedValueStates<AAType, StateType, IRAttributeKind, RecurseForSelectAndPHI>(
+    clampReturnedValueStates<AAType, StateType, IRAttributeKind,
+                             RecurseForSelectAndPHI>(
         A, *this, S,
         PropagateCallBaseContext ? this->getCallBaseContext() : nullptr);
     // TODO: If we know we visited all returned values, thus no are assumed
@@ -1083,6 +1084,10 @@ struct AAPointerInfoImpl
     return State::numOffsetBins();
   }
 
+  virtual const Access &getBinAccess(unsigned Index) const override {
+    return getAccess(Index);
+  }
+
   bool forallInterferingAccesses(
       AA::RangeTy Range,
       function_ref<bool(const AAPointerInfo::Access &, bool)> CB)
@@ -1429,7 +1434,7 @@ struct AAPointerInfoImpl
   void trackPointerInfoStatistics(const IRPosition &IRP) const {}
 
   /// Dump the state into \p O.
-  void dumpState(raw_ostream &O) {
+  virtual void dumpState(raw_ostream &O) const override {
     for (auto &It : OffsetBins) {
       O << "[" << It.first.Offset << "-" << It.first.Offset + It.first.Size
         << "] : " << It.getSecond().size() << "\n";
@@ -6969,10 +6974,9 @@ ChangeStatus AAHeapToStackFunction::updateImpl(Attributor &A) {
     if (AI.LibraryFunctionId != LibFunc___kmpc_alloc_shared) {
       Instruction *CtxI = isa<InvokeInst>(AI.CB) ? AI.CB : AI.CB->getNextNode();
       if (!Explorer || !Explorer->findInContextOf(UniqueFree, CtxI)) {
-        LLVM_DEBUG(
-            dbgs()
-            << "[H2S] unique free call might not be executed with the allocation "
-            << *UniqueFree << "\n");
+        LLVM_DEBUG(dbgs() << "[H2S] unique free call might not be executed "
+                             "with the allocation "
+                          << *UniqueFree << "\n");
         return false;
       }
     }
@@ -10425,11 +10429,12 @@ struct AANoFPClassFloating : public AANoFPClassImpl {
 
 struct AANoFPClassReturned final
     : AAReturnedFromReturnedValues<AANoFPClass, AANoFPClassImpl,
-                                   AANoFPClassImpl::StateType, false, Attribute::None, false> {
+                                   AANoFPClassImpl::StateType, false,
+                                   Attribute::None, false> {
   AANoFPClassReturned(const IRPosition &IRP, Attributor &A)
       : AAReturnedFromReturnedValues<AANoFPClass, AANoFPClassImpl,
-                                     AANoFPClassImpl::StateType, false, Attribute::None, false>(
-            IRP, A) {}
+                                     AANoFPClassImpl::StateType, false,
+                                     Attribute::None, false>(IRP, A) {}
 
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override {
@@ -12686,6 +12691,11 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
     return AssumedAllocatedSize;
   }
 
+  const NewOffsetsTy &getNewOffsets() const override {
+    assert(isValidState() && "the AA is invalid");
+    return NewComputedOffsets;
+  }
+
   std::optional<TypeSize> findInitialAllocationSize(Instruction *I,
                                                     const DataLayout &DL) {
 
@@ -12735,37 +12745,42 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
     if (*AllocationSize == 0)
       return indicatePessimisticFixpoint();
 
-    int64_t BinSize = PI->numOffsetBins();
-
-    // TODO: implement for multiple bins
-    if (BinSize > 1)
-      return indicatePessimisticFixpoint();
+    int64_t NumBins = PI->numOffsetBins();
 
-    if (BinSize == 0) {
+    if (NumBins == 0) {
       auto 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();
+    // For each access bin
+    // Compute its new start Offset and store the results in a new map
+    // (NewOffsetBins).
+    int64_t PrevBinEndOffset = 0;
+    bool ChangedOffsets = false;
+    for (AAPointerInfo::OffsetBinsTy::const_iterator It = PI->begin();
+         It != PI->end(); It++) {
+      const AA::RangeTy &OldRange = It->getFirst();
+      int64_t NewStartOffset = PrevBinEndOffset;
+      int64_t NewEndOffset = NewStartOffset + OldRange.Size;
+      PrevBinEndOffset = NewEndOffset;
 
-    // TODO: handle if Offset is not zero
-    if (It->first.Offset != 0)
-      return indicatePessimisticFixpoint();
-
-    uint64_t SizeOfBin = It->first.Offset + It->first.Size;
-
-    if (SizeOfBin >= *AllocationSize)
-      return indicatePessimisticFixpoint();
+      ChangedOffsets |= setNewOffsets(OldRange, OldRange.Offset, NewStartOffset,
+                                      OldRange.Size);
+    }
 
+    // Set the new size of the allocation, the new size of the Allocation should
+    // be the size of NewEndOffset * 8, in bits.
     auto NewAllocationSize =
-        std::optional<TypeSize>(TypeSize(SizeOfBin * 8, false));
+        std::optional<TypeSize>(TypeSize(PrevBinEndOffset * 8, false));
 
     if (!changeAllocationSize(NewAllocationSize))
       return ChangeStatus::UNCHANGED;
 
+    if (!ChangedOffsets)
+      return ChangeStatus::UNCHANGED;
+
     return ChangeStatus::CHANGED;
   }
 
@@ -12775,12 +12790,13 @@ 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();
 
     unsigned long NumBytesToAllocate = (FixedAllocatedSizeInBits + 7) / 8;
-
     switch (I->getOpcode()) {
     // TODO: add case for malloc like calls
     case Instruction::Alloca: {
@@ -12789,25 +12805,100 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
 
       Type *CharType = Type::getInt8Ty(I->getContext());
 
-      auto *NumBytesToValue =
-          ConstantInt::get(I->getContext(), APInt(32, NumBytesToAllocate));
+      Type *CharArrayType = ArrayType::get(CharType, NumBytesToAllocate);
 
       BasicBlock::iterator insertPt = AI->getIterator();
       insertPt = std::next(insertPt);
-      AllocaInst *NewAllocaInst =
-          new AllocaInst(CharType, AI->getAddressSpace(), NumBytesToValue,
-                         AI->getAlign(), AI->getName(), insertPt);
-
-      if (A.changeAfterManifest(IRPosition::inst(*AI), *NewAllocaInst))
-        return ChangeStatus::CHANGED;
+      AllocaInst *NewAllocaInst = new AllocaInst(
+          CharArrayType, AI->getAddressSpace(), AI->getName(), insertPt);
 
+      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);
+      int64_t ShiftValue = NewOffsetRange.Offset - OldOffsetRange.Offset;
+      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();
+          Type *PointeeTy = OldLoadInst->getPointerOperandType();
+
+          IntegerType *Int64TyInteger =
+              IntegerType::get(LocalInst->getContext(), 64);
+
+          Value *IndexList[1] = {ConstantInt::get(Int64TyInteger, ShiftValue)};
+          Value *GepToNewAddress = GetElementPtrInst::Create(
+              PointeeTy, PointerOperand, IndexList, "NewGep", OldLoadInst);
+
+          LoadInst *NewLoadInst = new LoadInst(
+              OldLoadInst->getType(), GepToNewAddress, OldLoadInst->getName(),
+              false, OldLoadInst->getAlign(), OldLoadInst);
+
+          Changed |= A.changeAfterManifest(IRPosition::inst(*OldLoadInst),
+                                           *NewLoadInst);
+
+          A.deleteAfterManifest(*OldLoadInst);
+          break;
+        }
+        case Instruction::Store: {
+          StoreInst *OldStoreInst = cast<StoreInst>(LocalInst);
+          Value *PointerOperand = OldStoreInst->getPointerOperand();
+          Type *PointeeTy = OldStoreInst->getPointerOperandType();
+
+          IntegerType *Int64TyInteger =
+              IntegerType::get(LocalInst->getContext(), 64);
+
+          Value *IndexList[1] = {ConstantInt::get(Int64TyInteger, ShiftValue)};
+          Value *GepToNewAddress = GetElementPtrInst::Create(
+              PointeeTy, PointerOperand, IndexList, "NewGep", OldStoreInst);
+
+          StoreInst *NewStoreInst =
+              new StoreInst(OldStoreInst->getValueOperand(), GepToNewAddress,
+                            false, OldStoreInst->getAlign(), OldStoreInst);
+
+          Changed |= A.changeAfterManifest(IRPosition::inst(*OldStoreInst),
+                                           *NewStoreInst);
+
+          A.deleteAfterManifest(*OldStoreInst);
+          break;
+        }
+        }
+      }
+    }
+
+    if (!Changed)
+      return ChangeStatus::UNCHANGED;
+    return ChangeStatus::CHANGED;
   }
 
   /// See AbstractAttribute::getAsStr().
@@ -12821,8 +12912,28 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
            ")";
   }
 
+  void dumpNewOffsetBins(raw_ostream &O) {
+
+    O << "Printing Map from [OldOffsetsRange] : [NewOffsetsRange] if the "
+         "offsets changed."
+      << "\n";
+    const auto &NewOffsetsMap = getNewOffsets();
+    for (auto It = NewOffsetsMap.begin(); It != NewOffsetsMap.end(); It++) {
+
+      const auto &OldRange = It->getFirst();
+      const auto &NewRange = It->getSecond();
+
+      O << "[" << OldRange.Offset << "," << OldRange.Offset + OldRange.Size
+        << "] : ";
+      O << "[" << NewRange.Offset << "," << NewRange.Offset + NewRange.Size
+        << "]";
+      O << "\n";
+    }
+  }
+
 private:
   std::optional<TypeSize> AssumedAllocatedSize = HasNoAllocationSize;
+  NewOffsetsTy NewComputedOffsets;
 
   // Maintain the computed allocation size of the object.
   // Returns (bool) weather the size of the allocation was modified or not.
@@ -12834,6 +12945,21 @@ struct AAAllocationInfoImpl : public AAAllocationInfo {
     }
     return false;
   }
+
+  // Maps an old byte range to its new Offset range in the new allocation.
+  // Returns (bool) weather the old byte range's offsets changed or not.
+  bool setNewOffsets(const AA::RangeTy &OldRange, int64_t OldOffset,
+                     int64_t NewComputedOffset, int64_t Size) {
+
+    if (OldOffset == NewComputedOffset)
+      return false;
+
+    AA::RangeTy &NewRange = NewComputedOffsets.getOrInsertDefault(OldRange);
+    NewRange.Offset = NewComputedOffset;
+    NewRange.Size = Size;
+
+    return true;
+  }
 };
 
 struct AAAllocationInfoFloating : AAAllocationInfoImpl {

llvm/test/Other/ChangePrinters/DotCfg/print-changed-dot-cfg.ll

@@ -1,3 +1,4 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 ; Simple checks of -print-changed=dot-cfg
 ;
 ; Note that (mostly) only the banners are checked.

llvm/test/Other/ChangePrinters/DotCfg/print-changed-dot-cfg.mir

@@ -1,3 +1,4 @@
+# NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 # REQUIRES: x86-registered-target
 # Simple functionality check.
 # RUN: rm -rf %t && mkdir -p %t
@@ -8,17 +9,17 @@
 name:            g
 body:             |
   bb.0.entry:
-    %0:gr32 = MOV32ri 5
-    $eax = COPY %0
-    RET 0, $eax
+  %0:gr32 = MOV32ri 5
+  $eax = COPY %0
+  RET 0, $eax
 
 ...
 ---
 name:            f
 body:             |
   bb.0.entry:
-    %0:gr32 = MOV32ri 7
-    $eax = COPY %0
-    RET 0, $eax
+  %0:gr32 = MOV32ri 7
+  $eax = COPY %0
+  RET 0, $eax
 
 ...