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

[via llvm-commits]

Author: Mircea Trofin
Date: 2024-06-25T12:36:40-07:00
New Revision: 600ff287722a15106e9a02c470b9865dda05980e

URL: https://github.com/llvm/llvm-project/commit/600ff287722a15106e9a02c470b9865dda05980e
DIFF: https://github.com/llvm/llvm-project/commit/600ff287722a15106e9a02c470b9865dda05980e.diff

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

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

Added: 
    llvm/test/Transforms/Inline/ML/avail-external.ll

Modified: 
    llvm/include/llvm/Analysis/InlineModelFeatureMaps.h
    llvm/lib/Analysis/MLInlineAdvisor.cpp
    llvm/lib/Analysis/models/gen-inline-oz-test-model.py

Removed: 
    


################################################################################
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..c6fcaef7c6ae6 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