FuncSpec] Invalidate analyses when deleting a fully specialised function

[chill]

Summary

Deleting a fully specialised function left dangling pointers in
FunctionAnalysisManager, which causes an internal compiler error
when the function's storage was reused.

Fixes bug #58759.

[chill]

Ideally, now that we have the FAM in FunctionSpecializer I would like to also remove the callbacks (GetTLI, GetTTi, etc) and also remove the Legacy pass manager support entirely
from SCCP/IPSCCP/FuncSpec. Unfortunately, those callbacks are needed by getInlineCost and a few other functions down the call chain.

[ChuanqiXu]

LGTM. And it should be good to remove the Legacy PMs in separate patches.

llvm/include/llvm/Transforms/Scalar/SCCP.h

public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 bool runIPSCCP(Module &M, const DataLayout &DL,
                std::function<const TargetLibraryInfo &(Function &)> GetTLI,
                function_ref<AnalysisResultsForFn(Function &)> getAnalysis);
 
 bool runFunctionSpecialization(
-    Module &M, const DataLayout &DL,
+    Module &M, FunctionAnalysisManager *FAM, const DataLayout &DL,
     std::function<TargetLibraryInfo &(Function &)> GetTLI,
     std::function<TargetTransformInfo &(Function &)> GetTTI,
     std::function<AssumptionCache &(Function &)> GetAC,
     function_ref<AnalysisResultsForFn(Function &)> GetAnalysis);
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_SCCP_H

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 }
 
 namespace {
 class FunctionSpecializer {
 
   /// The IPSCCP Solver.
   SCCPSolver &Solver;
 
+  /// Analysis manager, needed to invalidate analyses.
+  FunctionAnalysisManager *FAM;
+
   /// Analyses used to help determine if a function should be specialized.
   std::function<AssumptionCache &(Function &)> GetAC;
   std::function<TargetTransformInfo &(Function &)> GetTTI;
   std::function<TargetLibraryInfo &(Function &)> GetTLI;
 
   SmallPtrSet<Function *, 4> SpecializedFuncs;
   SmallPtrSet<Function *, 4> FullySpecialized;
   SmallVector<Instruction *> ReplacedWithConstant;
   DenseMap<Function *, CodeMetrics> FunctionMetrics;
 
 public:
-  FunctionSpecializer(SCCPSolver &Solver,
+  FunctionSpecializer(SCCPSolver &Solver, FunctionAnalysisManager *FAM,
                       std::function<AssumptionCache &(Function &)> GetAC,
                       std::function<TargetTransformInfo &(Function &)> GetTTI,
                       std::function<TargetLibraryInfo &(Function &)> GetTLI)
-      : Solver(Solver), GetAC(GetAC), GetTTI(GetTTI), GetTLI(GetTLI) {}
+      : Solver(Solver), FAM(FAM), GetAC(GetAC), GetTTI(GetTTI), GetTLI(GetTLI) {
+  }
 
   ~FunctionSpecializer() {
     // Eliminate dead code.
     removeDeadInstructions();
     removeDeadFunctions();
   }
 
   /// Attempt to specialize functions in the module to enable constant
void removeDeadInstructions() {
     }
     ReplacedWithConstant.clear();
   }
 
   void removeDeadFunctions() {
     for (auto *F : FullySpecialized) {
       LLVM_DEBUG(dbgs() << "FnSpecialization: Removing dead function "
                         << F->getName() << "\n");
+      if (FAM)
+        FAM->clear(*F, F->getName());
       F->eraseFromParent();
     }
     FullySpecialized.clear();
   }
 
   bool tryToReplaceWithConstant(Value *V) {
     if (!V->getType()->isSingleValueType() || isa<CallBase>(V) ||
         V->user_empty())
for (auto *F : WorkList) {
           LLVM_DEBUG(dbgs() << "FnSpecialization: Replaced constant argument: "
                             << Arg.getNameOrAsOperand() << "\n");
     }
   }
 };
 } // namespace
 
 bool llvm::runFunctionSpecialization(
-    Module &M, const DataLayout &DL,
+    Module &M, FunctionAnalysisManager *FAM, const DataLayout &DL,
     std::function<TargetLibraryInfo &(Function &)> GetTLI,
     std::function<TargetTransformInfo &(Function &)> GetTTI,
     std::function<AssumptionCache &(Function &)> GetAC,
     function_ref<AnalysisResultsForFn(Function &)> GetAnalysis) {
   SCCPSolver Solver(DL, GetTLI, M.getContext());
-  FunctionSpecializer FS(Solver, GetAC, GetTTI, GetTLI);
+  FunctionSpecializer FS(Solver, FAM, GetAC, GetTTI, GetTLI);
   bool Changed = false;
 
   // Loop over all functions, marking arguments to those with their addresses
   // taken or that are external as overdefined.
   for (Function &F : M) {
     if (F.isDeclaration())
       continue;
     if (F.hasFnAttribute(Attribute::NoDuplicate))

llvm/lib/Transforms/IPO/SCCP.cpp

PreservedAnalyses FunctionSpecializationPass::run(Module &M,
   auto GetAnalysis = [&FAM](Function &F) -> AnalysisResultsForFn {
     DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
     return {std::make_unique<PredicateInfo>(
                 F, DT, FAM.getResult<AssumptionAnalysis>(F)),
             &DT, FAM.getCachedResult<PostDominatorTreeAnalysis>(F),
             &FAM.getResult<LoopAnalysis>(F)};
   };
 
-  if (!runFunctionSpecialization(M, DL, GetTLI, GetTTI, GetAC, GetAnalysis))
+  if (!runFunctionSpecialization(M, &FAM, DL, GetTLI, GetTTI, GetAC, GetAnalysis))
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserve<DominatorTreeAnalysis>();
   PA.preserve<PostDominatorTreeAnalysis>();
   PA.preserve<FunctionAnalysisManagerModuleProxy>();
   return PA;
 }
auto GetAnalysis = [this](Function &F) -> AnalysisResultsForFn {
           std::make_unique<PredicateInfo>(
               F, DT,
               this->getAnalysis<AssumptionCacheTracker>().getAssumptionCache(
                   F)),
           nullptr, // We cannot preserve the LI, DT, or PDT with the legacy pass
           nullptr, // manager, so set them to nullptr.
           nullptr};
     };
-    return runFunctionSpecialization(M, DL, GetTLI, GetTTI, GetAC, GetAnalysis);
+    return runFunctionSpecialization(M, nullptr, DL, GetTLI, GetTTI, GetAC, GetAnalysis);
   }
 };
 } // namespace
 
 char FunctionSpecializationLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(
     FunctionSpecializationLegacyPass, "function-specialization",

llvm/test/Transforms/FunctionSpecialization/compiler-crash-58759.ll

+; RUN: opt -S --passes='default<O3>' -enable-function-specialization < %s | FileCheck %s
+
+define dso_local i32 @g0(i32 noundef %x) local_unnamed_addr {
+entry:
+  %call = tail call fastcc i32 @f(i32 noundef %x, ptr noundef nonnull @p0)
+  ret i32 %call
+}
+
+define internal fastcc i32 @f(i32 noundef %x, ptr nocapture noundef readonly %p) noinline {
+entry:
+  %call = tail call i32 %p(i32 noundef %x)
+  %add = add nsw i32 %call, %x
+  ret i32 %add
+}
+
+define dso_local i32 @g1(i32 noundef %x) {
+entry:
+  %call = tail call fastcc i32 @f(i32 noundef %x, ptr noundef nonnull @p1)
+  ret i32 %call
+}
+
+declare i32 @p0(i32 noundef)
+declare i32 @p1(i32 noundef)
+
+;; Tests that `f` has been fully specialize and it didn't cause compiler crash.
+;; CHECK-DAG: f.1
+;; CHECK-DAG: f.2