[mlgo] add 2 new features whether caller/callee is available_externally

[mtrofin]

AvailableExternally linkage is interesting because, in ThinLTO cases, it means the function may get elided if it survives inlining - see elim-avail-extern pass.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis
@llvm/pr-subscribers-mlgo

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

AvailableExternally linkage is interesting because, in ThinLTO cases, it means the function may get elided if it survives inlining - see elim-avail-extern pass.

---

Full diff: https://github.com/llvm/llvm-project/pull/96585.diff

4 Files Affected:

• (modified) llvm/include/llvm/Analysis/InlineModelFeatureMaps.h (+9-1)
• (modified) llvm/lib/Analysis/MLInlineAdvisor.cpp (+4)
• (modified) llvm/lib/Analysis/models/gen-inline-oz-test-model.py (+2)
• (added) llvm/test/Transforms/Inline/ML/avail-external.ll (+22)

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

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.

[boomanaiden154]

LGTM.