[llvm] [llvm][GlobalOpt] Optimize statically resolvable IFuncs (PR #80606)

[Jon Roelofs via llvm-commits]

https://github.com/jroelofs updated https://github.com/llvm/llvm-project/pull/80606

>From 5318a4dd324a3e74130c7bd9f0b0031e1740d847 Mon Sep 17 00:00:00 2001
From: Jon Roelofs <jonathan_roelofs at apple.com>
Date: Sun, 4 Feb 2024 08:41:11 -0800
Subject: [PATCH] [llvm][GlobalOpt] Optimize statically resolvable IFuncs

---
 llvm/lib/Transforms/IPO/GlobalOpt.cpp         | 40 +++++++++
 .../GlobalOpt/resolve-static-ifunc.ll         | 81 +++++++++++++++++++
 2 files changed, 121 insertions(+)
 create mode 100644 llvm/test/Transforms/GlobalOpt/resolve-static-ifunc.ll

diff --git a/llvm/lib/Transforms/IPO/GlobalOpt.cpp b/llvm/lib/Transforms/IPO/GlobalOpt.cpp
index 951372adcfa93..32f23ffcefb70 100644
--- a/llvm/lib/Transforms/IPO/GlobalOpt.cpp
+++ b/llvm/lib/Transforms/IPO/GlobalOpt.cpp
@@ -89,6 +89,7 @@ STATISTIC(NumAliasesRemoved, "Number of global aliases eliminated");
 STATISTIC(NumCXXDtorsRemoved, "Number of global C++ destructors removed");
 STATISTIC(NumInternalFunc, "Number of internal functions");
 STATISTIC(NumColdCC, "Number of functions marked coldcc");
+STATISTIC(NumIFuncsResolved, "Number of statically resolved IFuncs");
 
 static cl::opt<bool>
     EnableColdCCStressTest("enable-coldcc-stress-test",
@@ -2404,6 +2405,42 @@ static bool OptimizeEmptyGlobalCXXDtors(Function *CXAAtExitFn) {
   return Changed;
 }
 
+static Function *hasSideeffectFreeStaticResolution(GlobalIFunc &IF) {
+  Function *Resolver = IF.getResolverFunction();
+  if (!Resolver)
+    return nullptr;
+
+  Function *Callee = nullptr;
+  for (BasicBlock &BB : *Resolver) {
+    if (any_of(BB, [](Instruction &I) { return I.mayHaveSideEffects(); }))
+      return nullptr;
+
+    if (auto *Ret = dyn_cast<ReturnInst>(BB.getTerminator()))
+      if (auto *F = dyn_cast<Function>(Ret->getReturnValue())) {
+        if (!Callee)
+          Callee = F;
+        else if (F != Callee)
+          return nullptr;
+      }
+  }
+
+  return Callee;
+}
+
+/// Find IFuncs that have resolvers that always point at the same statically
+/// known callee, and replace their callers with a direct call.
+static bool OptimizeStaticIFuncs(Module &M) {
+  bool Changed = false;
+  for (GlobalIFunc &IF : M.ifuncs())
+    if (Function *Callee = hasSideeffectFreeStaticResolution(IF))
+      if (!IF.use_empty()) {
+        IF.replaceAllUsesWith(Callee);
+        NumIFuncsResolved++;
+        Changed = true;
+      }
+  return Changed;
+}
+
 static bool
 optimizeGlobalsInModule(Module &M, const DataLayout &DL,
                         function_ref<TargetLibraryInfo &(Function &)> GetTLI,
@@ -2464,6 +2501,9 @@ optimizeGlobalsInModule(Module &M, const DataLayout &DL,
     if (CXAAtExitFn)
       LocalChange |= OptimizeEmptyGlobalCXXDtors(CXAAtExitFn);
 
+    // Optimize IFuncs whose callee's are statically known.
+    LocalChange |= OptimizeStaticIFuncs(M);
+
     Changed |= LocalChange;
   }
 
diff --git a/llvm/test/Transforms/GlobalOpt/resolve-static-ifunc.ll b/llvm/test/Transforms/GlobalOpt/resolve-static-ifunc.ll
new file mode 100644
index 0000000000000..d51e9111ad2cd
--- /dev/null
+++ b/llvm/test/Transforms/GlobalOpt/resolve-static-ifunc.ll
@@ -0,0 +1,81 @@
+; RUN: opt --passes=globalopt -o - -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+$callee_with_trivial_resolver.resolver = comdat any
+ at callee_with_trivial_resolver.ifunc = weak_odr dso_local ifunc void (), ptr @callee_with_trivial_resolver.resolver
+define weak_odr ptr @callee_with_trivial_resolver.resolver() comdat {
+  ret ptr @callee_with_trivial_resolver._Msimd
+}
+define void @callee_with_trivial_resolver._Msimd() {
+  ret void
+}
+define void @callee_with_trivial_resolver.default() {
+  ret void
+}
+
+ at unknown_condition = external global i1
+$callee_with_complex_static_resolver.resolver = comdat any
+ at callee_with_complex_static_resolver.ifunc = weak_odr dso_local ifunc void (), ptr @callee_with_complex_static_resolver.resolver
+define weak_odr ptr @callee_with_complex_static_resolver.resolver() comdat {
+entry:
+  %v = load i1, ptr @unknown_condition
+  br i1 %v, label %fast, label %slow
+fast:
+  ret ptr @callee_with_complex_static_resolver._Msimd
+slow:
+  ret ptr @callee_with_complex_static_resolver._Msimd
+}
+define void @callee_with_complex_static_resolver._Msimd() {
+  ret void
+}
+define void @callee_with_complex_static_resolver.default() {
+  ret void
+}
+
+$callee_with_complex_dynamic_resolver.resolver = comdat any
+ at callee_with_complex_dynamic_resolver.ifunc = weak_odr dso_local ifunc void (), ptr @callee_with_complex_dynamic_resolver.resolver
+define weak_odr ptr @callee_with_complex_dynamic_resolver.resolver() comdat {
+entry:
+  %v = load i1, ptr @unknown_condition
+  br i1 %v, label %fast, label %slow
+fast:
+  ret ptr @callee_with_complex_dynamic_resolver._Msimd
+slow:
+  ret ptr @callee_with_complex_dynamic_resolver.default
+}
+define void @callee_with_complex_dynamic_resolver._Msimd() {
+  ret void
+}
+define void @callee_with_complex_dynamic_resolver.default() {
+  ret void
+}
+
+$callee_with_sideeffects_resolver.resolver = comdat any
+ at callee_with_sideeffects_resolver.ifunc = weak_odr dso_local ifunc void (), ptr @callee_with_sideeffects_resolver.resolver
+define weak_odr ptr @callee_with_sideeffects_resolver.resolver() comdat {
+  store i1 0, ptr @unknown_condition
+  ret ptr @callee_with_sideeffects_resolver.default
+}
+define void @callee_with_sideeffects_resolver._Msimd() {
+  ret void
+}
+define void @callee_with_sideeffects_resolver.default() {
+  ret void
+}
+
+define void @caller() {
+  call void @callee_with_trivial_resolver.ifunc()
+  call void @callee_with_complex_static_resolver.ifunc()
+  call void @callee_with_complex_dynamic_resolver.ifunc()
+  call void @callee_with_sideeffects_resolver.ifunc()
+  ret void
+}
+
+; CHECK-LABEL: define void @caller()
+; CHECK-NEXT:     call void @callee_with_trivial_resolver._Msimd()
+; CHECK-NEXT:     call void @callee_with_complex_static_resolver._Msimd()
+; CHECK-NEXT:     call void @callee_with_complex_dynamic_resolver.ifunc()
+; CHECK-NEXT:     call void @callee_with_sideeffects_resolver.ifunc()
+; CHECK-NEXT:     ret void