Add IRPGO and CSIRPGO options to Swift

include/swift/AST/DiagnosticsIRGen.def

@@ -41,6 +41,11 @@ ERROR(too_few_output_filenames,none,
 ERROR(no_input_files_for_mt,none,
       "no swift input files for multi-threaded compilation", ())
 
+ERROR(ir_profile_read_failed, none,
+      "error reading profile '%0': %1", (StringRef, StringRef))
+ERROR(ir_profile_invalid, none,
+      "invalid ir profile '%0'", (StringRef))
+
 ERROR(alignment_dynamic_type_layout_unsupported,none,
       "'@_alignment' is not supported on types with dynamic layout", ())
 ERROR(alignment_less_than_natural,none,

include/swift/AST/IRGenOptions.h

@@ -548,10 +548,23 @@ class IRGenOptions {
   /// Path to the profdata file to be used for PGO, or the empty string.
   std::string UseProfile = "";
 
+  /// Path to the profdata file to be used for IR/CS-IR PGO, or the empty string.
+  std::string UseIRProfile = "";
+
   /// Path to the data file to be used for sampling-based PGO,
   /// or the empty string.
   std::string UseSampleProfile = "";
 
+  /// Name of the profile file to use as output for -ir-profile-generate,
+  /// and -cs-profile-generate, or the default string.
+  std::string InstrProfileOutput = "default_%m.profraw";
+
+  /// Whether to enable context-sensitive IR PGO generation.
+  bool EnableCSIRProfileGen = false;
+
+  /// Whether to enable IR level instrumentation.
+  bool EnableIRProfileGen = false;
+
   /// Controls whether DWARF discriminators are added to the IR.
   unsigned DebugInfoForProfiling : 1;
 

include/swift/Option/Options.td

@@ -1608,6 +1608,29 @@ def profile_sample_use : Joined<["-"], "profile-sample-use=">,
   MetaVarName<"<profile data>">,
   HelpText<"Supply sampling-based profiling data from llvm-profdata to enable profile-guided optimization">;
 
+def cs_profile_generate : Flag<["-"], "cs-profile-generate">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  HelpText<"Generate instrumented code to collect context sensitive execution counts into default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
+def cs_profile_generate_EQ : Joined<["-"], "cs-profile-generate=">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  MetaVarName<"<directory>">,
+  HelpText<"Generate instrumented code to collect context sensitive execution counts into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
+def ir_profile_generate: Flag<["-"], "ir-profile-generate">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  HelpText<"Generate instrumented code to collect execution counts into default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
+def ir_profile_generate_EQ : Joined<["-"], "ir-profile-generate=">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  MetaVarName<"<directory>">,
+  HelpText<"Generate instrumented code to collect execution counts into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
+def ir_profile_use : CommaJoined<["-"], "ir-profile-use=">,
+  Flags<[FrontendOption, NoInteractiveOption, ArgumentIsPath]>,
+  MetaVarName<"<profdata>">,
+  HelpText<"Supply an IR-level PGO profdata file to enable profile-guided optimization">;
+
 def embed_bitcode : Flag<["-"], "embed-bitcode">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Embed LLVM IR bitcode as data">;

lib/Driver/DarwinToolChains.cpp

@@ -42,6 +42,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string
 toolchains::Darwin::findProgramRelativeToSwiftImpl(StringRef name) const {
@@ -470,7 +471,7 @@ void
 toolchains::Darwin::addProfileGenerationArgs(ArgStringList &Arguments,
                                              const JobContext &context) const {
   const llvm::Triple &Triple = getTriple();
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     SmallString<128> LibProfile;
     getClangLibraryPath(context.Args, LibProfile);
 

lib/Driver/Driver.cpp

@@ -175,19 +175,92 @@ static void validateWarningControlArgs(DiagnosticEngine &diags,
   }
 }
 
+/// Validates only *generate* profiling flags and their mutual conflicts.
+static void validateProfilingGenerateArgs(DiagnosticEngine &diags,
+                                          const ArgList &args) {
+  const Arg *ProfileGenerate = args.getLastArg(options::OPT_profile_generate);
+  const Arg *IRProfileGenerate =
+      args.getLastArg(options::OPT_ir_profile_generate);
+  const Arg *CSProfileGenerate =
+      args.getLastArg(options::OPT_cs_profile_generate);
+  const Arg *CSProfileGenerateEQ =
+      args.getLastArg(options::OPT_cs_profile_generate_EQ);
+  const Arg *IRProfileGenerateEQ =
+      args.getLastArg(options::OPT_ir_profile_generate_EQ);
+
+  // If both CS Profile forms were specified, report a clear conflict.
+  if (CSProfileGenerate && CSProfileGenerateEQ) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-cs-profile-generate", "-cs-profile-generate=");
+    CSProfileGenerateEQ = nullptr;
+  }
+  // If both IR Profile forms were specified, report a clear conflict.
+  if (IRProfileGenerate && IRProfileGenerateEQ) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-ir-profile-generate", "-ir-profile-generate=");
+    IRProfileGenerateEQ = nullptr;
+  }
+
+  llvm::SmallVector<std::pair<const Arg *, const char *>, 3> gens;
+  if (ProfileGenerate)
+    gens.push_back({ProfileGenerate, "-profile-generate"});
+  if (IRProfileGenerate)
+    gens.push_back({IRProfileGenerate, "-ir-profile-generate"});
+  if (IRProfileGenerateEQ)
+    gens.push_back({IRProfileGenerateEQ, "-ir-profile-generate="});
+  if (CSProfileGenerate)
+    gens.push_back({CSProfileGenerate, "-cs-profile-generate"});
+  else if (CSProfileGenerateEQ)
+    gens.push_back({CSProfileGenerateEQ, "-cs-profile-generate="});
+
+  // Emit pairwise conflicts if more than one generate-mode was selected
+  for (size_t i = 0; i + 1 < gens.size(); ++i) {
+    for (size_t j = i + 1; j < gens.size(); ++j) {
+      diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                     gens[i].second, gens[j].second);
+    }
+  }
+}
+
 static void validateProfilingArgs(DiagnosticEngine &diags,
                                   const ArgList &args) {
+  validateProfilingGenerateArgs(diags, args);
   const Arg *ProfileGenerate = args.getLastArg(options::OPT_profile_generate);
   const Arg *ProfileUse = args.getLastArg(options::OPT_profile_use);
+  const Arg *IRProfileGenerate =
+      args.getLastArg(options::OPT_ir_profile_generate);
+  const Arg *IRProfileGenerateEQ =
+      args.getLastArg(options::OPT_ir_profile_generate_EQ);
+  const Arg *IRProfileUse = args.getLastArg(options::OPT_ir_profile_use);
   if (ProfileGenerate && ProfileUse) {
     diags.diagnose(SourceLoc(), diag::error_conflicting_options,
                    "-profile-generate", "-profile-use");
   }
-
+  if (ProfileGenerate && IRProfileUse) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-profile-generate", "-ir-profile-use");
+  }
+  if (IRProfileGenerate && ProfileUse) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-ir-profile-generate", "-profile-use");
+  }
+  if (IRProfileGenerate && IRProfileUse) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-ir-profile-generate", "-ir-profile-use");
+  }
+  if (IRProfileGenerateEQ && ProfileUse) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-ir-profile-generate=", "-profile-use");
+  }
+  if (IRProfileGenerateEQ && IRProfileUse) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-ir-profile-generate=", "-ir-profile-use");
+  }
   // Check if the profdata is missing
