[llvm] [LLVM] Fix use-after-free in AlwaysInliner flatten worklist (PR #194485)

[Grigory Pastukhov via llvm-commits]

https://github.com/grigorypas updated https://github.com/llvm/llvm-project/pull/194485

>From a6a6db5a3ee574bc86ed57de944167b8ffb28f90 Mon Sep 17 00:00:00 2001
From: Grigory Pastukhov <gpastukhov at meta.com>
Date: Mon, 27 Apr 2026 16:03:04 -0700
Subject: [PATCH] [AlwaysInliner] Fix use-after-free in flatten worklist

---
 llvm/lib/Transforms/IPO/AlwaysInliner.cpp | 81 +++++++++++------------
 llvm/test/Transforms/Inline/flatten.ll    | 38 +++++++++++
 2 files changed, 77 insertions(+), 42 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/AlwaysInliner.cpp b/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
index 080cb8ddb33fd..bdba6a1dbca53 100644
--- a/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
+++ b/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
@@ -42,7 +42,6 @@ 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,
@@ -76,52 +75,16 @@ bool AlwaysInlineImpl(
     return true;
   };
 
-  for (Function &F : make_early_inc_range(M)) {
-    if (F.hasFnAttribute(Attribute::Flatten))
-      NeedFlattening.push_back(&F);
-
-    if (F.isPresplitCoroutine())
+  for (Function &F : M) {
+    if (!F.hasFnAttribute(Attribute::Flatten))
       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();
@@ -141,7 +104,7 @@ bool AlwaysInlineImpl(
         continue;
 
       // Detect recursion.
-      if (Callee == F) {
+      if (Callee == &F) {
         ORE.emit([&]() {
           return OptimizationRemarkMissed("inline", "NotInlined",
                                           CB->getDebugLoc(), CB->getParent())
@@ -179,6 +142,40 @@ 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 355739a99dac7..4e246ade6b1ae 100644
--- a/llvm/test/Transforms/Inline/flatten.ll
+++ b/llvm/test/Transforms/Inline/flatten.ll
@@ -203,5 +203,43 @@ 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