Restore "[MemProf] Use new option/pass for profile feedback and matching"

This restores commit b4a82b6, reverted
in 3ab7ef2 because it was thought to
cause a bot failure, which ended up being unrelated to this patch set.

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

Author: teresajohnson

clang/include/clang/Basic/CodeGenOptions.h

@@ -282,6 +282,9 @@ class CodeGenOptions : public CodeGenOptionsBase {
   /// Name of the profile file to use as output for with -fmemory-profile.
   std::string MemoryProfileOutput;
 
+  /// Name of the profile file to use as input for -fmemory-profile-use.
+  std::string MemoryProfileUsePath;
+
   /// Name of the profile file to use as input for -fprofile-instr-use
   std::string ProfileInstrumentUsePath;
 

clang/include/clang/Driver/Options.td

@@ -1772,6 +1772,10 @@ defm memory_profile : OptInCC1FFlag<"memory-profile", "Enable", "Disable", " hea
 def fmemory_profile_EQ : Joined<["-"], "fmemory-profile=">,
     Group<f_Group>, Flags<[CC1Option]>, MetaVarName<"<directory>">,
     HelpText<"Enable heap memory profiling and dump results into <directory>">;
+def fmemory_profile_use_EQ : Joined<["-"], "fmemory-profile-use=">,
+    Group<f_Group>, Flags<[CC1Option, CoreOption]>, MetaVarName<"<pathname>">,
+    HelpText<"Use memory profile for profile-guided memory optimization">,
+    MarshallingInfoString<CodeGenOpts<"MemoryProfileUsePath">>;
 
 // Begin sanitizer flags. These should all be core options exposed in all driver
 // modes.

clang/lib/CodeGen/BackendUtil.cpp

@@ -762,31 +762,37 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
     PGOOpt = PGOOptions(
         CodeGenOpts.InstrProfileOutput.empty() ? getDefaultProfileGenName()
                                                : CodeGenOpts.InstrProfileOutput,
-        "", "", nullptr, PGOOptions::IRInstr, PGOOptions::NoCSAction,
-        CodeGenOpts.DebugInfoForProfiling);
+        "", "", CodeGenOpts.MemoryProfileUsePath, nullptr, PGOOptions::IRInstr,
+        PGOOptions::NoCSAction, CodeGenOpts.DebugInfoForProfiling);
   else if (CodeGenOpts.hasProfileIRUse()) {
     // -fprofile-use.
     auto CSAction = CodeGenOpts.hasProfileCSIRUse() ? PGOOptions::CSIRUse
                                                     : PGOOptions::NoCSAction;
-    PGOOpt =
-        PGOOptions(CodeGenOpts.ProfileInstrumentUsePath, "",
-                   CodeGenOpts.ProfileRemappingFile, VFS, PGOOptions::IRUse,
-                   CSAction, CodeGenOpts.DebugInfoForProfiling);
+    PGOOpt = PGOOptions(
+        CodeGenOpts.ProfileInstrumentUsePath, "",
+        CodeGenOpts.ProfileRemappingFile, CodeGenOpts.MemoryProfileUsePath, VFS,
+        PGOOptions::IRUse, CSAction, CodeGenOpts.DebugInfoForProfiling);
   } else if (!CodeGenOpts.SampleProfileFile.empty())
     // -fprofile-sample-use
     PGOOpt = PGOOptions(
         CodeGenOpts.SampleProfileFile, "", CodeGenOpts.ProfileRemappingFile,
-        VFS, PGOOptions::SampleUse, PGOOptions::NoCSAction,
-        CodeGenOpts.DebugInfoForProfiling, CodeGenOpts.PseudoProbeForProfiling);
+        CodeGenOpts.MemoryProfileUsePath, VFS, PGOOptions::SampleUse,
+        PGOOptions::NoCSAction, CodeGenOpts.DebugInfoForProfiling,
+        CodeGenOpts.PseudoProbeForProfiling);
+  else if (!CodeGenOpts.MemoryProfileUsePath.empty())
+    // -fmemory-profile-use (without any of the above options)
+    PGOOpt = PGOOptions("", "", "", CodeGenOpts.MemoryProfileUsePath, VFS,
+                        PGOOptions::NoAction, PGOOptions::NoCSAction,
+                        CodeGenOpts.DebugInfoForProfiling);
   else if (CodeGenOpts.PseudoProbeForProfiling)
     // -fpseudo-probe-for-profiling
