[StaticDataLayout][MemProf]Introduce an LLVM option to specify one of read-only vs read-writeusers/mingmingl-llvm/profile

3 files changed, 46 insertions, 9 deletions

diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 1584d3087557..7df2dc888d66 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -1283,7 +1283,7 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
                                                EnableSampledInstr));
 
   if (IsMemprofUse)
-    MPM.addPass(MemProfUsePass(PGOOpt->MemoryProfile, FS));
+    MPM.addPass(MemProfUsePass(PGOOpt->MemoryProfile, TM, FS));
 
   if (PGOOpt && (PGOOpt->Action == PGOOptions::IRUse ||
                  PGOOpt->Action == PGOOptions::SampleUse))
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 2cfb5b259260..3e1a6e0bc564 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -243,7 +243,7 @@ MODULE_PASS_WITH_PARAMS(
     parseLoopExtractorPassOptions, "single")
 MODULE_PASS_WITH_PARAMS(
     "memprof-use", "MemProfUsePass",
-    [](std::string Opts) { return MemProfUsePass(Opts); },
+    [this](std::string Opts) { return MemProfUsePass(Opts, this->TM); },
     parseMemProfUsePassOptions, "profile-filename=S")
 MODULE_PASS_WITH_PARAMS(
     "msan", "MemorySanitizerPass",
diff --git a/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp b/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp
index 1a55021c5d3f..ca6d509912c7 100644
--- a/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemProfUse.cpp
@@ -32,6 +32,8 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/HashBuilder.h"
 #include "llvm/Support/VirtualFileSystem.h"
+#include "llvm/Target/TargetLoweringObjectFile.h"
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/LongestCommonSequence.h"
 #include <map>
 #include <set>
@@ -88,9 +90,24 @@ static cl::opt<unsigned> MinMatchedColdBytePercent(
     "memprof-matching-cold-threshold", cl::init(100), cl::Hidden,
     cl::desc("Min percent of cold bytes matched to hint allocation cold"));
 
-static cl::opt<bool> AnnotateStaticDataSectionPrefix(
-    "memprof-annotate-static-data-prefix", cl::init(false), cl::Hidden,
-    cl::desc("If true, annotate the static data section prefix"));
+enum class AnnotateDataType {
+  None,
+  ReadOnly,
+  ReadWrite,
+};
+static cl::opt<enum AnnotateDataType> AnnotateStaticDataType(
+    "memprof-annotate-static-data-type",
+    cl::values(
+        clEnumValN(AnnotateDataType::None, "none",
+                   "Do not annotate static data sections in this pass"),
+        clEnumValN(AnnotateDataType::ReadOnly, "readonly",
+                   "Annotate read-only static data sections (i.e., .rodata and "
+                   ".data.rel.ro)"),
+        clEnumValN(AnnotateDataType::ReadWrite, "readwrite",
+                   "Annotate both read-only (i.e., .rodata and .data.rel.ro) "
+                   "and read-write (.data and .bss) static data sections")),
+    cl::init(AnnotateDataType::None), cl::Hidden,
+    cl::desc("Specify the type of static data sections to annotate"));
 
 // Matching statistics
 STATISTIC(NumOfMemProfMissing, "Number of functions without memory profile.");
@@ -233,6 +250,20 @@ static void HandleUnsupportedAnnotationKinds(GlobalVariable &GVar,
                     << Reason << ".\n");
 }
 
+// Returns true if the global variable GV should be annotated based on
+// AnnotateStaticDataType and its section kind.
+static bool ShouldAnnotateStaticDataSection(const GlobalVariable &GV,
+                                            const TargetMachine &TM) {
+  SectionKind Kind = TargetLoweringObjectFile::getKindForGlobal(&GV, TM);
+  const bool IsReadOnlyData = Kind.isReadOnly() || Kind.isReadOnlyWithRel();
+  if (AnnotateStaticDataType == AnnotateDataType::ReadOnly) {
+    return IsReadOnlyData;
+  }
+  assert(AnnotateStaticDataType == AnnotateDataType::ReadWrite &&
+         "Unknown static data annotation type");
+  return IsReadOnlyData || Kind.isData() || Kind.isBSS();
+}
+
 // Structure for tracking info about matched allocation contexts for use with
 // -memprof-print-match-info and -memprof-print-matched-alloc-stack.
 struct AllocMatchInfo {
@@ -797,9 +828,9 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
   }
 }
 
-MemProfUsePass::MemProfUsePass(std::string MemoryProfileFile,
+MemProfUsePass::MemProfUsePass(std::string MemoryProfileFile, TargetMachine *TM,
                                IntrusiveRefCntPtr<vfs::FileSystem> FS)
-    : MemoryProfileFileName(MemoryProfileFile), FS(FS) {
+    : MemoryProfileFileName(MemoryProfileFile), TM(TM), FS(FS) {
   if (!FS)
     this->FS = vfs::getRealFileSystem();
 }
@@ -905,7 +936,7 @@ PreservedAnalyses MemProfUsePass::run(Module &M, ModuleAnalysisManager &AM) {
 
 bool MemProfUsePass::annotateGlobalVariables(
     Module &M, const memprof::DataAccessProfData *DataAccessProf) {
-  if (!AnnotateStaticDataSectionPrefix || M.globals().empty())
+  if (AnnotateStaticDataType == AnnotateDataType::None || M.globals().empty())
     return false;
 
   if (!DataAccessProf) {
@@ -913,7 +944,7 @@ bool MemProfUsePass::annotateGlobalVariables(
     M.getContext().diagnose(DiagnosticInfoPGOProfile(
         MemoryProfileFileName.data(),
         StringRef("Data access profiles not found in memprof. Ignore "
-                  "-memprof-annotate-static-data-prefix."),
+                  "-memprof-annotate-static-data-type."),
         DS_Warning));
     return false;
   }
@@ -934,6 +965,12 @@ bool MemProfUsePass::annotateGlobalVariables(
       continue;
     }
 
+    if (!ShouldAnnotateStaticDataSection(GVar, *TM)) {
+      LLVM_DEBUG(dbgs() << "Skip annotating global variable " << GVar.getName()
+                        << "\n");
+      continue;
+    }
+
     StringRef Name = GVar.getName();
     // Skip string literals as their mangled names don't stay stable across
     // binary releases.