-  if (ProfileUse && !llvm::sys::fs::exists(ProfileUse->getValue())) {
-    diags.diagnose(SourceLoc(), diag::error_profile_missing,
-                   ProfileUse->getValue());
+  for (const Arg *use : {ProfileUse, IRProfileUse}) {
+    if (use && !llvm::sys::fs::exists(use->getValue())) {
+      diags.diagnose(SourceLoc(), diag::error_profile_missing, use->getValue());
+    }
   }
 }
 

lib/Driver/ToolChains.cpp

@@ -160,6 +160,16 @@ bool containsValue(
 
 }
 
+namespace swift::driver::toolchains {
+bool needsInstrProfileRuntime(const llvm::opt::ArgList &Args) {
+  return Args.hasArg(options::OPT_profile_generate) ||
+         Args.hasArg(options::OPT_cs_profile_generate) ||
+         Args.hasArg(options::OPT_cs_profile_generate_EQ) ||
+         Args.hasArg(options::OPT_ir_profile_generate) ||
+         Args.hasArg(options::OPT_ir_profile_generate_EQ);
+}
+} // namespace swift::driver::toolchains
+
 void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
                                       const CommandOutput &output,
                                       const ArgList &inputArgs,
@@ -323,6 +333,11 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
   inputArgs.AddLastArg(arguments, options::OPT_PackageCMO);
   inputArgs.AddLastArg(arguments, options::OPT_profile_generate);
   inputArgs.AddLastArg(arguments, options::OPT_profile_use);