-    PGOOpt = PGOOptions("", "", "", nullptr, PGOOptions::NoAction,
-                        PGOOptions::NoCSAction,
+    PGOOpt = PGOOptions("", "", "", /*MemoryProfile=*/"", nullptr,
+                        PGOOptions::NoAction, PGOOptions::NoCSAction,
                         CodeGenOpts.DebugInfoForProfiling, true);
   else if (CodeGenOpts.DebugInfoForProfiling)
     // -fdebug-info-for-profiling
-    PGOOpt = PGOOptions("", "", "", nullptr, PGOOptions::NoAction,
-                        PGOOptions::NoCSAction, true);
+    PGOOpt = PGOOptions("", "", "", /*MemoryProfile=*/"", nullptr,
+                        PGOOptions::NoAction, PGOOptions::NoCSAction, true);
 
   // Check to see if we want to generate a CS profile.
   if (CodeGenOpts.hasProfileCSIRInstr()) {
@@ -808,8 +814,8 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
                      CodeGenOpts.InstrProfileOutput.empty()
                          ? getDefaultProfileGenName()
                          : CodeGenOpts.InstrProfileOutput,
-                     "", nullptr, PGOOptions::NoAction, PGOOptions::CSIRInstr,
-                     CodeGenOpts.DebugInfoForProfiling);
+                     "", /*MemoryProfile=*/"", nullptr, PGOOptions::NoAction,
+                     PGOOptions::CSIRInstr, CodeGenOpts.DebugInfoForProfiling);
   }
   if (TM)
     TM->setPGOOption(PGOOpt);

clang/lib/Driver/ToolChains/Clang.cpp

@@ -4946,6 +4946,18 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
       !MemProfArg->getOption().matches(options::OPT_fno_memory_profile))
     MemProfArg->render(Args, CmdArgs);
 
