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

[Johannes Doerfert via llvm-commits]

================
@@ -12658,6 +12669,268 @@ 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;
+  }
+
+  bool isaMallocInst(Instruction *I) {
+    CallInst *Call = dyn_cast<CallInst>(I);
+    auto CallName = Call->getCalledFunction()->getName();
+    if (CallName.str() == "malloc")
+      return true;
+
+    return false;
+  };
+
+  ChangeStatus updateImpl(Attributor &A) override {
+
+    Instruction *I = getIRPosition().getCtxI();
+
+    const IRPosition &IRP = getIRPosition();
+
+    if (!(isa<AllocaInst>(I) || isaMallocInst(I)))
+      return llvm::IntegerStateBase<bool, true,
+                                    false>::indicatePessimisticFixpoint();
+
+    const AAPointerInfo *PI =
+        A.getOrCreateAAFor<AAPointerInfo>(IRP, *this, DepClassTy::REQUIRED);
+
+    if (!PI)
+      return indicatePessimisticFixpoint();
+
+    if (!PI->getState().isValidState())
+      return indicatePessimisticFixpoint();
+
+    int64_t BinSize = PI->numOffsetBins();
+
+    // TODO: Handle for more than one Bin
+    if (BinSize > 1)
+      return indicatePessimisticFixpoint();
+
+    const auto &It = PI->begin();
+
+    if (BinSize == 0)
+      return indicatePessimisticFixpoint();
+
+    if (It->first.Offset != 0)
+      return indicatePessimisticFixpoint();
+
+    uint64_t OffsetEnd = It->first.Offset + It->first.Size;
+    const DataLayout &DL = A.getDataLayout();
+
+    switch (I->getOpcode()) {
+    case Instruction::Alloca: {
+      AllocaInst *AI = dyn_cast<AllocaInst>(getIRPosition().getCtxI());
+      if (!AI)
+        return indicatePessimisticFixpoint();
+      const auto &AllocationSize = AI->getAllocationSize(DL);
+
+      if (!AllocationSize || AllocationSize == 0)
+        return indicatePessimisticFixpoint();
+
+      if (OffsetEnd == *AllocationSize)
+        return indicatePessimisticFixpoint();
+
+      break;
+    }
+    case Instruction::Call: {
+
+      if (!isaMallocInst(I))
+        return llvm::IntegerStateBase<bool, true,
+                                      false>::indicatePessimisticFixpoint();
+
+      Value *ValueOperand = I->getOperand(0);
+
+      if (!ValueOperand)
+        return llvm::IntegerStateBase<bool, true,
+                                      false>::indicatePessimisticFixpoint();
+
+      ConstantInt *IntOperand = dyn_cast<ConstantInt>(ValueOperand);
+
+      if (!IntOperand) {
+        const AAPotentialConstantValues *PotentialConstant =
+            A.getOrCreateAAFor<AAPotentialConstantValues>(
+                IRPosition::value(*ValueOperand), *this, DepClassTy::REQUIRED);
+
+        if (!PotentialConstant || !PotentialConstant->isValidState())
+          return llvm::IntegerStateBase<bool, true,
+                                        false>::indicatePessimisticFixpoint();
+
+        Value *GetConstantAsValue =
+            PotentialConstant->getAssumedConstant(A).value();
+
+        IntOperand = dyn_cast<ConstantInt>(GetConstantAsValue);
+
+        if (!IntOperand)
+          return llvm::IntegerStateBase<bool, true,
+                                        false>::indicatePessimisticFixpoint();
+      }
+
+      if (OffsetEnd == IntOperand->getZExtValue())
+        return llvm::IntegerStateBase<bool, true,
+                                      false>::indicatePessimisticFixpoint();
+
+      break;
+    }
+    default:
+      return indicatePessimisticFixpoint();
+    }
+
+    auto SizeOfTypeInBits =
+        std::optional<TypeSize>(TypeSize(OffsetEnd * 8, false));
+
+    if (!changeAllocationSize(SizeOfTypeInBits))
+      return ChangeStatus::UNCHANGED;
+
+    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()) {
+    case Instruction::Alloca: {
+
+      AllocaInst *AI = dyn_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;
+    }
+    case Instruction::Call: {
+
+      if (!isaMallocInst(I))
+        return llvm::IntegerStateBase<bool, true,
+                                      false>::indicatePessimisticFixpoint();
+
+      Type *IntPtrTy = IntegerType::getInt32Ty(I->getContext());
+      Instruction *NewMallocInst = CallInst::CreateMalloc(
+          I->getNextNode(), IntPtrTy, CharType, NumBytesToValue, nullptr,
+          nullptr, I->getName());
+
+      if (A.changeAfterManifest(IRPosition::inst(*I), *NewMallocInst)) {
+        A.deleteAfterManifest(*I);
+        return ChangeStatus::CHANGED;
+      }
+
+      break;
+    }
+    default:
+      break;
+    }
----------------
jdoerfert wrote:

See the comments on the switch above.

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