MemoryBuiltins: accept non-TLI funcs with attribs as allocator funcs

This allows us to accept annotations from out-of-tree languages (the
example test is derived from Rust) so they can enjoy the benefits of
LLVM's optimizations without requiring LLVM to have language-specific
knowledge.

Differential Revision: https://reviews.llvm.org/D123091

llvm/lib/Analysis/MemoryBuiltins.cpp

@@ -540,14 +540,17 @@ Optional<StringRef> llvm::getAllocationFamily(const Value *I,
   if (Callee == nullptr || IsNoBuiltin)
     return None;
   LibFunc TLIFn;
-  if (!TLI || !TLI->getLibFunc(*Callee, TLIFn) || !TLI->has(TLIFn))
-    return None;
-  const auto AllocData = getAllocationDataForFunction(Callee, AnyAlloc, TLI);
-  if (AllocData)
-    return mangledNameForMallocFamily(AllocData.value().Family);
-  const auto FreeData = getFreeFunctionDataForFunction(Callee, TLIFn);
-  if (FreeData)
-    return mangledNameForMallocFamily(FreeData.value().Family);
+
+  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.value().Family);
+    const auto FreeData = getFreeFunctionDataForFunction(Callee, TLIFn);
+    if (FreeData)
+      return mangledNameForMallocFamily(FreeData.value().Family);
+  }
+  // Callee isn't a known library function, still check attributes.
   if (checkFnAllocKind(I, AllocFnKind::Free | AllocFnKind::Alloc |
                               AllocFnKind::Realloc)) {
     Attribute Attr = cast<CallBase>(I)->getFnAttr("alloc-family");

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

@@ -0,0 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=instcombine %s | FileCheck %s
+
+; Function Attrs: nonlazybind uwtable
+define void @alloc_elides_test(i32 %data){
+; CHECK-LABEL: @alloc_elides_test(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    ret void
+;
+start:
+  %0 = tail call i8* @__rust_alloc(i64 4, i64 32)
+  tail call void @__rust_dealloc(i8* nonnull %0, i64 4, i64 32)
+  ret void
+}
+
+declare noalias i8* @__rust_alloc(i64, i64 allocalign) unnamed_addr nounwind nonlazybind allocsize(0) uwtable allockind("alloc,uninitialized,aligned") "alloc-family"="__rust_alloc"
+
+declare void @__rust_dealloc(i8* allocptr, i64, i64) unnamed_addr nounwind nonlazybind uwtable allockind("free") "alloc-family"="__rust_alloc"