[llvm] [TypeMetadataUtil][CallPromtionUtil]Add utility function getFunctionAtVTableOffset that finds functions through alias. Use it in CallPromotionUtils which didn't promote aliasee previously (PR #80736)

[Mingming Liu via llvm-commits]

https://github.com/minglotus-6 created https://github.com/llvm/llvm-project/pull/80736

- The utility function is extracted from the implementation of [DevirtModule::tryFindVirtualCallTargets](https://github.com/llvm/llvm-project/blob/b99163fe8feeacba7797d5479bbcd5d8f327dd2d/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp#L1069-L1082)

>From f1721444c762e8d1494da789d61e621ea7367e19 Mon Sep 17 00:00:00 2001
From: mingmingl <mingmingl at google.com>
Date: Mon, 5 Feb 2024 11:45:36 -0800
Subject: [PATCH] [TypeMetadataUtil][CallPromtionUtil]Add utility function
 getFunctionAtVTableOffset that finds functions through alias. Use it in
 CallPromotionUtils which didn't promote aliasee previously

- The utility function is extracted from the implementation of
  'DevirtModule::tryFindVirtualCallTargets'.
---
 llvm/include/llvm/Analysis/TypeMetadataUtils.h  |  5 +++++
 llvm/lib/Analysis/TypeMetadataUtils.cpp         | 17 +++++++++++++++++
 llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp  | 10 ++++------
 .../lib/Transforms/Utils/CallPromotionUtils.cpp | 11 ++++-------
 4 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/Analysis/TypeMetadataUtils.h b/llvm/include/llvm/Analysis/TypeMetadataUtils.h
index dab67aad1ab0e..9f8c364b49375 100644
--- a/llvm/include/llvm/Analysis/TypeMetadataUtils.h
+++ b/llvm/include/llvm/Analysis/TypeMetadataUtils.h
@@ -14,6 +14,7 @@
 #ifndef LLVM_ANALYSIS_TYPEMETADATAUTILS_H
 #define LLVM_ANALYSIS_TYPEMETADATAUTILS_H
 
+#include "llvm/IR/GlobalVariable.h"
 #include <cstdint>
 
 namespace llvm {
@@ -77,6 +78,10 @@ void findDevirtualizableCallsForTypeCheckedLoad(
 Constant *getPointerAtOffset(Constant *I, uint64_t Offset, Module &M,
                              Constant *TopLevelGlobal = nullptr);
 
+// Given a vtable, returns the function pointer specified by Offset.
+std::pair<Function *, Constant *>
+getFunctionAtVTableOffset(GlobalVariable *GV, uint64_t Offset, Module &M);
+
 /// Finds the same "relative pointer" pattern as described above, where the
 /// target is `F`, and replaces the entire pattern with a constant zero.
 void replaceRelativePointerUsersWithZero(Function *F);
diff --git a/llvm/lib/Analysis/TypeMetadataUtils.cpp b/llvm/lib/Analysis/TypeMetadataUtils.cpp
index bbaee06ed8a55..0e6859f622a7d 100644
--- a/llvm/lib/Analysis/TypeMetadataUtils.cpp
+++ b/llvm/lib/Analysis/TypeMetadataUtils.cpp
@@ -201,6 +201,23 @@ Constant *llvm::getPointerAtOffset(Constant *I, uint64_t Offset, Module &M,
   return nullptr;
 }
 
+std::pair<Function *, Constant *>
+llvm::getFunctionAtVTableOffset(GlobalVariable *GV, uint64_t Offset,
+                                Module &M) {
+  Constant *Ptr = getPointerAtOffset(GV->getInitializer(), Offset, M, GV);
+  if (!Ptr)
+    return std::pair<Function *, Constant *>(nullptr, nullptr);
+
+  auto C = Ptr->stripPointerCasts();
+  // Make sure this is a function or alias to a function.
+  auto Fn = dyn_cast<Function>(C);
+  auto A = dyn_cast<GlobalAlias>(C);
+  if (!Fn && A)
+    Fn = dyn_cast<Function>(A->getAliasee());
+
+  return std::pair<Function *, Constant *>(Fn, C);
+}
+
 void llvm::replaceRelativePointerUsersWithZero(Function *F) {
   for (auto *U : F->users()) {
     auto *PtrExpr = dyn_cast<ConstantExpr>(U);
diff --git a/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp b/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
index 01aba47cdbfff..154ff876a53cd 100644
--- a/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
+++ b/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
@@ -1071,12 +1071,10 @@ bool DevirtModule::tryFindVirtualCallTargets(
     if (!Ptr)
       return false;
 
-    auto C = Ptr->stripPointerCasts();
-    // Make sure this is a function or alias to a function.
-    auto Fn = dyn_cast<Function>(C);
-    auto A = dyn_cast<GlobalAlias>(C);
-    if (!Fn && A)
-      Fn = dyn_cast<Function>(A->getAliasee());
+    Function *Fn = nullptr;
+    Constant *C = nullptr;
+    std::tie(Fn, C) =
+        getFunctionAtVTableOffset(TM.Bits->GV, TM.Offset + ByteOffset, M);
 
     if (!Fn)
       return false;
diff --git a/llvm/lib/Transforms/Utils/CallPromotionUtils.cpp b/llvm/lib/Transforms/Utils/CallPromotionUtils.cpp
index e42cdab64446e..4e84927f1cfc9 100644
--- a/llvm/lib/Transforms/Utils/CallPromotionUtils.cpp
+++ b/llvm/lib/Transforms/Utils/CallPromotionUtils.cpp
@@ -597,16 +597,13 @@ bool llvm::tryPromoteCall(CallBase &CB) {
     // Not in the form of a global constant variable with an initializer.
     return false;
 
-  Constant *VTableGVInitializer = GV->getInitializer();
   APInt VTableGVOffset = VTableOffsetGVBase + VTableOffset;
   if (!(VTableGVOffset.getActiveBits() <= 64))
     return false; // Out of range.
-  Constant *Ptr = getPointerAtOffset(VTableGVInitializer,
-                                     VTableGVOffset.getZExtValue(),
-                                     *M);
-  if (!Ptr)
-    return false; // No constant (function) pointer found.
-  Function *DirectCallee = dyn_cast<Function>(Ptr->stripPointerCasts());
+
+  Function *DirectCallee = nullptr;
+  std::tie(DirectCallee, std::ignore) =
+      getFunctionAtVTableOffset(GV, VTableGVOffset.getZExtValue(), *M);
   if (!DirectCallee)
     return false; // No function pointer found.