[MemoryBuiltins] Handle allocator attributes on call-site

We should handle allocator attributes not only on function
declarations, but also on the call-site. That way we can e.g.
also optimize cases where the allocator function is a virtual
function call.

This was already supported in some of the MemoryBuiltins helpers,
but not all of them. This adds support for allocsize, alloc-family
and allockind("free").

llvm/lib/Analysis/MemoryBuiltins.cpp

@@ -240,21 +240,18 @@ getAllocationData(const Value *V, AllocType AllocTy,
 }
 
 static std::optional<AllocFnsTy>
-getAllocationSize(const Value *V, const TargetLibraryInfo *TLI) {
+getAllocationSize(const CallBase *CB, const TargetLibraryInfo *TLI) {
   bool IsNoBuiltinCall;
-  const Function *Callee =
-      getCalledFunction(V, IsNoBuiltinCall);
-  if (!Callee)
-    return std::nullopt;
-
-  // Prefer to use existing information over allocsize. This will give us an
-  // accurate AllocTy.
-  if (!IsNoBuiltinCall)
+  const Function *Callee = getCalledFunction(CB, IsNoBuiltinCall);
+  if (Callee && !IsNoBuiltinCall) {
+    // Prefer to use existing information over allocsize. This will give us an
+    // accurate AllocTy.
     if (std::optional<AllocFnsTy> Data =
             getAllocationDataForFunction(Callee, AnyAlloc, TLI))
       return Data;
+  }
 
-  Attribute Attr = Callee->getFnAttribute(Attribute::AllocSize);
+  Attribute Attr = CB->getFnAttr(Attribute::AllocSize);
   if (Attr == Attribute())
     return std::nullopt;
 
@@ -264,7 +261,7 @@ getAllocationSize(const Value *V, const TargetLibraryInfo *TLI) {
   // Because allocsize only tells us how many bytes are allocated, we're not
   // really allowed to assume anything, so we use MallocLike.
   Result.AllocTy = MallocLike;
-  Result.NumParams = Callee->getNumOperands();
+  Result.NumParams = CB->arg_size();
   Result.FstParam = Args.first;
   Result.SndParam = Args.second.value_or(-1);
   // Allocsize has no way to specify an alignment argument
@@ -512,19 +509,20 @@ std::optional<StringRef>
 llvm::getAllocationFamily(const Value *I, const TargetLibraryInfo *TLI) {
   bool IsNoBuiltin;
   const Function *Callee = getCalledFunction(I, IsNoBuiltin);
-  if (Callee == nullptr || IsNoBuiltin)
-    return std::nullopt;
-  LibFunc TLIFn;
-
-  if (TLI && TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn)) {
-    // Callee is some known library function.
-    const auto AllocData = getAllocationDataForFunction(Callee, AnyAlloc, TLI);
-    if (AllocData)
-      return mangledNameForMallocFamily(AllocData->Family);
-    const auto FreeData = getFreeFunctionDataForFunction(Callee, TLIFn);
-    if (FreeData)
-      return mangledNameForMallocFamily(FreeData->Family);
+  if (Callee && !IsNoBuiltin) {
+    LibFunc TLIFn;
+    if (TLI && TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn)) {
+      // Callee is some known library function.
+      const auto AllocData =
+          getAllocationDataForFunction(Callee, AnyAlloc, TLI);
+      if (AllocData)
+        return mangledNameForMallocFamily(AllocData->Family);
+      const auto FreeData = getFreeFunctionDataForFunction(Callee, TLIFn);
+      if (FreeData)
+        return mangledNameForMallocFamily(FreeData->Family);
+    }
   }
+
   // Callee isn't a known library function, still check attributes.
   if (checkFnAllocKind(I, AllocFnKind::Free | AllocFnKind::Alloc |
                               AllocFnKind::Realloc)) {
@@ -558,14 +556,13 @@ bool llvm::isLibFreeFunction(const Function *F, const LibFunc TLIFn) {
 Value *llvm::getFreedOperand(const CallBase *CB, const TargetLibraryInfo *TLI) {
   bool IsNoBuiltinCall;
   const Function *Callee = getCalledFunction(CB, IsNoBuiltinCall);
-  if (Callee == nullptr || IsNoBuiltinCall)
-    return nullptr;
-
-  LibFunc TLIFn;
-  if (TLI && TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn) &&
-      isLibFreeFunction(Callee, TLIFn)) {
-    // All currently supported free functions free the first argument.
-    return CB->getArgOperand(0);
+  if (Callee && !IsNoBuiltinCall) {
+    LibFunc TLIFn;
+    if (TLI && TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn) &&
+        isLibFreeFunction(Callee, TLIFn)) {
+      // All currently supported free functions free the first argument.
+      return CB->getArgOperand(0);
+    }
   }
 
   if (checkFnAllocKind(CB, AllocFnKind::Free))

llvm/test/Transforms/InstCombine/deref-alloc-fns.ll

@@ -372,3 +372,12 @@ define ptr @my_calloc_constant_size() {
   %call = call ptr @my_calloc(i64 32, i64 4)
   ret ptr %call
 }
+
+define ptr @virtual_constant_size(ptr %alloc) {
+; CHECK-LABEL: @virtual_constant_size(
+; CHECK-NEXT:    [[CALL:%.*]] = call dereferenceable_or_null(16) ptr [[ALLOC:%.*]](i64 16) #[[ATTR5:[0-9]+]]
+; CHECK-NEXT:    ret ptr [[CALL]]
+;
+  %call = call ptr %alloc(i64 16) allocsize(0)
+  ret ptr %call
+}

llvm/test/Transforms/InstCombine/out-of-tree-allocator-optimizes-away.ll

@@ -15,6 +15,20 @@ start:
   ret void
 }
 
+define void @alloc_elides_test_virtual(i32 %data, ptr %alloc, ptr %realloc, ptr %dealloc) {
+; CHECK-LABEL: @alloc_elides_test_virtual(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    ret void
+;
+start:
+  %a = call noalias ptr %alloc(i64 4, i64 allocalign 32) nounwind allocsize(0) allockind("alloc,uninitialized,aligned") "alloc-family"="__rust_alloc"
+  store i32 0, ptr %a
+  %a2 = call noalias ptr %realloc(ptr allocptr %a, i64 4, i64 allocalign 32, i64 8) nounwind allocsize(3) allockind("alloc,uninitialized,aligned") "alloc-family"="__rust_alloc"
+  store i32 1, ptr %a2
+  call void %dealloc(ptr allocptr %a2, i64 4, i64 32) nounwind allockind("free") "alloc-family"="__rust_alloc"
+  ret void
+}
+
 declare noalias ptr @__rust_alloc(i64, i64 allocalign) nounwind allocsize(0) allockind("alloc,uninitialized,aligned") "alloc-family"="__rust_alloc"
 
 declare noalias ptr @__rust_realloc(ptr allocptr, i64, i64 allocalign, i64) nounwind allocsize(3) allockind("alloc,uninitialized,aligned") "alloc-family"="__rust_alloc"