[StackSafety] Pass summary into codegen

Summary:
The patch wraps ThinLTO index into immutable
pass which can be used by StackSafety analysis.

Reviewers: eugenis, pcc

Reviewed By: eugenis

Subscribers: hiraditya, steven_wu, dexonsmith, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D80985

include/llvm/Analysis/ModuleSummaryAnalysis.h

@@ -78,6 +78,27 @@ public:
 //
 ModulePass *createModuleSummaryIndexWrapperPass();
 
+/// Legacy wrapper pass to provide the ModuleSummaryIndex object.
+class ImmutableModuleSummaryIndexWrapperPass : public ImmutablePass {
+  const ModuleSummaryIndex *Index;
+
+public:
+  static char ID;
+
+  ImmutableModuleSummaryIndexWrapperPass(
+      const ModuleSummaryIndex *Index = nullptr);
+  const ModuleSummaryIndex *getIndex() const { return Index; }
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+};
+
+//===--------------------------------------------------------------------===//
+//
+// ImmutableModuleSummaryIndexWrapperPass - This pass wrap provided
+// ModuleSummaryIndex object for the module, to be used by other passes.
+//
+ImmutablePass *
+createImmutableModuleSummaryIndexWrapperPass(const ModuleSummaryIndex *Index);
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_MODULESUMMARYANALYSIS_H

include/llvm/Analysis/StackSafetyAnalysis.h

@@ -55,13 +55,15 @@ public:
 private:
   Module *M = nullptr;
   std::function<const StackSafetyInfo &(Function &F)> GetSSI;
+  const ModuleSummaryIndex *Index = nullptr;
   mutable std::unique_ptr<InfoTy> Info;
   const InfoTy &getInfo() const;
 
 public:
   StackSafetyGlobalInfo();
   StackSafetyGlobalInfo(
-      Module *M, std::function<const StackSafetyInfo &(Function &F)> GetSSI);
+      Module *M, std::function<const StackSafetyInfo &(Function &F)> GetSSI,
+      const ModuleSummaryIndex *Index);
   StackSafetyGlobalInfo(StackSafetyGlobalInfo &&);
   StackSafetyGlobalInfo &operator=(StackSafetyGlobalInfo &&);
   ~StackSafetyGlobalInfo();

include/llvm/InitializePasses.h

@@ -183,6 +183,7 @@ void initializeIRCELegacyPassPass(PassRegistry&);
 void initializeIRTranslatorPass(PassRegistry&);
 void initializeIVUsersWrapperPassPass(PassRegistry&);
 void initializeIfConverterPass(PassRegistry&);
+void initializeImmutableModuleSummaryIndexWrapperPassPass(PassRegistry&);
 void initializeImplicitNullChecksPass(PassRegistry&);
 void initializeIndVarSimplifyLegacyPassPass(PassRegistry&);
 void initializeIndirectBrExpandPassPass(PassRegistry&);

lib/Analysis/ModuleSummaryAnalysis.cpp

@@ -909,3 +909,25 @@ void ModuleSummaryIndexWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<ProfileSummaryInfoWrapperPass>();
   AU.addRequired<StackSafetyInfoWrapperPass>();
 }
+
+char ImmutableModuleSummaryIndexWrapperPass::ID = 0;
+
+ImmutableModuleSummaryIndexWrapperPass::ImmutableModuleSummaryIndexWrapperPass(
+    const ModuleSummaryIndex *Index)
+    : ImmutablePass(ID), Index(Index) {
+  initializeImmutableModuleSummaryIndexWrapperPassPass(
+      *PassRegistry::getPassRegistry());
+}
+
+void ImmutableModuleSummaryIndexWrapperPass::getAnalysisUsage(
+    AnalysisUsage &AU) const {
+  AU.setPreservesAll();
+}
+
+ImmutablePass *llvm::createImmutableModuleSummaryIndexWrapperPass(
+    const ModuleSummaryIndex *Index) {
+  return new ImmutableModuleSummaryIndexWrapperPass(Index);
+}
+
+INITIALIZE_PASS(ImmutableModuleSummaryIndexWrapperPass, "module-summary-info",
+                "Module summary info", false, true)

