[llvm-branch-commits] [llvm] fe69360 - Revert "[LLVM] Fix use-after-free in AlwaysInliner flatten worklist (#194485)"

[via llvm-branch-commits]

Author: Alexander Kornienko
Date: 2026-05-11T15:35:14+02:00
New Revision: fe6936030b320f34252f5d253b397da2bf538c34

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

LOG: Revert "[LLVM] Fix use-after-free in AlwaysInliner flatten worklist (#194485)"

This reverts commit b40c1d511b2e84842707939a1332b90ebb1a50a0.

Added: 
    

Modified: 
    llvm/lib/Transforms/IPO/AlwaysInliner.cpp
    llvm/test/Transforms/Inline/flatten.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/AlwaysInliner.cpp b/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
index bdba6a1dbca53..080cb8ddb33fd 100644
--- a/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
+++ b/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
@@ -42,6 +42,7 @@ bool AlwaysInlineImpl(
   SmallSetVector<CallBase *, 16> Calls;
   bool Changed = false;
   SmallVector<Function *, 16> InlinedComdatFunctions;
+  SmallVector<Function *, 4> NeedFlattening;
 
   auto TryInline = [&](CallBase &CB, Function &Callee,
                        OptimizationRemarkEmitter &ORE, const char *InlineReason,
@@ -75,16 +76,52 @@ bool AlwaysInlineImpl(
     return true;
   };
 
-  for (Function &F : M) {
-    if (!F.hasFnAttribute(Attribute::Flatten))
+  for (Function &F : make_early_inc_range(M)) {
+    if (F.hasFnAttribute(Attribute::Flatten))
+      NeedFlattening.push_back(&F);
+
+    if (F.isPresplitCoroutine())
       continue;
+
+    if (F.isDeclaration() || !isInlineViable(F).isSuccess())
+      continue;
+
+    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) {
+      OptimizationRemarkEmitter ORE(CB->getCaller());
+      Changed |= TryInline(*CB, F, ORE, "always inline attribute");
+    }
+
+    F.removeDeadConstantUsers();
+    if (F.hasFnAttribute(Attribute::AlwaysInline) && F.isDefTriviallyDead()) {
+      if (F.hasComdat()) {
+        InlinedComdatFunctions.push_back(&F);
+      } else {
+        if (FAM)
+          FAM->clear(F, F.getName());
+        M.getFunctionList().erase(F);
+        Changed = true;
+      }
+    }
+  }
+
+  // Flatten functions with the flatten attribute using a local worklist.
+  for (Function *F : NeedFlattening) {
     SmallVector<std::pair<CallBase *, int>, 16> Worklist;
     SmallVector<std::pair<Function *, int>, 16> InlineHistory;
     SmallVector<CallBase *> NewCallSites;
-    OptimizationRemarkEmitter ORE(&F);
+    OptimizationRemarkEmitter ORE(F);
 
     // Collect initial calls.
-    for (BasicBlock &BB : F) {
+    for (BasicBlock &BB : *F) {
       for (Instruction &I : BB) {
         if (auto *CB = dyn_cast<CallBase>(&I)) {
           Function *Callee = CB->getCalledFunction();
@@ -104,7 +141,7 @@ bool AlwaysInlineImpl(
         continue;
 
       // Detect recursion.
-      if (Callee == &F) {
+      if (Callee == F) {
         ORE.emit([&]() {
           return OptimizationRemarkMissed("inline", "NotInlined",
                                           CB->getDebugLoc(), CB->getParent())
@@ -142,40 +179,6 @@ bool AlwaysInlineImpl(
     }
   }
 
-  for (Function &F : make_early_inc_range(M)) {
-    if (F.isPresplitCoroutine())
-      continue;
-
-    if (F.isDeclaration() || !isInlineViable(F).isSuccess())
-      continue;
-
-    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) {
-      OptimizationRemarkEmitter ORE(CB->getCaller());
-      Changed |= TryInline(*CB, F, ORE, "always inline attribute");
-    }
-
-    F.removeDeadConstantUsers();
-    if (F.hasFnAttribute(Attribute::AlwaysInline) && F.isDefTriviallyDead()) {
-      if (F.hasComdat()) {
-        InlinedComdatFunctions.push_back(&F);
-      } else {
-        if (FAM)
-          FAM->clear(F, F.getName());
-        M.getFunctionList().erase(F);
-        Changed = true;
-      }
-    }
-  }
-
   if (!InlinedComdatFunctions.empty()) {
     // Now we just have the comdat functions. Filter out the ones whose comdats
     // are not actually dead.

diff  --git a/llvm/test/Transforms/Inline/flatten.ll b/llvm/test/Transforms/Inline/flatten.ll
index 4e246ade6b1ae..355739a99dac7 100644
--- a/llvm/test/Transforms/Inline/flatten.ll
+++ b/llvm/test/Transforms/Inline/flatten.ll
@@ -203,43 +203,5 @@ define i32 @test_mutual_recursion() flatten {
   ret i32 %r
 }
 
-; Always-inline function with flatten that becomes dead after inlining.
-define internal i32 @alwaysinline_flatten_callee() alwaysinline flatten {
-  ret i32 5
-}
-
-define i32 @test_alwaysinline_flatten() {
-; ALWAYS-LABEL: define i32 @test_alwaysinline_flatten() {
-; ALWAYS-NEXT:    ret i32 5
-;
-; INLINE-LABEL: define i32 @test_alwaysinline_flatten() {
-; INLINE-NEXT:    ret i32 5
-;
-; MANDATORY-LABEL: define i32 @test_alwaysinline_flatten() {
-; MANDATORY-NEXT:    ret i32 5
-;
-  %r = call i32 @alwaysinline_flatten_callee() alwaysinline
-  ret i32 %r
-}
-
-; Flatten with alwaysinline: callees are flattened first, then the
-; flattened function is always-inlined into callers.
-define internal i32 @inner() {
-  ret i32 7
-}
-
-define internal i32 @alwaysinline_flatten_two_levels() alwaysinline flatten {
-  %r = call i32 @inner()
-  ret i32 %r
-}
-
-define i32 @test_alwaysinline_flatten_two_levels() {
-; CHECK-LABEL: define i32 @test_alwaysinline_flatten_two_levels() {
-; CHECK-NEXT:    ret i32 7
-;
-  %r = call i32 @alwaysinline_flatten_two_levels() alwaysinline
-  ret i32 %r
-}
-
 ; Check that optimization remark is emitted for recursive calls during flattening.
 ; REMARK: remark: {{.*}} 'test_direct_recursion' is not inlined into 'test_direct_recursion': recursive call during flattening