Optimize the speed runtime with a committed runtime profile

native/llvm-codegen/src/emit.cpp

@@ -1340,6 +1340,44 @@ emitSplitCodegen(const EmitOptions &Options, Module &M, TargetMachine &Machine,
 
 } // namespace
 
+std::optional<std::string> optimizeModule(Module &M, TargetMachine &Machine,
+                                          const std::string &OptLevel,
+                                          bool Instrument,
+                                          const std::string &ProfileGenerate,
+                                          const std::string &ProfileUse,
+                                          ModuleProfileMatch &Result) {
+  OptimizationLevel Level = optimizationLevel(OptLevel);
+  std::optional<PGOOptions> Profile;
+  if (Instrument) {
+    Profile = PGOOptions(ProfileGenerate, "", "", "", PGOOptions::IRInstr);
+  } else if (!ProfileUse.empty()) {
+    EmitOptions Inspected;
+    Inspected.ProfileUse = ProfileUse;
+    if (std::optional<EmitFailure> Failure = inspectProfile(Inspected))
+      return Failure->Message;
+    reportMissingProfiles();
+    Profile = PGOOptions(ProfileUse, "", "", "", PGOOptions::IRUse);
+  }
+  bool Annotating = Profile && Profile->Action == PGOOptions::IRUse;
+  ProfileMatch Match;
+  if (Annotating)
+    M.getContext().setDiagnosticHandler(
+        std::make_unique<ProfileMatchHandler>(Match));
+  {
+    Pipeline P(Machine, Level, std::move(Profile));
+    if (Annotating)
+      P.countAnnotated(Match.Annotated);
+    P.PB.buildPerModuleDefaultPipeline(Level).run(M, P.MAM);
+  }
+  Result.Annotated = Match.Annotated.load();
+  Result.Mismatched = Match.Mismatched.load();
+  Result.Missing = Match.Missing.load();
+  if (std::optional<EmitFailure> Failure =
+          verify(M, "post_optimization_verification_failed"))
+    return Failure->Message;
+  return std::nullopt;
+}
+
 int emit(const EmitOptions &Options) {
   if (!supportsTarget(Options.Target))
     return reportError("unsupported_target",

native/llvm-codegen/src/emit.h

@@ -1,9 +1,15 @@
 #pragma once
 
+#include <cstdint>
 #include <optional>
 #include <string>
 #include <vector>
 
+namespace llvm {
+class Module;
+class TargetMachine;
+} // namespace llvm
+
 namespace scriptc {
 
 struct EmitOptions {
@@ -36,4 +42,27 @@ struct EmitOptions {
 std::optional<EmitOptions> parseEmitOptions(int Argc, char **Argv);
 int emit(const EmitOptions &Options);
 
+// How a profile matched the functions of one optimized module: functions
+// that received its counts, functions whose control flow changed since it
+// was collected, and functions it has no record of.
+struct ModuleProfileMatch {
+  uint64_t Annotated = 0;
+  uint64_t Mismatched = 0;
+  uint64_t Missing = 0;
+};
+
+// Optimizes `M` with the per-module default pipeline at `OptLevel`, as `emit`
+// does for a whole program, with IR instrumentation when `Instrument` is set
+// or annotation from the indexed `ProfileUse` profile. Instrumentation writes
+// raw profiles to the `ProfileGenerate` pattern; an empty pattern names no
+// file, so the program that links the module names it. Returns an error
+// message, or nothing on success.
+std::optional<std::string> optimizeModule(llvm::Module &M,
+                                          llvm::TargetMachine &Machine,
+                                          const std::string &OptLevel,
+                                          bool Instrument,
+                                          const std::string &ProfileGenerate,
+                                          const std::string &ProfileUse,
+                                          ModuleProfileMatch &Match);
+
 } // namespace scriptc

native/llvm-codegen/src/main.cpp

@@ -58,6 +58,9 @@ int version(int Argc, char **Argv) {
       {"supported_targets", std::move(SupportedTargets)},
       {"default_target", scriptc::DefaultTarget},
       {"data_layout", Machine->createDataLayout().getStringRepresentation()},
+      // Commands and options added within one protocol version, for callers
+      // that also run against older helpers.
+      {"features", json::Array{"runtime-unit-profile", "profile-functions"}},
   };
   outs() << formatv("{0}\n", json::Value(std::move(Response)));
   return 0;
@@ -80,6 +83,8 @@ int main(int Argc, char **Argv) {
     return scriptc::pruneImport(Argc, Argv);
   if (Argc >= 2 && StringRef(Argv[1]) == "merge-profiles")
     return scriptc::mergeProfiles(Argc, Argv);
+  if (Argc >= 2 && StringRef(Argv[1]) == "profile-functions")
+    return scriptc::profileFunctions(Argc, Argv);
   if (Argc >= 2 && StringRef(Argv[1]) == "emit") {
     std::optional<scriptc::EmitOptions> Options =
         scriptc::parseEmitOptions(Argc, Argv);
@@ -95,5 +100,6 @@ int main(int Argc, char **Argv) {
   }
   return scriptc::reportError("usage",
                               "expected 'version --format=json', 'emit', 'runtime-unit', "
-                              "'prune-import', or 'merge-profiles'");
+                              "'prune-import', 'merge-profiles', or "
+                              "'profile-functions'");
 }

native/llvm-codegen/src/profile.cpp

@@ -2,10 +2,14 @@
 #include "diagnostics.h"
 
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/InstrProfWriter.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -19,24 +23,43 @@ namespace scriptc {
 int mergeProfiles(int Argc, char **Argv) {
   std::string Output;
   bool Sparse = false;
+  std::string KeepFunctions;
   std::vector<std::string> Inputs;
   for (int I = 2; I < Argc; ++I) {
     StringRef Arg(Argv[I]);
     if (Arg == "--output" && I + 1 < Argc)
       Output = Argv[++I];
     else if (Arg == "--sparse")
       Sparse = true;
+    else if (Arg == "--keep-functions" && I + 1 < Argc)
+      KeepFunctions = Argv[++I];
     else if (!Arg.starts_with("--"))
       Inputs.push_back(Arg.str());
     else
       return reportError("usage", "merge-profiles requires --output <file> "
                                   "and one or more input profiles, and "
-                                  "accepts --sparse");
+                                  "accepts --sparse and --keep-functions "
+                                  "<file>");
   }
   if (Output.empty() || Inputs.empty())
     return reportError("usage", "merge-profiles requires --output <file> and "
                                 "one or more input profiles");
 
+  // --keep-functions names the functions to keep, one profile name per
+  // line; the records of every other function are left out.
+  StringSet<> Kept;
+  if (!KeepFunctions.empty()) {
+    ErrorOr<std::unique_ptr<MemoryBuffer>> List =
+        MemoryBuffer::getFile(KeepFunctions);
+    if (!List)
+      return reportError("profile_read_failed",
+                         KeepFunctions + ": " + List.getError().message());
+    SmallVector<StringRef, 0> Names;
+    (*List)->getBuffer().split(Names, '\n', -1, /*KeepEmpty=*/false);
+    for (StringRef Name : Names)
+      Kept.insert(Name);
+  }
+
   // Sparse profiles leave out the functions that never ran, which then
   // optimize as without a profile instead of as cold.
   InstrProfWriter Writer(Sparse);
@@ -52,6 +75,8 @@ int mergeProfiles(int Argc, char **Argv) {
                          Path + ": " + toString(std::move(E)));
     std::string RecordError;
     for (NamedInstrProfRecord &Record : **Reader) {
+      if (!KeepFunctions.empty() && !Kept.contains(Record.Name))
+        continue;
       Writer.addRecord(std::move(Record), /*Weight=*/1, [&](Error E) {
         if (RecordError.empty())
           RecordError = toString(std::move(E));
@@ -100,4 +125,26 @@ int mergeProfiles(int Argc, char **Argv) {
   return 0;
 }
 
+int profileFunctions(int Argc, char **Argv) {
+  if (Argc != 4 || StringRef(Argv[2]) != "--input")
+    return reportError("usage",
+                       "profile-functions requires --input <file.profdata>");
+  std::string Path = Argv[3];
+  IntrusiveRefCntPtr<vfs::FileSystem> FS = vfs::getRealFileSystem();
+  Expected<std::unique_ptr<InstrProfReader>> Reader =
+      InstrProfReader::create(Path, *FS);
+  if (!Reader)
+    return reportError("profile_read_failed",
+                       Path + ": " + toString(Reader.takeError()));
+  json::Array Functions;
+  for (const NamedInstrProfRecord &Record : **Reader)
+    Functions.push_back(Record.Name.str());
+  if ((*Reader)->hasError())
+    return reportError("profile_read_failed",
+                       Path + ": " + toString((*Reader)->getError()));
+  json::Object Response{{"ok", true}, {"functions", std::move(Functions)}};
+  outs() << formatv("{0}\n", json::Value(std::move(Response)));
+  return 0;
+}
+
 } // namespace scriptc

native/llvm-codegen/src/profile.h

@@ -2,9 +2,14 @@
 
 namespace scriptc {
 
-// `merge-profiles --output <file.profdata> [--sparse] <input>...`: merges raw
-// (.profraw) or indexed (.profdata) IR profiles into one indexed profile.
-// --sparse leaves out functions that never ran.
+// `merge-profiles --output <file.profdata> [--sparse] [--keep-functions
+// <file>] <input>...`: merges raw (.profraw) or indexed (.profdata) IR
+// profiles into one indexed profile. --sparse leaves out functions that never
+// ran, and --keep-functions every function the file does not name.
 int mergeProfiles(int Argc, char **Argv);
 
+// `profile-functions --input <file.profdata>`: prints the names of the
+// functions an indexed profile has records for, as {"ok":true,"functions":[]}.
+int profileFunctions(int Argc, char **Argv);
+
 } // namespace scriptc

native/llvm-codegen/src/runtime-unit.cpp

@@ -13,6 +13,7 @@
 #include "runtime-unit.h"
 
 #include "diagnostics.h"
+#include "emit.h"
 #include "runtime-import.h"
 #include "target.h"
 
@@ -27,12 +28,16 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/IRReader/IRReader.h"
+#include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/JSON.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/SizeOpts.h"
 
 #include <memory>
 #include <optional>
@@ -53,6 +58,16 @@ struct UnitOptions {
   std::string Tag;
   bool FunctionSections = false;
   bool DataSections = false;
+  // When set, the input is unoptimized and the helper optimizes it at this
+  // level, with the profile-guided optimization below; otherwise the input
+  // is already optimized and is only promoted and emitted.
+  std::string OptLevel;
+  // Instruments the unit for IR profiling; raw profiles go to this pattern,
+  // or, when it is empty, to the file the program that links the unit names.
+  bool ProfileGenerate = false;
+  std::string ProfileGeneratePattern;
+  // Indexed profile (.profdata) that guides the unit's optimization.
+  std::string ProfileUse;
 };
 
 std::optional<UnitOptions> parse(int Argc, char **Argv) {
@@ -80,12 +95,26 @@ std::optional<UnitOptions> parse(int Argc, char **Argv) {
       Options.Target = Value.str();
     else if (Arg == "--tag")
       Options.Tag = Value.str();
+    else if (Arg == "--opt-level")
+      Options.OptLevel = Value.str();
+    else if (Arg == "--profile-generate") {
+      Options.ProfileGenerate = true;
+      Options.ProfileGeneratePattern = Value.str();
+    } else if (Arg == "--profile-use")
+      Options.ProfileUse = Value.str();
     else
       return std::nullopt;
   }
   if (Options.Input.empty() || Options.Object.empty() ||
       Options.Bitcode.empty() || Options.Target.empty() || Options.Tag.empty())
     return std::nullopt;
+  // A profile only guides an optimization the helper runs, and a unit is
+  // either instrumented or annotated.
+  if ((Options.ProfileGenerate || !Options.ProfileUse.empty()) &&
+      Options.OptLevel.empty())
+    return std::nullopt;
+  if (Options.ProfileGenerate && !Options.ProfileUse.empty())
+    return std::nullopt;
   return Options;
 }
 
@@ -98,10 +127,22 @@ bool keepLocal(const GlobalValue &GV) {
   return false;
 }
 
+// The profile's counters and per-function records (instrumented units only)
+// stay private to the unit: the profile runtime finds them by section.
+bool profileData(const GlobalValue &GV) {
+  if (const auto *Object = dyn_cast<GlobalObject>(&GV))
+    if (Object->getSection().starts_with("__llvm_prf"))
+      return true;
+  return GV.getName().starts_with("__llvm_prf") ||
+         GV.getName().starts_with("__profc_") ||
+         GV.getName().starts_with("__profd_") ||
+         GV.getName().starts_with("__profvp_");
+}
+
 void promoteLocals(Module &M, StringRef Tag) {
   unsigned Anonymous = 0;
   auto Promote = [&](GlobalValue &GV) {
-    if (!GV.hasLocalLinkage() || keepLocal(GV))
+    if (!GV.hasLocalLinkage() || keepLocal(GV) || profileData(GV))
       return;
     std::string Base =
         GV.hasName() ? GV.getName().str() : ("anon." + Twine(Anonymous++)).str();
@@ -157,7 +198,9 @@ int runtimeUnit(int Argc, char **Argv) {
     return fail("usage",
                 "runtime-unit requires --input <bc> --object <o> --bitcode <bc> "
                 "--target <triple> --tag <unit> and accepts "
-                "--function-sections --data-sections");
+                "--function-sections --data-sections, and --opt-level <level> "
+                "with either --profile-generate <pattern> or --profile-use "
+                "<profdata>");
   LLVMContext Context;
   SMDiagnostic Diagnostic;
   std::unique_ptr<Module> M = parseIRFile(Options->Input, Diagnostic, Context);
@@ -208,6 +251,42 @@ int runtimeUnit(int Argc, char **Argv) {
     if (F.getFnAttribute("probe-stack").getValueAsString() == "__chkstk_darwin")
       F.addFnAttr("probe-stack", "inline-asm");
 
+  // Profiles name a unit's internal functions after its source file. The
+  // unit's own name keeps those names the same in every checkout, so a
+  // committed profile matches wherever the pack is built.
+  if (!Options->OptLevel.empty()) {
+    M->setSourceFileName(Options->Tag + ".c");
+    // The runtime profile comes from a fixed set of training programs, and a
+    // runtime function they rarely call can be hot in another program (a
+    // long UTF-8 decode, a busy worker channel). Code the profile finds cold
+    // is therefore optimized for speed as without a profile, not for size:
+    // the profile still guides inlining, block layout and branch weights.
+    EnablePGSO.setValue(false);
+    // The profile names of the functions an instrumented unit defines, so
+    // a merged profile can keep the runtime's records and drop the
+    // training programs'.
+    json::Array Functions;
+    if (Options->ProfileGenerate)
+      for (const Function &F : *M)
+        if (!F.isDeclaration() && !F.hasAvailableExternallyLinkage())
+          Functions.push_back(getIRPGOFuncName(F));
+    ModuleProfileMatch Match;
+    if (std::optional<std::string> Failure = optimizeModule(
+            *M, *Machine, Options->OptLevel, Options->ProfileGenerate,
+            Options->ProfileGeneratePattern, Options->ProfileUse, Match))
+      return fail("optimization_failed", *Failure);
+    if (!Options->ProfileUse.empty()) {
+      json::Object Profile{{"annotated", Match.Annotated},
+                           {"mismatched", Match.Mismatched},
+                           {"missing", Match.Missing}};
+      json::Object Response{{"ok", true}, {"profile", std::move(Profile)}};
+      outs() << formatv("{0}\n", json::Value(std::move(Response)));
+    } else if (Options->ProfileGenerate) {
+      json::Object Response{{"ok", true}, {"functions", std::move(Functions)}};
+      outs() << formatv("{0}\n", json::Value(std::move(Response)));
+    }
+  }
+
   promoteLocals(*M, Options->Tag);
   std::string VerifyError;
   raw_string_ostream VerifyStream(VerifyError);