lib/Analysis/StackSafetyAnalysis.cpp

@@ -12,6 +12,7 @@
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/ModuleSummaryAnalysis.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -571,7 +572,8 @@ template <typename CalleeTy> void resolveAllCalls(UseInfo<CalleeTy> &Use) {
 }
 
 GVToSSI createGlobalStackSafetyInfo(
-    std::map<const GlobalValue *, FunctionInfo<GlobalValue>> Functions) {
+    std::map<const GlobalValue *, FunctionInfo<GlobalValue>> Functions,
+    const ModuleSummaryIndex *Index) {
   GVToSSI SSI;
   if (Functions.empty())
     return SSI;
@@ -647,8 +649,8 @@ const StackSafetyGlobalInfo::InfoTy &StackSafetyGlobalInfo::getInfo() const {
         Functions.emplace(&F, std::move(FI));
       }
     }
-    Info.reset(
+    Info.reset(new InfoTy{
-        new InfoTy{createGlobalStackSafetyInfo(std::move(Functions)), {}});
+        createGlobalStackSafetyInfo(std::move(Functions), Index), {}});
     for (auto &FnKV : Info->Info) {
       for (auto &KV : FnKV.second.Allocas) {
         ++NumAllocaTotal;
@@ -695,8 +697,9 @@ StackSafetyInfo::getParamAccesses() const {
 StackSafetyGlobalInfo::StackSafetyGlobalInfo() = default;
 
 StackSafetyGlobalInfo::StackSafetyGlobalInfo(
-    Module *M, std::function<const StackSafetyInfo &(Function &F)> GetSSI)
+    Module *M, std::function<const StackSafetyInfo &(Function &F)> GetSSI,
-    : M(M), GetSSI(GetSSI) {
+    const ModuleSummaryIndex *Index)
+    : M(M), GetSSI(GetSSI), Index(Index) {
   if (StackSafetyRun)
     getInfo();
 }
@@ -770,11 +773,14 @@ AnalysisKey StackSafetyGlobalAnalysis::Key;
 
 StackSafetyGlobalInfo
 StackSafetyGlobalAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
+  // FIXME: Lookup Module Summary.
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
-  return {&M, [&FAM](Function &F) -> const StackSafetyInfo & {
+  return {&M,
+          [&FAM](Function &F) -> const StackSafetyInfo & {
             return FAM.getResult<StackSafetyAnalysis>(F);
-          }};
+          },
+          nullptr};
 }
 
 PreservedAnalyses StackSafetyGlobalPrinterPass::run(Module &M,
@@ -806,13 +812,22 @@ void StackSafetyGlobalInfoWrapperPass::getAnalysisUsage(
 }
 
 bool StackSafetyGlobalInfoWrapperPass::runOnModule(Module &M) {
-  SSGI = {&M, [this](Function &F) -> const StackSafetyInfo & {
+  const ModuleSummaryIndex *ImportSummary = nullptr;
+  if (auto *IndexWrapperPass =
+          getAnalysisIfAvailable<ImmutableModuleSummaryIndexWrapperPass>())
+    ImportSummary = IndexWrapperPass->getIndex();
+
+  SSGI = {&M,
+          [this](Function &F) -> const StackSafetyInfo & {
             return getAnalysis<StackSafetyInfoWrapperPass>(F).getResult();
-          }};
+          },
+          ImportSummary};
   return false;
 }
 
 bool llvm::needsParamAccessSummary(const Module &M) {
+  if (StackSafetyRun)
+    return true;
   for (auto &F : M.functions())
     if (F.hasFnAttribute(Attribute::SanitizeMemTag))
       return true;
@@ -831,5 +846,6 @@ static const char GlobalPassName[] = "Stack Safety Analysis";
 INITIALIZE_PASS_BEGIN(StackSafetyGlobalInfoWrapperPass, DEBUG_TYPE,
                       GlobalPassName, false, true)
 INITIALIZE_PASS_DEPENDENCY(StackSafetyInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(ImmutableModuleSummaryIndexWrapperPass)
 INITIALIZE_PASS_END(StackSafetyGlobalInfoWrapperPass, DEBUG_TYPE,
                     GlobalPassName, false, true)

lib/LTO/LTOBackend.cpp

@@ -16,6 +16,7 @@
 #include "llvm/LTO/LTOBackend.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/ModuleSummaryAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Bitcode/BitcodeReader.h"
@@ -363,7 +364,8 @@ static void EmitBitcodeSection(Module &M, const Config &Conf) {
 }
 
 void codegen(const Config &Conf, TargetMachine *TM, AddStreamFn AddStream,
-             unsigned Task, Module &Mod) {
+             unsigned Task, Module &Mod,
+             const ModuleSummaryIndex &CombinedIndex) {
   if (Conf.PreCodeGenModuleHook && !Conf.PreCodeGenModuleHook(Task, Mod))
     return;
 
@@ -392,6 +394,8 @@ void codegen(const Config &Conf, TargetMachine *TM, AddStreamFn AddStream,
 
   auto Stream = AddStream(Task);
   legacy::PassManager CodeGenPasses;
+  CodeGenPasses.add(
+      createImmutableModuleSummaryIndexWrapperPass(&CombinedIndex));
   if (TM->addPassesToEmitFile(CodeGenPasses, *Stream->OS,
                               DwoOut ? &DwoOut->os() : nullptr,
                               Conf.CGFileType))
@@ -404,7 +408,8 @@ void codegen(const Config &Conf, TargetMachine *TM, AddStreamFn AddStream,
 
 void splitCodeGen(const Config &C, TargetMachine *TM, AddStreamFn AddStream,
                   unsigned ParallelCodeGenParallelismLevel,
-                  std::unique_ptr<Module> Mod) {
+                  std::unique_ptr<Module> Mod,
+                  const ModuleSummaryIndex &CombinedIndex) {
   ThreadPool CodegenThreadPool(
       heavyweight_hardware_concurrency(ParallelCodeGenParallelismLevel));
   unsigned ThreadCount = 0;
@@ -437,7 +442,8 @@ void splitCodeGen(const Config &C, TargetMachine *TM, AddStreamFn AddStream,
               std::unique_ptr<TargetMachine> TM =
                   createTargetMachine(C, T, *MPartInCtx);
 
-              codegen(C, TM.get(), AddStream, ThreadId, *MPartInCtx);
+              codegen(C, TM.get(), AddStream, ThreadId, *MPartInCtx,
+                      CombinedIndex);
             },
             // Pass BC using std::move to ensure that it get moved rather than
             // copied into the thread's context.
@@ -493,10 +499,10 @@ Error lto::backend(const Config &C, AddStreamFn AddStream,
   }
 
   if (ParallelCodeGenParallelismLevel == 1) {
-    codegen(C, TM.get(), AddStream, 0, *Mod);
+    codegen(C, TM.get(), AddStream, 0, *Mod, CombinedIndex);
   } else {
     splitCodeGen(C, TM.get(), AddStream, ParallelCodeGenParallelismLevel,
-                 std::move(Mod));
+                 std::move(Mod), CombinedIndex);
   }
   return Error::success();
 }
@@ -546,7 +552,7 @@ Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   Mod.setPartialSampleProfileRatio(CombinedIndex);
 
   if (Conf.CodeGenOnly) {
-    codegen(Conf, TM.get(), AddStream, Task, Mod);
+    codegen(Conf, TM.get(), AddStream, Task, Mod, CombinedIndex);
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
   }
 
@@ -597,6 +603,6 @@ Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
            /*ExportSummary=*/nullptr, /*ImportSummary=*/&CombinedIndex))
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 
-  codegen(Conf, TM.get(), AddStream, Task, Mod);
+  codegen(Conf, TM.get(), AddStream, Task, Mod, CombinedIndex);
   return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 }