[StackSafety] Run ThinLTO

Summary:
ThinLTO linking runs dataflow processing on collected
function parameters. Then StackSafetyGlobalInfoWrapperPass
in ThinLTO backend will run as usual looking up to external
symbol in the summary if needed.

Depends on D80985.

Reviewers: eugenis, pcc

Reviewed By: eugenis

Subscribers: inglorion, hiraditya, steven_wu, dexonsmith, cfe-commits, llvm-commits

Tags: #clang, #llvm

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

1 files changed, 111 insertions, 4 deletions

diff --git a/llvm/lib/Analysis/StackSafetyAnalysis.cpp b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
index 3fb5c29..a59d564 100644
--- a/llvm/lib/Analysis/StackSafetyAnalysis.cpp
+++ b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
@@ -544,6 +544,22 @@ StackSafetyDataFlowAnalysis<CalleeTy>::run() {
   return Functions;
 }
 
+FunctionSummary *resolveCallee(GlobalValueSummary *S) {
+  while (S) {
+    if (!S->isLive() || !S->isDSOLocal())
+      return nullptr;
+    if (FunctionSummary *FS = dyn_cast<FunctionSummary>(S))
+      return FS;
+    AliasSummary *AS = dyn_cast<AliasSummary>(S);
+    if (!AS)
+      return nullptr;
+    S = AS->getBaseObject();
+    if (S == AS)
+      return nullptr;
+  }
+  return nullptr;
+}
+
 const Function *findCalleeInModule(const GlobalValue *GV) {
   while (GV) {
     if (GV->isDeclaration() || GV->isInterposable() || !GV->isDSOLocal())
@@ -560,7 +576,28 @@ const Function *findCalleeInModule(const GlobalValue *GV) {
   return nullptr;
 }
 
-template <typename CalleeTy> void resolveAllCalls(UseInfo<CalleeTy> &Use) {
+GlobalValueSummary *getGlobalValueSummary(const ModuleSummaryIndex *Index,
+                                          uint64_t ValueGUID) {
+  auto VI = Index->getValueInfo(ValueGUID);
+  if (!VI || VI.getSummaryList().empty())
+    return nullptr;
+  assert(VI.getSummaryList().size() == 1);
+  auto &Summary = VI.getSummaryList()[0];
+  return Summary.get();
+}
+
+const ConstantRange *findParamAccess(const FunctionSummary &FS,
+                                     uint32_t ParamNo) {
+  assert(FS.isLive());
+  assert(FS.isDSOLocal());
+  for (auto &PS : FS.paramAccesses())
+    if (ParamNo == PS.ParamNo)
+      return &PS.Use;
+  return nullptr;
+}
+
+void resolveAllCalls(UseInfo<GlobalValue> &Use,
+                     const ModuleSummaryIndex *Index) {
   ConstantRange FullSet(Use.Range.getBitWidth(), true);
   for (auto &C : Use.Calls) {
     const Function *F = findCalleeInModule(C.Callee);
@@ -569,8 +606,24 @@ template <typename CalleeTy> void resolveAllCalls(UseInfo<CalleeTy> &Use) {
       continue;
     }
 
-    return Use.updateRange(FullSet);
+    if (!Index)
+      return Use.updateRange(FullSet);
+    GlobalValueSummary *GVS = getGlobalValueSummary(Index, C.Callee->getGUID());
+
+    FunctionSummary *FS = resolveCallee(GVS);
+    if (!FS)
+      return Use.updateRange(FullSet);
+    const ConstantRange *Found = findParamAccess(*FS, C.ParamNo);
+    if (!Found)
+      return Use.updateRange(FullSet);
+    ConstantRange Access = Found->sextOrTrunc(Use.Range.getBitWidth());
+    Use.updateRange(addOverflowNever(Access, C.Offset));
+    C.Callee = nullptr;
   }
+
+  Use.Calls.erase(std::remove_if(Use.Calls.begin(), Use.Calls.end(),
+                                 [](auto &T) { return !T.Callee; }),
+                  Use.Calls.end());
 }
 
 GVToSSI createGlobalStackSafetyInfo(
@@ -585,7 +638,7 @@ GVToSSI createGlobalStackSafetyInfo(
 
   for (auto &FnKV : Copy)
     for (auto &KV : FnKV.second.Params)
-      resolveAllCalls(KV.second);
+      resolveAllCalls(KV.second, Index);
 
   uint32_t PointerSize = Copy.begin()
                              ->first->getParent()
@@ -598,7 +651,7 @@ GVToSSI createGlobalStackSafetyInfo(
     auto &SrcF = Functions[F.first];
     for (auto &KV : FI.Allocas) {
       auto &A = KV.second;
-      resolveAllCalls(A);
+      resolveAllCalls(A, Index);
       for (auto &C : A.Calls) {
         A.updateRange(
             SSDFA.getArgumentAccessRange(C.Callee, C.ParamNo, C.Offset));
@@ -836,6 +889,60 @@ bool llvm::needsParamAccessSummary(const Module &M) {
   return false;
 }
 
+void llvm::generateParamAccessSummary(ModuleSummaryIndex &Index) {
+  const ConstantRange FullSet(FunctionSummary::ParamAccess::RangeWidth, true);
+  std::map<const FunctionSummary *, FunctionInfo<FunctionSummary>> Functions;
+
+  // Convert the ModuleSummaryIndex to a FunctionMap
+  for (auto &GVS : Index) {
+    for (auto &GV : GVS.second.SummaryList) {
+      FunctionSummary *FS = dyn_cast<FunctionSummary>(GV.get());
+      if (!FS)
+        continue;
+      if (FS->isLive() && FS->isDSOLocal()) {
+        FunctionInfo<FunctionSummary> FI;
+        for (auto &PS : FS->paramAccesses()) {
+          auto &US =
+              FI.Params
+                  .emplace(PS.ParamNo, FunctionSummary::ParamAccess::RangeWidth)
+                  .first->second;
+          US.Range = PS.Use;
+          for (auto &Call : PS.Calls) {
+            assert(!Call.Offsets.isFullSet());
+            FunctionSummary *S = resolveCallee(
+                Index.findSummaryInModule(Call.Callee, FS->modulePath()));
+            if (!S) {
+              US.Range = FullSet;
+              US.Calls.clear();
+              break;
+            }
+            US.Calls.emplace_back(S, Call.ParamNo, Call.Offsets);
+          }
+        }
+        Functions.emplace(FS, std::move(FI));
+      }
+      // Reset data for all summaries. Alive and DSO local will be set back from
+      // of data flow results below. Anything else will not be accessed
+      // by ThinLTO backend, so we can save on bitcode size.
+      FS->setParamAccesses({});
+    }
+  }
+  StackSafetyDataFlowAnalysis<FunctionSummary> SSDFA(
+      FunctionSummary::ParamAccess::RangeWidth, std::move(Functions));
+  for (auto &KV : SSDFA.run()) {
+    std::vector<FunctionSummary::ParamAccess> NewParams;
+    NewParams.reserve(KV.second.Params.size());
+    for (auto &KV : KV.second.Params) {
+      NewParams.emplace_back();
+      FunctionSummary::ParamAccess &New = NewParams.back();
+      New.ParamNo = KV.first;
+      New.Use = KV.second.Range; // Only range is needed.
+    }
+    const_cast<FunctionSummary *>(KV.first)->setParamAccesses(
+        std::move(NewParams));
+  }
+}
+
 static const char LocalPassArg[] = "stack-safety-local";
 static const char LocalPassName[] = "Stack Safety Local Analysis";
 INITIALIZE_PASS_BEGIN(StackSafetyInfoWrapperPass, LocalPassArg, LocalPassName,