[llvm] [mlgo] add 2 new features whether caller/callee is `available_externally` (PR #96585)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/96585

>From 7a9c3f8ce1be5877e56b605e5dd695d767fb7279 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 24 Jun 2024 20:10:22 -0700
Subject: [PATCH 1/2] [mlgo] add 2 new features whether caller/callee is
 `available_externally`

AvailableExternally linkage is interesting because, in ThinLTO cases,
it means the function may get elided if it survives inlining - see
`elim-avail-extern` pass.
---
 .../llvm/Analysis/InlineModelFeatureMaps.h    | 10 ++++++++-
 llvm/lib/Analysis/MLInlineAdvisor.cpp         |  4 ++++
 .../models/gen-inline-oz-test-model.py        |  2 ++
 .../Transforms/Inline/ML/avail-external.ll    | 22 +++++++++++++++++++
 4 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/Inline/ML/avail-external.ll

diff --git a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h
index d62ec9ce49119..acb6c21e48038 100644
--- a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h
+++ b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h
@@ -121,7 +121,15 @@ constexpr bool isHeuristicInlineCostFeature(InlineCostFeatureIndex Feature) {
     "number of blocks reached from a conditional instruction, in the callee")  \
   M(int64_t, {1}, callee_users,                                                \
     "number of module-internal users of the callee, +1 if the callee is "      \
-    "exposed externally")
+    "exposed externally")                                                      \
+  M(int64_t, {1}, is_callee_avail_external,                                    \
+    "Is callee an available-externally linkage type (i.e. could be DCEd if "   \
+    "not "                                                                     \
+    "fully inlined by ElimAvailExtern)")                                       \
+  M(int64_t, {1}, is_caller_avail_external,                                    \
+    "Is caller an available-externally linkage type (i.e. could be DCEd if "   \
+    "not "                                                                     \
+    "fully inlined by ElimAvailExtern)")
 
 // clang-format off
 enum class FeatureIndex : size_t {
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 14eadf6be372a..ec6f3780fe2e4 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -427,6 +427,10 @@ std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   *ModelRunner->getTensor<int64_t>(FeatureIndex::callee_users) =
       CalleeBefore.Uses;
   *ModelRunner->getTensor<int64_t>(FeatureIndex::cost_estimate) = CostEstimate;
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::is_callee_avail_external) =
+      Callee.hasAvailableExternallyLinkage();
+  *ModelRunner->getTensor<int64_t>(FeatureIndex::is_caller_avail_external) =
+      Caller.hasAvailableExternallyLinkage();
 
   // Add the cost features
   for (size_t I = 0;
diff --git a/llvm/lib/Analysis/models/gen-inline-oz-test-model.py b/llvm/lib/Analysis/models/gen-inline-oz-test-model.py
index 3846e7a3cee76..fa20f0ba21110 100644
--- a/llvm/lib/Analysis/models/gen-inline-oz-test-model.py
+++ b/llvm/lib/Analysis/models/gen-inline-oz-test-model.py
@@ -71,6 +71,8 @@ def get_input_signature():
             "nested_inlines",
             "nested_inline_cost_estimate",
             "threshold",
+            "is_callee_avail_external",
+            "is_caller_avail_external"
         ]
     ]
 
diff --git a/llvm/test/Transforms/Inline/ML/avail-external.ll b/llvm/test/Transforms/Inline/ML/avail-external.ll
new file mode 100644
index 0000000000000..1913f096e4031
--- /dev/null
+++ b/llvm/test/Transforms/Inline/ML/avail-external.ll
@@ -0,0 +1,22 @@
+; REQUIRES: x86_64-linux
+; RUN: rm -rf %t.rundir
+; RUN: rm -rf %t.channel-basename.*
+; RUN: mkdir %t.rundir
+; RUN: cp %S/../../../../lib/Analysis/models/log_reader.py %t.rundir
+; RUN: cp %S/../../../../lib/Analysis/models/interactive_host.py %t.rundir
+; RUN: cp %S/Inputs/interactive_main.py %t.rundir
+; RUN: %python %t.rundir/interactive_main.py %t.channel-basename \
+; RUN:    opt -passes=scc-oz-module-inliner -interactive-model-runner-echo-reply \
+; RUN:    -enable-ml-inliner=release -inliner-interactive-channel-base=%t.channel-basename %s -S -o /dev/null | FileCheck %s
+
+define available_externally void @g() {
+  ret void
+}
+
+define void @f(){
+  call void @g()
+  ret void
+}
+
+; CHECK: is_callee_avail_external: 1
+; CHECK: is_caller_avail_external: 0

>From 5bd4d3b39ced88b22a531fc2284bad7e77c01f60 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 24 Jun 2024 20:29:21 -0700
Subject: [PATCH 2/2] fix python

---
 llvm/lib/Analysis/models/gen-inline-oz-test-model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/models/gen-inline-oz-test-model.py b/llvm/lib/Analysis/models/gen-inline-oz-test-model.py
index fa20f0ba21110..c6fcaef7c6ae6 100644
--- a/llvm/lib/Analysis/models/gen-inline-oz-test-model.py
+++ b/llvm/lib/Analysis/models/gen-inline-oz-test-model.py
@@ -72,7 +72,7 @@ def get_input_signature():
             "nested_inline_cost_estimate",
             "threshold",
             "is_callee_avail_external",
-            "is_caller_avail_external"
+            "is_caller_avail_external",
         ]
     ]