+  if (auto *MemProfUseArg =
+          Args.getLastArg(options::OPT_fmemory_profile_use_EQ)) {
+    if (MemProfArg)
+      D.Diag(diag::err_drv_argument_not_allowed_with)
+          << MemProfUseArg->getAsString(Args) << MemProfArg->getAsString(Args);
+    if (auto *PGOInstrArg = Args.getLastArg(options::OPT_fprofile_generate,
+                                            options::OPT_fprofile_generate_EQ))
+      D.Diag(diag::err_drv_argument_not_allowed_with)
+          << MemProfUseArg->getAsString(Args) << PGOInstrArg->getAsString(Args);
+    MemProfUseArg->render(Args, CmdArgs);
+  }
+
   // Embed-bitcode option.
   // Only white-listed flags below are allowed to be embedded.
   if (C.getDriver().embedBitcodeInObject() && !IsUsingLTO &&

clang/test/CodeGen/memprof.cpp

@@ -16,8 +16,8 @@
 
 // Profile use:
 // Ensure Pass PGOInstrumentationUse is invoked with the memprof-only profile.
-// RUN: %clang_cc1 -O2 -fprofile-instrument-use-path=%t.memprofdata %s -fdebug-pass-manager  -emit-llvm -o - 2>&1 | FileCheck %s -check-prefix=USE
-// USE: Running pass: PGOInstrumentationUse on [module]
+// RUN: %clang_cc1 -O2 -fmemory-profile-use=%t.memprofdata %s -fdebug-pass-manager  -emit-llvm -o - 2>&1 | FileCheck %s -check-prefix=USE
+// USE: Running pass: MemProfUsePass on [module]
 
 char *foo() {
   return new char[10];

clang/test/Driver/fmemprof.cpp

@@ -8,3 +8,12 @@
 // DIR: ld{{.*}}libclang_rt.memprof{{.*}}libclang_rt.memprof_cxx
 // OFF-NOT: "-fmemory-profile"
 // OFF-NOT: libclang_rt.memprof
+
+// RUN: %clangxx -target x86_64-linux-gnu -fmemory-profile-use=foo %s -### 2>&1 | FileCheck %s --check-prefix=USE
+// USE: "-cc1" {{.*}} "-fmemory-profile-use=foo"
+
+// RUN: %clangxx -target x86_64-linux-gnu -fmemory-profile -fmemory-profile-use=foo %s -### 2>&1 | FileCheck %s --check-prefix=CONFLICTWITHMEMPROFINSTR
+// CONFLICTWITHMEMPROFINSTR: error: invalid argument '-fmemory-profile-use=foo' not allowed with '-fmemory-profile'
+
+// RUN: %clangxx -target x86_64-linux-gnu -fprofile-generate -fmemory-profile-use=foo %s -### 2>&1 | FileCheck %s --check-prefix=CONFLICTWITHPGOINSTR
+// CONFLICTWITHPGOINSTR: error: invalid argument '-fmemory-profile-use=foo' not allowed with '-fprofile-generate'

llvm/include/llvm/Support/PGOOptions.h

@@ -28,7 +28,7 @@ struct PGOOptions {
   enum PGOAction { NoAction, IRInstr, IRUse, SampleUse };
   enum CSPGOAction { NoCSAction, CSIRInstr, CSIRUse };
   PGOOptions(std::string ProfileFile, std::string CSProfileGenFile,
-             std::string ProfileRemappingFile,
+             std::string ProfileRemappingFile, std::string MemoryProfile,
              IntrusiveRefCntPtr<vfs::FileSystem> FS,
              PGOAction Action = NoAction, CSPGOAction CSAction = NoCSAction,
              bool DebugInfoForProfiling = false,
@@ -40,6 +40,7 @@ struct PGOOptions {
   std::string ProfileFile;
   std::string CSProfileGenFile;
   std::string ProfileRemappingFile;
+  std::string MemoryProfile;
   PGOAction Action;
   CSPGOAction CSAction;
   bool DebugInfoForProfiling;

llvm/include/llvm/Transforms/Instrumentation/MemProfiler.h

@@ -12,6 +12,7 @@
 #ifndef LLVM_TRANSFORMS_INSTRUMENTATION_MEMPROFILER_H
 #define LLVM_TRANSFORMS_INSTRUMENTATION_MEMPROFILER_H
 
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
@@ -20,6 +21,10 @@ class FunctionPass;
 class Module;
 class ModulePass;
 
+namespace vfs {
+class FileSystem;
+} // namespace vfs
+
 /// Public interface to the memory profiler pass for instrumenting code to
 /// profile memory accesses.
 ///
@@ -43,12 +48,16 @@ class ModuleMemProfilerPass : public PassInfoMixin<ModuleMemProfilerPass> {
   static bool isRequired() { return true; }
 };
 
-// TODO: Remove this declaration and make readMemprof static once the matching
-// is moved into its own pass.
-class IndexedInstrProfReader;
-class TargetLibraryInfo;
-void readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
-                 const TargetLibraryInfo &TLI);
+class MemProfUsePass : public PassInfoMixin<MemProfUsePass> {
+public:
+  explicit MemProfUsePass(std::string MemoryProfileFile,
+                          IntrusiveRefCntPtr<vfs::FileSystem> FS = nullptr);
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+
+private:
+  std::string MemoryProfileFileName;
+  IntrusiveRefCntPtr<vfs::FileSystem> FS;
+};
 
 } // namespace llvm
 

llvm/lib/LTO/LTOBackend.cpp

@@ -236,20 +236,21 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
   auto FS = vfs::getRealFileSystem();
   std::optional<PGOOptions> PGOOpt;
   if (!Conf.SampleProfile.empty())
-    PGOOpt = PGOOptions(Conf.SampleProfile, "", Conf.ProfileRemapping, FS,
-                        PGOOptions::SampleUse, PGOOptions::NoCSAction, true);
+    PGOOpt = PGOOptions(Conf.SampleProfile, "", Conf.ProfileRemapping,
+                        /*MemoryProfile=*/"", FS, PGOOptions::SampleUse,
+                        PGOOptions::NoCSAction, true);
   else if (Conf.RunCSIRInstr) {
-    PGOOpt = PGOOptions("", Conf.CSIRProfile, Conf.ProfileRemapping, FS,
-                        PGOOptions::IRUse, PGOOptions::CSIRInstr,
-                        Conf.AddFSDiscriminator);
+    PGOOpt = PGOOptions("", Conf.CSIRProfile, Conf.ProfileRemapping,
+                        /*MemoryProfile=*/"", FS, PGOOptions::IRUse,
+                        PGOOptions::CSIRInstr, Conf.AddFSDiscriminator);
   } else if (!Conf.CSIRProfile.empty()) {
-    PGOOpt = PGOOptions(Conf.CSIRProfile, "", Conf.ProfileRemapping, FS,
-                        PGOOptions::IRUse, PGOOptions::CSIRUse,
-                        Conf.AddFSDiscriminator);
+    PGOOpt = PGOOptions(Conf.CSIRProfile, "", Conf.ProfileRemapping,
+                        /*MemoryProfile=*/"", FS, PGOOptions::IRUse,
+                        PGOOptions::CSIRUse, Conf.AddFSDiscriminator);
     NoPGOWarnMismatch = !Conf.PGOWarnMismatch;
   } else if (Conf.AddFSDiscriminator) {
-    PGOOpt = PGOOptions("", "", "", nullptr, PGOOptions::NoAction,
-                        PGOOptions::NoCSAction, true);
+    PGOOpt = PGOOptions("", "", "", /*MemoryProfile=*/"", nullptr,
+                        PGOOptions::NoAction, PGOOptions::NoCSAction, true);
   }
   TM->setPGOOption(PGOOpt);
 

llvm/lib/Passes/PassBuilder.cpp

@@ -1071,6 +1071,23 @@ Expected<bool> parseMemorySSAPrinterPassOptions(StringRef Params) {
                                "MemorySSAPrinterPass");
 }
 
+Expected<std::string> parseMemProfUsePassOptions(StringRef Params) {
+  std::string Result;
+  while (!Params.empty()) {
+    StringRef ParamName;
+    std::tie(ParamName, Params) = Params.split(';');
+
+    if (ParamName.consume_front("profile-filename=")) {
+      Result = ParamName.str();
+    } else {
+      return make_error<StringError>(
+          formatv("invalid MemProfUse pass parameter '{0}' ", ParamName).str(),
+          inconvertibleErrorCode());
+    }
+  }
+  return Result;
+}
+
 } // namespace
 
 /// Tests whether a pass name starts with a valid prefix for a default pipeline

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -1102,6 +1102,10 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
       PGOOpt->CSAction == PGOOptions::CSIRInstr)
     MPM.addPass(PGOInstrumentationGenCreateVar(PGOOpt->CSProfileGenFile));
 