+  inputArgs.AddLastArg(arguments, options::OPT_ir_profile_generate);
+  inputArgs.AddLastArg(arguments, options::OPT_ir_profile_generate_EQ);
+  inputArgs.AddLastArg(arguments, options::OPT_ir_profile_use);
+  inputArgs.AddLastArg(arguments, options::OPT_cs_profile_generate);
+  inputArgs.AddLastArg(arguments, options::OPT_cs_profile_generate_EQ);
   inputArgs.AddLastArg(arguments, options::OPT_profile_coverage_mapping);
   inputArgs.AddAllArgs(arguments, options::OPT_warning_treating_Group);
   inputArgs.AddLastArg(arguments, options::OPT_sanitize_EQ);

lib/Driver/ToolChains.h

@@ -25,6 +25,11 @@ class DiagnosticEngine;
 namespace driver {
 namespace toolchains {
 
+/// True if any *generation* mode of instrumentation-based profile is enabled.
+///
+/// This is used to determine if the profiler runtime should be linked.
+bool needsInstrProfileRuntime(const llvm::opt::ArgList &Args);
+
 class LLVM_LIBRARY_VISIBILITY Darwin : public ToolChain {
 protected:
 

lib/Driver/UnixToolChains.cpp

@@ -40,6 +40,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string
 toolchains::GenericUnix::sanitizerRuntimeLibName(StringRef Sanitizer,
@@ -330,7 +331,7 @@ toolchains::GenericUnix::constructInvocation(const DynamicLinkJobAction &job,
     }
   }
 
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     SmallString<128> LibProfile(SharedResourceDirPath);
     llvm::sys::path::remove_filename(LibProfile); // remove platform name
     llvm::sys::path::append(LibProfile, "clang", "lib");

lib/Driver/WebAssemblyToolChains.cpp

@@ -38,6 +38,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string
 toolchains::WebAssembly::sanitizerRuntimeLibName(StringRef Sanitizer,
@@ -168,7 +169,7 @@ toolchains::WebAssembly::constructInvocation(const DynamicLinkJobAction &job,
         "-fsanitize=" + getSanitizerList(context.OI.SelectedSanitizers)));
   }
 
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     SmallString<128> LibProfile(SharedResourceDirPath);
     llvm::sys::path::remove_filename(LibProfile); // remove platform name
     llvm::sys::path::append(LibProfile, "clang", "lib");

lib/Driver/WindowsToolChains.cpp

@@ -36,6 +36,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string toolchains::Windows::sanitizerRuntimeLibName(StringRef Sanitizer,
                                                          bool shared) const {
@@ -89,7 +90,7 @@ toolchains::Windows::constructInvocation(const DynamicLinkJobAction &job,
     // for now, which supports the behavior via a flag.
     // TODO: Once we've changed coverage to no longer rely on emitting
     // duplicate weak symbols (rdar://131295678), we can remove this.
-    if (context.Args.getLastArg(options::OPT_profile_generate)) {
+    if (swift::driver::toolchains::needsInstrProfileRuntime(context.Args)) {
       return true;
     }
     return false;
@@ -186,7 +187,7 @@ toolchains::Windows::constructInvocation(const DynamicLinkJobAction &job,
                         sanitizerRuntimeLibName("ubsan"));
   }
 
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     Arguments.push_back(context.Args.MakeArgString("-Xlinker"));
     Arguments.push_back(context.Args.MakeArgString(
         Twine({"-include:", llvm::getInstrProfRuntimeHookVarName()})));

lib/Frontend/CompilerInvocation.cpp

@@ -3497,6 +3497,17 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
   const Arg *ProfileSampleUse = Args.getLastArg(OPT_profile_sample_use);
   Opts.UseSampleProfile = ProfileSampleUse ? ProfileSampleUse->getValue() : "";
 
+  Opts.EnableIRProfileGen = Args.hasArg(OPT_ir_profile_generate) ||
+                            Args.hasArg(OPT_ir_profile_generate_EQ);
+  if (auto V = Args.getLastArgValue(OPT_ir_profile_generate_EQ); !V.empty())
+    Opts.InstrProfileOutput = V.str();
+  Opts.EnableCSIRProfileGen = Args.hasArg(OPT_cs_profile_generate) ||
+                              Args.hasArg(OPT_cs_profile_generate_EQ);
+  if (auto V = Args.getLastArgValue(OPT_cs_profile_generate_EQ); !V.empty())
+    Opts.InstrProfileOutput = V.str();
+  const Arg *IRProfileUse = Args.getLastArg(OPT_ir_profile_use);
+  Opts.UseIRProfile = IRProfileUse ? IRProfileUse->getValue() : "";
+
   Opts.DebugInfoForProfiling |= Args.hasArg(OPT_debug_info_for_profiling);