[NFC][AlwaysInliner] Reduce AlwaysInliner memory consumption. (#96958)

Refactored AlwaysInliner to remove some of inlined functions earlier.

Before the change AlwaysInliner walked through all functions in the
module and inlined them into calls where it is appropriate. Removing of
the dead inlined functions was performed only after all of inlining. For
the test case from the issue
[59126](#59126) compiler
consumes all of the memory on 64GB machine, so is killed.

The change checks if just inlined function can be removed from the
module and removes it.

Author: dfukalov

llvm/lib/Transforms/IPO/AlwaysInliner.cpp

@@ -15,12 +15,12 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/InlineAdvisor.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
-#include "llvm/Transforms/IPO/Inliner.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
@@ -37,86 +37,73 @@ bool AlwaysInlineImpl(
     function_ref<BlockFrequencyInfo &(Function &)> GetBFI) {
   SmallSetVector<CallBase *, 16> Calls;
   bool Changed = false;
-  SmallVector<Function *, 16> InlinedFunctions;
-  for (Function &F : M) {
-    // When callee coroutine function is inlined into caller coroutine function
-    // before coro-split pass,
-    // coro-early pass can not handle this quiet well.
-    // So we won't inline the coroutine function if it have not been unsplited
+  SmallVector<Function *, 16> InlinedComdatFunctions;
+
+  for (Function &F : make_early_inc_range(M)) {
     if (F.isPresplitCoroutine())
       continue;
 
-    if (!F.isDeclaration() && isInlineViable(F).isSuccess()) {
-      Calls.clear();
-
-      for (User *U : F.users())
-        if (auto *CB = dyn_cast<CallBase>(U))
-          if (CB->getCalledFunction() == &F &&
-                CB->hasFnAttr(Attribute::AlwaysInline) &&
-                !CB->getAttributes().hasFnAttr(Attribute::NoInline))
-              Calls.insert(CB);
-
-      for (CallBase *CB : Calls) {
-        Function *Caller = CB->getCaller();
-        OptimizationRemarkEmitter ORE(Caller);
-        DebugLoc DLoc = CB->getDebugLoc();
-        BasicBlock *Block = CB->getParent();
-
-        InlineFunctionInfo IFI(GetAssumptionCache, &PSI,
-                               GetBFI ? &GetBFI(*Caller) : nullptr,
-                               GetBFI ? &GetBFI(F) : nullptr);
-
-        InlineResult Res = InlineFunction(*CB, IFI, /*MergeAttributes=*/true,
-                                          &GetAAR(F), InsertLifetime);
-        if (!Res.isSuccess()) {
-          ORE.emit([&]() {
-            return OptimizationRemarkMissed(DEBUG_TYPE, "NotInlined", DLoc,
-                                            Block)
-                   << "'" << ore::NV("Callee", &F) << "' is not inlined into '"
-                   << ore::NV("Caller", Caller)
-                   << "': " << ore::NV("Reason", Res.getFailureReason());
-          });
-          continue;
-        }
-
-        emitInlinedIntoBasedOnCost(
-            ORE, DLoc, Block, F, *Caller,
-            InlineCost::getAlways("always inline attribute"),
-            /*ForProfileContext=*/false, DEBUG_TYPE);
+    if (F.isDeclaration() || !isInlineViable(F).isSuccess())
+      continue;
 
-        Changed = true;
+    Calls.clear();
+
+    for (User *U : F.users())
+      if (auto *CB = dyn_cast<CallBase>(U))
+        if (CB->getCalledFunction() == &F &&
+            CB->hasFnAttr(Attribute::AlwaysInline) &&
+            !CB->getAttributes().hasFnAttr(Attribute::NoInline))
+          Calls.insert(CB);
+
+    for (CallBase *CB : Calls) {
+      Function *Caller = CB->getCaller();
+      OptimizationRemarkEmitter ORE(Caller);
+      DebugLoc DLoc = CB->getDebugLoc();
+      BasicBlock *Block = CB->getParent();
+
+      InlineFunctionInfo IFI(GetAssumptionCache, &PSI,
+                             GetBFI ? &GetBFI(*Caller) : nullptr,
+                             GetBFI ? &GetBFI(F) : nullptr);
+
+      InlineResult Res = InlineFunction(*CB, IFI, /*MergeAttributes=*/true,
+                                        &GetAAR(F), InsertLifetime);
+      if (!Res.isSuccess()) {
+        ORE.emit([&]() {
+          return OptimizationRemarkMissed(DEBUG_TYPE, "NotInlined", DLoc, Block)
+                 << "'" << ore::NV("Callee", &F) << "' is not inlined into '"
+                 << ore::NV("Caller", Caller)
+                 << "': " << ore::NV("Reason", Res.getFailureReason());
+        });
+        continue;
       }
 
-      if (F.hasFnAttribute(Attribute::AlwaysInline)) {
-        // Remember to try and delete this function afterward. This both avoids
-        // re-walking the rest of the module and avoids dealing with any
-        // iterator invalidation issues while deleting functions.
-        InlinedFunctions.push_back(&F);
-      }
+      emitInlinedIntoBasedOnCost(
+          ORE, DLoc, Block, F, *Caller,
+          InlineCost::getAlways("always inline attribute"),
+          /*ForProfileContext=*/false, DEBUG_TYPE);
+
+      Changed = true;
     }
-  }
 
-  // Remove any live functions.
-  erase_if(InlinedFunctions, [&](Function *F) {
-    F->removeDeadConstantUsers();
-    return !F->isDefTriviallyDead();
-  });
-
-  // Delete the non-comdat ones from the module and also from our vector.
-  auto NonComdatBegin = partition(
-      InlinedFunctions, [&](Function *F) { return F->hasComdat(); });
-  for (Function *F : make_range(NonComdatBegin, InlinedFunctions.end())) {
-    M.getFunctionList().erase(F);
-    Changed = true;
+    F.removeDeadConstantUsers();
+    if (F.hasFnAttribute(Attribute::AlwaysInline) && F.isDefTriviallyDead()) {
+      // Remember to try and delete this function afterward. This allows to call
+      // filterDeadComdatFunctions() only once.
+      if (F.hasComdat()) {
+        InlinedComdatFunctions.push_back(&F);
+      } else {
+        M.getFunctionList().erase(F);
+        Changed = true;
+      }
+    }
   }
-  InlinedFunctions.erase(NonComdatBegin, InlinedFunctions.end());
 
-  if (!InlinedFunctions.empty()) {
+  if (!InlinedComdatFunctions.empty()) {
     // Now we just have the comdat functions. Filter out the ones whose comdats
     // are not actually dead.
-    filterDeadComdatFunctions(InlinedFunctions);
+    filterDeadComdatFunctions(InlinedComdatFunctions);
     // The remaining functions are actually dead.
-    for (Function *F : InlinedFunctions) {
+    for (Function *F : InlinedComdatFunctions) {
       M.getFunctionList().erase(F);
       Changed = true;
     }

llvm/test/Transforms/Inline/always-inline-phase-ordering.ll

@@ -6,7 +6,6 @@ target triple = "arm64e-apple-macosx13"
 ; CHECK: remark: <unknown>:0:0: 'wibble' inlined into 'pluto' with (cost=always): always inline attribute
 ; CHECK: remark: <unknown>:0:0: 'snork' inlined into 'blam' with (cost=always): always inline attribute
 ; CHECK: remark: <unknown>:0:0: 'wobble' inlined into 'blam' with (cost=always): always inline attribute
-; CHECK: remark: <unknown>:0:0: 'wobble' inlined into 'snork' with (cost=always): always inline attribute
 ; CHECK: remark: <unknown>:0:0: 'spam' inlined into 'blam' with (cost=65, threshold=75)
 ; CHECK: remark: <unknown>:0:0: 'wibble.1' inlined into 'widget' with (cost=30, threshold=75)
 ; CHECK: remark: <unknown>:0:0: 'widget' inlined into 'bar.8' with (cost=30, threshold=75)