From bea8c444141c9cbbf4a3f6515ce572a28be3c0a1 Mon Sep 17 00:00:00 2001
From: pchintalapudi <34727397+pchintalapudi@users.noreply.github.com>
Date: Wed, 2 Aug 2023 21:51:31 +0000
Subject: [PATCH] Don't partition alwaysinline functions (#50766)

---
 src/aotcompile.cpp | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/aotcompile.cpp b/src/aotcompile.cpp
index c04f1954ce1e4..966f44a8316f1 100644
--- a/src/aotcompile.cpp
+++ b/src/aotcompile.cpp
@@ -723,13 +723,18 @@ static inline bool verify_partitioning(const SmallVectorImpl<Partition> &partiti
             gvars[gvar.second] = i+1;
         }
     }
-    for (auto &GV : M.globals()) {
+    for (auto &GV : M.global_values()) {
         if (GV.isDeclaration()) {
             if (GVNames.count(GV.getName())) {
                 bad = true;
                 dbgs() << "Global " << GV.getName() << " is a declaration but is in partition " << GVNames[GV.getName()] << "\n";
             }
         } else {
+            if (auto F = dyn_cast<Function>(&GV)) {
+                // Ignore alwaysinline functions
+                if (F->hasFnAttribute(Attribute::AlwaysInline))
+                    continue;
+            }
             if (!GVNames.count(GV.getName())) {
                 bad = true;
                 dbgs() << "Global " << GV << " not in any partition\n";
@@ -809,8 +814,12 @@ static SmallVector<Partition, 32> partitionModule(Module &M, unsigned threads) {
     for (auto &G : M.global_values()) {
         if (G.isDeclaration())
             continue;
-        if (isa<Function>(G)) {
-            partitioner.make(&G, getFunctionWeight(cast<Function>(G)).weight);
+        if (auto F = dyn_cast<Function>(&G)) {
+            // alwaysinline functions cannot be partitioned,
+            // they must remain in every module in order to be inlined
+            if (F->hasFnAttribute(Attribute::AlwaysInline))
+                continue;
+            partitioner.make(&G, getFunctionWeight(*F).weight);
         } else {
             partitioner.make(&G, 1);
         }
@@ -1109,6 +1118,12 @@ static void materializePreserved(Module &M, Partition &partition) {
     for (auto &F : M.functions()) {
         if (!F.isDeclaration()) {
             if (!Preserve.contains(&F)) {
+                if (F.hasFnAttribute(Attribute::AlwaysInline)) {
+                    F.setLinkage(GlobalValue::InternalLinkage);
+                    F.setVisibility(GlobalValue::DefaultVisibility);
+                    F.setDSOLocal(true);
+                    continue;
+                }
                 F.deleteBody();
                 F.setLinkage(GlobalValue::ExternalLinkage);
                 F.setVisibility(GlobalValue::HiddenVisibility);