[llvm] [Attributor] New attribute to identify what byte ranges are alive for an allocation (PR #66148)

[Vidhush Singhal via llvm-commits]

================
@@ -12658,6 +12665,232 @@ struct AAAddressSpaceCallSiteArgument final : AAAddressSpaceImpl {
 };
 } // namespace
 
+/// ----------- Allocation Info ----------
+namespace {
+struct AAAllocationInfoImpl : public AAAllocationInfo {
+  AAAllocationInfoImpl(const IRPosition &IRP, Attributor &A)
+      : AAAllocationInfo(IRP, A) {}
+
+  std::optional<TypeSize> getAllocatedSize() const override {
+    assert(isValidState() && "the AA is invalid");
+    return AssumedAllocatedSize;
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+
+    const IRPosition &IRP = getIRPosition();
+    Instruction *I = IRP.getCtxI();
+
+    /*TODO: update check for malloc like calls.*/
+    if (!isa<AllocaInst>(I))
+      return indicatePessimisticFixpoint();
+
+    bool IsKnownNoCapture;
+    if (!AA::hasAssumedIRAttr<Attribute::NoCapture>(
+            A, this, IRP, DepClassTy::OPTIONAL, IsKnownNoCapture))
+      return indicatePessimisticFixpoint();
+
+    if (IsKnownNoCapture)
+      return indicatePessimisticFixpoint();
+
+    const AAPointerInfo *PI =
+        A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::REQUIRED);
+
+    if (!PI)
+      return indicatePessimisticFixpoint();
+
+    if (!PI->getState().isValidState())
+      return indicatePessimisticFixpoint();
+
+    const DataLayout &DL = A.getDataLayout();
+    int64_t BinSize = PI->numOffsetBins();
+    switch (BinSize) {
+    case 0: {
+      switch (I->getOpcode()) {
+      /*TODO: add case for malloc like calls*/
+      case Instruction::Alloca: {
+        AllocaInst *AI = cast<AllocaInst>(I);
+        const auto AllocationSize = AI->getAllocationSize(DL);
+
+        if (!AllocationSize || *AllocationSize == 0)
+          return indicatePessimisticFixpoint();
+
+        break;
+      }
+      default:
+        return indicatePessimisticFixpoint();
+      }
+
+      auto NewAllocationSize = std::optional<TypeSize>(TypeSize(0, false));
+
+      if (!changeAllocationSize(NewAllocationSize))
+        return ChangeStatus::UNCHANGED;
+      break;
+    }
+    case 1: {
+      const auto &It = PI->begin();
+      if (It->first.Offset == 0) {
+
+        uint64_t SizeOfBin = It->first.Offset + It->first.Size;
+        /*TODO: add case for malloc like calls*/
+        switch (I->getOpcode()) {
+        case Instruction::Alloca: {
+          AllocaInst *AI = cast<AllocaInst>(I);
+          const auto AllocationSize = AI->getAllocationSize(DL);
+
+          if (!AllocationSize || *AllocationSize == 0)
+            return indicatePessimisticFixpoint();
+
+          if (SizeOfBin == *AllocationSize)
+            return indicatePessimisticFixpoint();
+
+          break;
+        }
+        default:
+          return indicatePessimisticFixpoint();
+        }
+
+        auto NewAllocationSize =
+            std::optional<TypeSize>(TypeSize(SizeOfBin * 8, false));
+
+        if (!changeAllocationSize(NewAllocationSize))
+          return ChangeStatus::UNCHANGED;
+      } else {
+        /*TODO: when access does not start at the 0th byte of the bin*/
+        return indicatePessimisticFixpoint();
+      }
+      break;
+    }
+    default: {
+      /*TODO: Handle for multiple Bins*/
+      return indicatePessimisticFixpoint();
+    }
+    }
+
+    return ChangeStatus::CHANGED;
+  }
+
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+
+    assert(isValidState() &&
+           "Manifest should only be called if the state is valid.");
+
+    Instruction *I = getIRPosition().getCtxI();
+
+    auto FixedAllocatedSizeInBits = getAllocatedSize()->getFixedValue();
+
+    int NumBytesToAllocate = (FixedAllocatedSizeInBits + 7) / 8;
+
+    Type *CharType = Type::getInt8Ty(I->getContext());
+
+    auto *NumBytesToValue = llvm::ConstantInt::get(
+        I->getContext(), llvm::APInt(32, NumBytesToAllocate));
+
+    switch (I->getOpcode()) {
+    /*TODO: add case for malloc like calls*/
+    case Instruction::Alloca: {
+
+      AllocaInst *AI = cast<AllocaInst>(I);
+      AllocaInst *NewAllocaInst =
+          new AllocaInst(CharType, AI->getAddressSpace(), NumBytesToValue,
+                         AI->getAlign(), AI->getName(), AI->getNextNode());
+
+      if (A.changeAfterManifest(IRPosition::inst(*AI), *NewAllocaInst))
+        return ChangeStatus::CHANGED;
+
+      break;
+    }
+    default:
+      break;
+    }
+
+    return llvm::ChangeStatus::UNCHANGED;
+  }
+
+  /// See AbstractAttribute::getAsStr().
+  const std::string getAsStr(Attributor *A) const override {
+    if (!isValidState())
+      return "allocationinfo(<invalid>)";
+    return "allocationinfo(" +
+           (AssumedAllocatedSize == HasNoAllocationSize
+                ? "none"
+                : std::to_string(AssumedAllocatedSize->getFixedValue())) +
+           ")";
+  }
+
+private:
+  std::optional<TypeSize> AssumedAllocatedSize = HasNoAllocationSize;
+
+  bool changeAllocationSize(std::optional<TypeSize> Size) {
+    if (AssumedAllocatedSize == HasNoAllocationSize ||
+        AssumedAllocatedSize != Size) {
----------------
vidsinghal wrote:

this only checks if the allocation size changed or not. 
The return value is a bool signifying if the allocation size changed or not changed. 

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