[AlwaysInline] Fix analysis invalidation

[nikic]

This is a followup to #117750. Currently, AlwaysInline only invalidates analyses at the end, by returning that no analyses are preserved. However, this means that analyses fetched during inlining may be outdated. The aforementioned PR exposed this issue.

Instead, bring the logic closer to what the normal inliner does, by directly invalidating the caller in FAM. This should make sure that we don't receive any outdated analyses even if they are fetched during inlining.

Also drop the BFI updating entirely -- there's no point in doing it if we're going to invalidate everything anyway.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

This is a followup to #117750. Currently, AlwaysInline only invalidates analyses at the end, by returning that no analyses are preserved. However, this means that analyses fetched during inlining may be outdated. The aforementioned PR exposed this issue.

Instead, bring the logic closer to what the normal inliner does, by directly invalidating the caller in FAM. This should make sure that we don't receive any outdated analyses even if they are fetched during inlining.

Also drop the BFI updating entirely -- there's no point in doing it if we're going to invalidate everything anyway.

---

Full diff: https://github.com/llvm/llvm-project/pull/119566.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/AlwaysInliner.cpp (+19-23)
• (added) llvm/test/Transforms/Inline/always-inline-bfi.ll (+100)

diff --git a/llvm/lib/Transforms/IPO/AlwaysInliner.cpp b/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
index 0baa34d50abf35..20fc630a74a86b 100644
--- a/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
+++ b/llvm/lib/Transforms/IPO/AlwaysInliner.cpp
@@ -32,10 +32,9 @@ namespace {
 
 bool AlwaysInlineImpl(
     Module &M, bool InsertLifetime, ProfileSummaryInfo &PSI,
+    FunctionAnalysisManager *FAM,
     function_ref<AssumptionCache &(Function &)> GetAssumptionCache,
-    function_ref<AAResults &(Function &)> GetAAR,
-    function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
-    function_ref<BlockFrequencyInfo *(Function &)> GetCachedBFI) {
+    function_ref<AAResults &(Function &)> GetAAR) {
   SmallSetVector<CallBase *, 16> Calls;
   bool Changed = false;
   SmallVector<Function *, 16> InlinedComdatFunctions;
@@ -62,12 +61,7 @@ bool AlwaysInlineImpl(
       DebugLoc DLoc = CB->getDebugLoc();
       BasicBlock *Block = CB->getParent();
 
-      // Only update CallerBFI if already available. The CallerBFI update
-      // requires CalleeBFI.
-      BlockFrequencyInfo *CallerBFI = GetCachedBFI(*Caller);
-      InlineFunctionInfo IFI(GetAssumptionCache, &PSI, CallerBFI,
-                             CallerBFI ? &GetBFI(F) : nullptr);
-
+      InlineFunctionInfo IFI(GetAssumptionCache, &PSI, nullptr, nullptr);
       InlineResult Res = InlineFunction(*CB, IFI, /*MergeAttributes=*/true,
                                         &GetAAR(F), InsertLifetime);
       if (!Res.isSuccess()) {
@@ -86,6 +80,8 @@ bool AlwaysInlineImpl(
           /*ForProfileContext=*/false, DEBUG_TYPE);
 
       Changed = true;
+      if (FAM)
+        FAM->invalidate(*Caller, PreservedAnalyses::none());
     }
 
     F.removeDeadConstantUsers();
@@ -95,6 +91,8 @@ bool AlwaysInlineImpl(
       if (F.hasComdat()) {
         InlinedComdatFunctions.push_back(&F);
       } else {
+        if (FAM)
+          FAM->clear(F, F.getName());
         M.getFunctionList().erase(F);
         Changed = true;
       }
@@ -107,6 +105,8 @@ bool AlwaysInlineImpl(
     filterDeadComdatFunctions(InlinedComdatFunctions);
     // The remaining functions are actually dead.
     for (Function *F : InlinedComdatFunctions) {
+      if (FAM)
+        FAM->clear(*F, F->getName());
       M.getFunctionList().erase(F);
       Changed = true;
     }
@@ -136,12 +136,9 @@ struct AlwaysInlinerLegacyPass : public ModulePass {
     auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
       return getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
     };
-    auto GetCachedBFI = [](Function &) -> BlockFrequencyInfo * {
-      return nullptr;
-    };
 
-    return AlwaysInlineImpl(M, InsertLifetime, PSI, GetAssumptionCache, GetAAR,
-                            /*GetBFI=*/nullptr, GetCachedBFI);
+    return AlwaysInlineImpl(M, InsertLifetime, PSI, /*FAM=*/nullptr,
+                            GetAssumptionCache, GetAAR);
   }
 
   static char ID; // Pass identification, replacement for typeid
@@ -175,19 +172,18 @@ PreservedAnalyses AlwaysInlinerPass::run(Module &M,
   auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
     return FAM.getResult<AssumptionAnalysis>(F);
   };
-  auto GetBFI = [&](Function &F) -> BlockFrequencyInfo & {
-    return FAM.getResult<BlockFrequencyAnalysis>(F);
-  };
-  auto GetCachedBFI = [&](Function &F) -> BlockFrequencyInfo * {
-    return FAM.getCachedResult<BlockFrequencyAnalysis>(F);
-  };
   auto GetAAR = [&](Function &F) -> AAResults & {
     return FAM.getResult<AAManager>(F);
   };
   auto &PSI = MAM.getResult<ProfileSummaryAnalysis>(M);
 
-  bool Changed = AlwaysInlineImpl(M, InsertLifetime, PSI, GetAssumptionCache,
-                                  GetAAR, GetBFI, GetCachedBFI);
+  bool Changed = AlwaysInlineImpl(M, InsertLifetime, PSI, &FAM,
+                                  GetAssumptionCache, GetAAR);
+  if (!Changed)
+    return PreservedAnalyses::all();
 
-  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  // We have already invalidated all analyses on modified functions.
+  PA.preserveSet<AllAnalysesOn<Function>>();
+  return PA;
 }
diff --git a/llvm/test/Transforms/Inline/always-inline-bfi.ll b/llvm/test/Transforms/Inline/always-inline-bfi.ll
new file mode 100644
index 00000000000000..1e2ba03e513fea
--- /dev/null
+++ b/llvm/test/Transforms/Inline/always-inline-bfi.ll
@@ -0,0 +1,100 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes="function(require<block-freq>,loop(loop-unroll-full)),always-inline" < %s | FileCheck %s
+
+; Make sure this does not crash.
+
+define void @f_116_0() alwaysinline {
+; CHECK-LABEL: define void @f_116_0(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[DOTPRE:%.*]] = load i16, ptr null, align 1
+; CHECK-NEXT:    br label %[[FOR_COND:.*]]
+; CHECK:       [[FOR_COND]]:
+; CHECK-NEXT:    [[CMP3:%.*]] = icmp ult i16 [[DOTPRE]], 1
+; CHECK-NEXT:    br i1 [[CMP3]], label %[[FOR_BODY:.*]], label %[[FOR_COND_CLEANUP:.*]]
+; CHECK:       [[FOR_COND_CLEANUP]]:
+; CHECK-NEXT:    ret void
+; CHECK:       [[FOR_BODY]]:
+; CHECK-NEXT:    br label %[[FOR_COND]]
+;
+entry:
+  %.pre = load i16, ptr null, align 1
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.body, %entry
+  %cmp3 = icmp ult i16 %.pre, 1
+  br i1 %cmp3, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond
+  ret void
+
+for.body:                                         ; preds = %for.cond
+  br label %for.cond
+}
+
+define void @f_321_0() alwaysinline {
+; CHECK-LABEL: define void @f_321_0(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[FOR_COND:.*]]
+; CHECK:       [[FOR_COND]]:
+; CHECK-NEXT:    br i1 false, label %[[CRIT_EDGE:.*]], label %[[FOR_COND_CLEANUP:.*]]
+; CHECK:       [[CRIT_EDGE]]:
+; CHECK-NEXT:    unreachable
+; CHECK:       [[FOR_COND_CLEANUP]]:
+; CHECK-NEXT:    [[DOTPRE_I:%.*]] = load i16, ptr null, align 1
+; CHECK-NEXT:    br label %[[FOR_COND_I:.*]]
+; CHECK:       [[FOR_COND_I]]:
+; CHECK-NEXT:    [[CMP3_I:%.*]] = icmp ult i16 [[DOTPRE_I]], 1
+; CHECK-NEXT:    br i1 [[CMP3_I]], label %[[FOR_BODY_I:.*]], label %[[F_116_0_EXIT:.*]]
+; CHECK:       [[FOR_BODY_I]]:
+; CHECK-NEXT:    br label %[[FOR_COND_I]]
+; CHECK:       [[F_116_0_EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %crit_edge, %entry
+  br i1 false, label %crit_edge, label %for.cond.cleanup
+
+crit_edge:                                        ; preds = %for.cond
+  br label %for.cond
+
+for.cond.cleanup:                                 ; preds = %for.cond
+  call void @f_116_0()
+  ret void
+}
+
+define i16 @main() {
+; CHECK-LABEL: define i16 @main() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[FOR_COND:.*]]
+; CHECK:       [[FOR_COND]]:
+; CHECK-NEXT:    br label %[[FOR_COND]]
+; CHECK:       [[IF_ELSE:.*:]]
+; CHECK-NEXT:    [[DOTPRE_I_I:%.*]] = load i16, ptr null, align 1
+; CHECK-NEXT:    br label %[[FOR_COND_I_I:.*]]
+; CHECK:       [[FOR_COND_I_I]]:
+; CHECK-NEXT:    [[CMP3_I_I:%.*]] = icmp ult i16 [[DOTPRE_I_I]], 1
+; CHECK-NEXT:    br i1 [[CMP3_I_I]], label %[[FOR_BODY_I_I:.*]], label %[[F_321_0_EXIT:.*]]
+; CHECK:       [[FOR_BODY_I_I]]:
+; CHECK-NEXT:    br label %[[FOR_COND_I_I]]
+; CHECK:       [[F_321_0_EXIT]]:
+; CHECK-NEXT:    br label %[[FOR_COND115:.*]]
+; CHECK:       [[FOR_COND115]]:
+; CHECK-NEXT:    br label %[[FOR_COND115]]
+;
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.cond, %entry
+  br label %for.cond
+
+if.else:                                          ; No predecessors!
+  call void @f_321_0()
+  br label %for.cond115
+
+for.cond115:                                      ; preds = %for.cond115, %if.else
+  br label %for.cond115
+}