diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -119,6 +119,12 @@
 /// Gets the alignment argument for an aligned_alloc-like function
 Value *getAllocAlignment(const CallBase *V, const TargetLibraryInfo *TLI);
 
+/// Returns the size arguments for a allocation function. If there is only one
+/// size argument, the second pointer is null. getObjectSize is preferred if
+/// only constant sizes are needed
+std::pair<Value *, Value *> getAllocSizeArgs(const CallBase *V,
+                                             const TargetLibraryInfo *TLI);
+
 /// If this allocation function initializes memory to a fixed value, return
 /// said value in the requested type.  Otherwise, return nullptr.
 Constant *getInitialValueOfAllocation(const CallBase *Alloc,
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -309,6 +309,24 @@
   return V->getOperand(FnData->AlignParam);
 }
 
+std::pair<Value *, Value *>
+llvm::getAllocSizeArgs(const CallBase *V, const TargetLibraryInfo *TLI) {
+  assert((isMallocOrCallocLikeFn(V, TLI) || isReallocLikeFn(V, TLI)) &&
+         "getAllocSizeArgs can only be called on an allocation function that "
+         "isn't like strdup");
+  const Optional<AllocFnsTy> FnData = getAllocationSize(V, TLI);
+  assert(FnData.hasValue());
+  if (FnData->FstParam < 0) {
+    return std::pair<Value *, Value *>(nullptr, nullptr);
+  }
+  if (FnData->SndParam < 0) {
+    return std::pair<Value *, Value *>(V->getOperand(FnData->FstParam),
+                                       nullptr);
+  }
+  return std::pair<Value *, Value *>(V->getOperand(FnData->FstParam),
+                                     V->getOperand(FnData->SndParam));
+}
+
 Constant *llvm::getInitialValueOfAllocation(const CallBase *Alloc,
                                             const TargetLibraryInfo *TLI,
                                             Type *Ty) {
diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -5771,13 +5771,6 @@
     /// The call that allocates the memory.
     CallBase *const CB;
 
-    /// The kind of allocation.
-    const enum class AllocationKind {
-      MALLOC,
-      CALLOC,
-      ALIGNED_ALLOC,
-    } Kind;
-
     /// The library function id for the allocation.
     LibFunc LibraryFunctionId = NotLibFunc;
 
@@ -5834,18 +5827,16 @@
         DeallocationInfos[CB] = new (A.Allocator) DeallocationInfo{CB};
         return true;
       }
-      bool IsMalloc = isMallocLikeFn(CB, TLI);
-      bool IsAlignedAllocLike = !IsMalloc && isAlignedAllocLikeFn(CB, TLI);
-      bool IsCalloc =
-          !IsMalloc && !IsAlignedAllocLike && isCallocLikeFn(CB, TLI);
-      if (!IsMalloc && !IsAlignedAllocLike && !IsCalloc)
+
+      if (!(isMallocOrCallocLikeFn(CB, TLI) && isAllocRemovable(CB, TLI)))
         return true;
-      auto Kind =
-          IsMalloc ? AllocationInfo::AllocationKind::MALLOC
-                   : (IsCalloc ? AllocationInfo::AllocationKind::CALLOC
-                               : AllocationInfo::AllocationKind::ALIGNED_ALLOC);
 
-      AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB, Kind};
+      auto *AllocValue = getInitialValueOfAllocation(CB, TLI, CB->getType());
+      if (!AllocValue) {
+        return true; // Allocation has unknown but defined initial value
+      }
+
+      AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB};
       AllocationInfos[CB] = AI;
       TLI->getLibFunc(*CB, AI->LibraryFunctionId);
       return true;
@@ -6017,17 +6008,24 @@
 
   Optional<APInt> getSize(Attributor &A, const AbstractAttribute &AA,
                           AllocationInfo &AI) {
+    Function *F = getAnchorScope();
+    const auto *TLI = A.getInfoCache().getTargetLibraryInfoForFunction(*F);
 
-    if (AI.Kind == AllocationInfo::AllocationKind::MALLOC)
-      return getAPInt(A, AA, *AI.CB->getArgOperand(0));
+    // Only if the alignment is also constant we return a size.
+    Value *Align;
+    if ((Align = getAllocAlignment(AI.CB, TLI)) &&
+        !getAPInt(A, AA, *Align).hasValue()) {
+      return llvm::None;
+    }
 
-    if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC)
-      return getAPInt(A, AA, *AI.CB->getArgOperand(1));
+    auto Args = getAllocSizeArgs(AI.CB, TLI);
+
+    if (!Args.second) {
+      return getAPInt(A, AA, *Args.first);
+    }
 
-    assert(AI.Kind == AllocationInfo::AllocationKind::CALLOC &&
-           "Expected only callocs are left");
-    Optional<APInt> Num = getAPInt(A, AA, *AI.CB->getArgOperand(0));
-    Optional<APInt> Size = getAPInt(A, AA, *AI.CB->getArgOperand(1));
+    Optional<APInt> Num = getAPInt(A, AA, *Args.first);
+    Optional<APInt> Size = getAPInt(A, AA, *Args.second);
     if (!Num.hasValue() || !Size.hasValue())
       return llvm::None;
     bool Overflow = false;