+  if (PGOOpt && Phase != ThinOrFullLTOPhase::ThinLTOPostLink &&
+      !PGOOpt->MemoryProfile.empty())
+    MPM.addPass(MemProfUsePass(PGOOpt->MemoryProfile, PGOOpt->FS));
+
   // Synthesize function entry counts for non-PGO compilation.
   if (EnableSyntheticCounts && !PGOOpt)
     MPM.addPass(SyntheticCountsPropagation());

llvm/lib/Passes/PassRegistry.def

@@ -183,6 +183,13 @@ MODULE_PASS_WITH_PARAMS("embed-bitcode",
                         },
                         parseEmbedBitcodePassOptions,
                         "thinlto;emit-summary")
+MODULE_PASS_WITH_PARAMS("memprof-use",
+                         "MemProfUsePass",
+                        [](std::string Opts) {
+                          return MemProfUsePass(Opts);
+                        },
+                        parseMemProfUsePassOptions,
+                        "profile-filename=S")
 #undef MODULE_PASS_WITH_PARAMS
 
 #ifndef CGSCC_ANALYSIS

llvm/lib/Support/PGOOptions.cpp

@@ -13,12 +13,13 @@ using namespace llvm;
 
 PGOOptions::PGOOptions(std::string ProfileFile, std::string CSProfileGenFile,
                        std::string ProfileRemappingFile,
+                       std::string MemoryProfile,
                        IntrusiveRefCntPtr<vfs::FileSystem> FS, PGOAction Action,
                        CSPGOAction CSAction, bool DebugInfoForProfiling,
                        bool PseudoProbeForProfiling)
     : ProfileFile(ProfileFile), CSProfileGenFile(CSProfileGenFile),
-      ProfileRemappingFile(ProfileRemappingFile), Action(Action),
-      CSAction(CSAction),
+      ProfileRemappingFile(ProfileRemappingFile), MemoryProfile(MemoryProfile),
+      Action(Action), CSAction(CSAction),
       DebugInfoForProfiling(DebugInfoForProfiling ||
                             (Action == SampleUse && !PseudoProbeForProfiling)),
       PseudoProbeForProfiling(PseudoProbeForProfiling), FS(std::move(FS)) {
@@ -36,13 +37,18 @@ PGOOptions::PGOOptions(std::string ProfileFile, std::string CSProfileGenFile,
   // a profile.
   assert(this->CSAction != CSIRUse || this->Action == IRUse);
 
-  // If neither Action nor CSAction, DebugInfoForProfiling or
-  // PseudoProbeForProfiling needs to be true.
+  // Cannot optimize with MemProf profile during IR instrumentation.
+  assert(this->MemoryProfile.empty() || this->Action != PGOOptions::IRInstr);
+
+  // If neither Action nor CSAction nor MemoryProfile are set,
+  // DebugInfoForProfiling or PseudoProbeForProfiling needs to be true.
   assert(this->Action != NoAction || this->CSAction != NoCSAction ||
-         this->DebugInfoForProfiling || this->PseudoProbeForProfiling);
+         !this->MemoryProfile.empty() || this->DebugInfoForProfiling ||
+         this->PseudoProbeForProfiling);
 
   // If we need to use the profile, the VFS cannot be nullptr.
-  assert(this->FS || !(this->Action == IRUse || this->CSAction == CSIRUse));
+  assert(this->FS || !(this->Action == IRUse || this->CSAction == CSIRUse ||
+                       !this->MemoryProfile.empty()));
 }
 
 PGOOptions::PGOOptions(const PGOOptions &) = default;

llvm/lib/Transforms/Instrumentation/MemProfiler.cpp

@@ -38,6 +38,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/HashBuilder.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
@@ -673,9 +674,9 @@ stackFrameIncludesInlinedCallStack(ArrayRef<Frame> ProfileCallStack,
   return InlCallStackIter == InlinedCallStack.end();
 }
 
-void llvm::readMemprof(Module &M, Function &F,
-                       IndexedInstrProfReader *MemProfReader,
-                       const TargetLibraryInfo &TLI) {
+static void readMemprof(Module &M, Function &F,
+                        IndexedInstrProfReader *MemProfReader,
+                        const TargetLibraryInfo &TLI) {
   auto &Ctx = M.getContext();
 
   auto FuncName = getPGOFuncName(F);
@@ -865,3 +866,49 @@ void llvm::readMemprof(Module &M, Function &F,
     }
   }
 }
+
+MemProfUsePass::MemProfUsePass(std::string MemoryProfileFile,
+                               IntrusiveRefCntPtr<vfs::FileSystem> FS)
+    : MemoryProfileFileName(MemoryProfileFile), FS(FS) {
+  if (!FS)
+    this->FS = vfs::getRealFileSystem();
+}
+
+PreservedAnalyses MemProfUsePass::run(Module &M, ModuleAnalysisManager &AM) {
+  LLVM_DEBUG(dbgs() << "Read in memory profile:");
+  auto &Ctx = M.getContext();
+  auto ReaderOrErr = IndexedInstrProfReader::create(MemoryProfileFileName, *FS);
+  if (Error E = ReaderOrErr.takeError()) {
+    handleAllErrors(std::move(E), [&](const ErrorInfoBase &EI) {
+      Ctx.diagnose(
+          DiagnosticInfoPGOProfile(MemoryProfileFileName.data(), EI.message()));
+    });
+    return PreservedAnalyses::all();
+  }
+
+  std::unique_ptr<IndexedInstrProfReader> MemProfReader =
+      std::move(ReaderOrErr.get());
+  if (!MemProfReader) {
+    Ctx.diagnose(DiagnosticInfoPGOProfile(
+        MemoryProfileFileName.data(), StringRef("Cannot get MemProfReader")));
+    return PreservedAnalyses::all();
+  }
+
+  if (!MemProfReader->hasMemoryProfile()) {
+    Ctx.diagnose(DiagnosticInfoPGOProfile(MemoryProfileFileName.data(),
+                                          "Not a memory profile"));
+    return PreservedAnalyses::all();
+  }
+
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+
+  for (auto &F : M) {
+    if (F.isDeclaration())
+      continue;
+
+    const TargetLibraryInfo &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
+    readMemprof(M, F, MemProfReader.get(), TLI);
+  }
+
+  return PreservedAnalyses::none();
+}